import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Initialize model and tokenizer
model_id = "Tesslate/WEBGEN-4B-Preview"

# Load model and tokenizer once during app initialization
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

def generate_code(prompt):
    inputs = tok(prompt, return_tensors="pt").to(model.device)
    
    # Generate with streaming
    from transformers import TextIteratorStreamer
    from threading import Thread
    
    streamer = TextIteratorStreamer(tok, skip_special_tokens=True)
    
    generation_kwargs = dict(
        **inputs,
        max_new_tokens=10000,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tok.eos_token_id,
        streamer=streamer
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    generated_text = ""
    for new_text in streamer:
        generated_text += new_text
        
        # Extract only the code portion (remove prompt and any non-code text)
        if "```html" in generated_text:
            code_start = generated_text.find("```html") + 7
            code_end = generated_text.find("```", code_start)
            if code_end != -1:
                clean_code = generated_text[code_start:code_end].strip()
            else:
                clean_code = generated_text[code_start:].strip()
        elif "<html" in generated_text or "<!DOCTYPE" in generated_text:
            # Find the start of HTML code
            html_start = generated_text.find("<!DOCTYPE")
            if html_start == -1:
                html_start = generated_text.find("<html")
            if html_start != -1:
                clean_code = generated_text[html_start:].strip()
            else:
                clean_code = generated_text
        else:
            clean_code = generated_text
        
        yield clean_code
    
    thread.join()

with gr.Blocks() as demo:
    gr.Markdown("# 🧪 Text-to-Code Generator")
    gr.Markdown("Generate HTML code from natural language prompts with WEBGEN-4B Preview model")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt",
                value="Make a single-file landing page for 'LatticeDB'. Style: modern, generous whitespace, Tailwind, rounded-xl, soft gradients. Sections: navbar, hero (headline + 2 CTAs), features grid, pricing (3 tiers), FAQ accordion, footer. Constraints: semantic HTML, no external JS.",
                lines=5,
                max_lines=10
            )
            generate_button = gr.Button("Generate Code")
            code_output = gr.Code(label="Generated HTML", language="html", lines=20, interactive=False)
        with gr.Column():
            html_output = gr.HTML(label="Preview")
    
    # When button is clicked, generate code and update both outputs
    generate_button.click(
        fn=generate_code,
        inputs=prompt_input,
        outputs=code_output
    ).then(
        fn=lambda code: code,
        inputs=code_output,
        outputs=html_output
    )

demo.launch()