import os
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

hf_token = os.getenv("HF_TOKEN")
model_id = "ZyperAI/Z-AI-0.1-1.1B-Code.web"

print("Loading model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    model_id, 
    token=hf_token, 
    use_fast=False
)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float32, 
    device_map="cpu",
    token=hf_token
)
print("Model loaded successfully.")

def generate_code(prompt, history):
    # Fix 1: Properly structure history for Gradio 6's list-of-dicts style
    messages = []
    for msg in history:
        # Prevent appending empty or broken dictionary structures
        if msg.get("content"):
            messages.append({"role": msg["role"], "content": msg["content"]})
            
    messages.append({"role": "user", "content": prompt})
    
    # Fix 2: Explicitly handle chat template errors if tokens are missing
    try:
        inputs = tokenizer.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            return_tensors="pt"
        ).to("cpu")
    except Exception:
        # Fallback if the specific model lacks a pre-configured chat template
        fallback_prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages]) + "\nassistant:"
        inputs = tokenizer(fallback_prompt, return_tensors="pt").input_ids.to("cpu")
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        input_ids=inputs, # Fix 3: Transformers generation kwargs expects 'input_ids', not 'inputs'
        streamer=streamer,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    # Fix 4: Gradio 6 gr.Chatbot (type="messages") yields back the full history list, 
    # not just a single raw string.
    updated_history = messages.copy()
    updated_history.append({"role": "assistant", "content": ""})
    
    for new_text in streamer:
        updated_history[-1]["content"] += new_text
        yield updated_history

# Gradio 6.x UI setup
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# ⚡ **Z-AI Web Coder**")
    
    chatbot = gr.Chatbot(height=500, show_copy_button=True, type="messages")
    with gr.Row():
        msg = gr.Textbox(
            placeholder="E.g., Create a responsive navigation bar with CSS...", 
            show_label=False, 
            scale=9
        )
        submit = gr.Button("Build", variant="primary", scale=1)
    
    # Fix 5: Use a unified event pipeline so input clearing 
    # doesn't disrupt the streaming text generator.
    submit_click = submit.click(
        generate_code, 
        inputs=[msg, chatbot], 
        outputs=[chatbot]
    ).then(lambda: "", None, [msg])
    
    msg_submit = msg.submit(
        generate_code, 
        inputs=[msg, chatbot], 
        outputs=[chatbot]
    ).then(lambda: "", None, [msg])

if __name__ == "__main__":
    demo.launch()