import os
import re

import gradio as gr
import requests
from gradio.utils import get_upload_folder
from gradio_client import utils as client_utils

MODEL = "google/gemini-4-argon"
URL = "https://openrouter.ai/api/v1/chat/completions"
SYSTEM = (
    "You are a front-end generator. Reply with ONE complete, self-contained HTML file "
    "(<!DOCTYPE html> ... </html>) with all CSS in <style> and all JS in <script>. "
    "No external assets unless from a public CDN. No explanation, only the HTML."
)


def generate_html(prompt: str, openrouter_api_key: str) -> str:
    key = (openrouter_api_key or "").strip()
    if not key:
        raise gr.Error("Paste your OpenRouter API key.")
    if not (prompt or "").strip():
        raise gr.Error("Write a prompt.")
    r = requests.post(
        URL,
        timeout=300,
        headers={"Authorization": f"Bearer {key}"},
        json={
            "model": MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": prompt},
            ],
        },
    )
    if not r.ok:
        if r.status_code == 402:
            raise gr.Error("OpenRouter error 402: out of credits or key limit reached")
        raise gr.Error(f"OpenRouter error {r.status_code}")
    return r.json()["choices"][0]["message"]["content"] or ""


def strip_fences(text: str) -> str:
    m = re.search(r"```(?:html)?\s*(.*?)```", text or "", re.S | re.I)
    html = m.group(1) if m else (text or "")
    low = html.lower()
    i = low.find("<!doctype")
    if i < 0:
        i = low.find("<html")
    return (html[i:] if i >= 0 else html).strip()


def save_html(html: str) -> dict:
    # Gradio's upload folder is always allowed for /gradio_api/file= serving.
    directory = get_upload_folder()
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"page_{os.urandom(6).hex()}.html")
    with open(path, "w", encoding="utf-8") as f:
        f.write(html)
    url = f"/gradio_api/file={client_utils.encode_file_path(path)}"
    return {"path": path, "url": url, "is_file": True}


demo = gr.Workflow(
    graph=os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow.json"),
    bind={"generate_html": generate_html, "strip_fences": strip_fences, "save_html": save_html},
)

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