import html
import re

import gradio as gr
import spaces
import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

MODEL_ID = "google/diffusiongemma-26B-A4B-it"

if not torch.cuda.is_available():
    raise RuntimeError("CUDA is not available — model would silently load on CPU.")

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = DiffusionGemmaForBlockDiffusion.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="auto",
)

PROMPT_TEMPLATE = (
    "Eres un desarrollador front-end experto. Genera una página HTML completa y "
    "autocontenida en un solo archivo (CSS y JavaScript inline, sin recursos "
    "externos) según la siguiente descripción. Responde ÚNICAMENTE con el código "
    "HTML, sin explicaciones ni bloques de markdown.\n\n"
    "Descripción: {description}"
)


def extract_html(text):
    # Keep only the last model turn and drop the thought channel.
    if "<|turn>model\n" in text:
        text = text.rsplit("<|turn>model\n", 1)[-1]
    if "<channel|>" in text:
        text = text.rsplit("<channel|>", 1)[-1]
    for tok in ("<eos>", "<pad>", "<turn|>", "<bos>"):
        text = text.replace(tok, "")
    text = text.strip()
    # Unwrap a markdown code fence if the model added one anyway.
    fenced = re.search(r"```(?:html)?\s*(.*?)```", text, re.DOTALL)
    if fenced:
        text = fenced.group(1).strip()
    return text


@spaces.GPU(duration=120)
def generate(description, max_new_tokens):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": PROMPT_TEMPLATE.format(description=description)}
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=max_new_tokens)

    decoded = processor.decode(output[0], skip_special_tokens=False)
    if isinstance(decoded, list):
        decoded = decoded[0]
    code = extract_html(decoded)

    preview = (
        f'<iframe srcdoc="{html.escape(code, quote=True)}" sandbox="allow-scripts" '
        'style="width:100%;height:70vh;border:1px solid #ccc;border-radius:8px;'
        'background:white"></iframe>'
    )
    return preview, code


with gr.Blocks(title="DiffusionGemma HTML Generator") as demo:
    gr.Markdown(
        "# DiffusionGemma → HTML\n"
        "Describe lo que quieres y [google/diffusiongemma-26B-A4B-it]"
        "(https://huggingface.co/google/diffusiongemma-26B-A4B-it) genera la página. "
        "El resultado se renderiza abajo en un iframe aislado."
    )
    with gr.Row():
        description = gr.Textbox(
            label="¿Qué quieres construir?",
            placeholder="ej: una landing page para una cafetería, con menú y formulario de contacto",
            lines=2,
            scale=4,
        )
        btn = gr.Button("Generar", variant="primary", scale=1)
    max_tokens = gr.Slider(
        minimum=256, maximum=4096, value=2048, step=256, label="Max new tokens"
    )
    with gr.Tab("Vista previa"):
        preview = gr.HTML()
    with gr.Tab("Código"):
        code = gr.Code(language="html")

    btn.click(generate, [description, max_tokens], [preview, code])
    description.submit(generate, [description, max_tokens], [preview, code])

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