import gradio as gr
import zipfile
import os
from langchain_core.messages import HumanMessage
from utils.langgraph_pipeline import run_pipeline_and_save

TEMP_DIR = "generated_output"
os.makedirs(TEMP_DIR, exist_ok=True)

def save_files(html: str, log: str):
    TEMP_DIR = "generated_output"
    os.makedirs(TEMP_DIR, exist_ok=True)

    # Extract CSS
    import re
    css_match = re.search(r"<style[^>]*>(.*?)</style>", html, flags=re.DOTALL)
    css = css_match.group(1).strip() if css_match else ""

    # Remove <style> tag and replace with <link>
    html_clean = re.sub(r"<style[^>]*>.*?</style>", '<link rel="stylesheet" href="styles.css">', html, flags=re.DOTALL)

    # Paths
    html_path = os.path.join(TEMP_DIR, "ui.html")
    css_path = os.path.join(TEMP_DIR, "styles.css")
    log_path = os.path.join(TEMP_DIR, "agent_log.txt")
    zip_path = os.path.join(TEMP_DIR, "output_bundle.zip")

    # Save files
    with open(html_path, "w", encoding="utf-8") as f:
        f.write(html_clean)

    with open(css_path, "w", encoding="utf-8") as f:
        f.write(css)

    with open(log_path, "w", encoding="utf-8") as f:
        f.write(log)

    with zipfile.ZipFile(zip_path, "w") as zipf:
        zipf.write(html_path, arcname="ui.html")
        zipf.write(css_path, arcname="styles.css")
        zipf.write(log_path, arcname="agent_log.txt")

    return html_path, css_path, log_path, zip_path


def process(user_input):
    messages = [HumanMessage(content=user_input)]
    final_state = run_pipeline_and_save(messages)

    html = final_state["html_output"]
    log = "\n\n".join([f"{msg.type.upper()}: {msg.content}" for msg in final_state["messages"]])

    html_file, css_file, log_file, zip_file = save_files(html, log)
    preview = f"<iframe src='file/{html_file}' width='100%' height='500px' style='border:1px solid #ccc;'></iframe>"

    return preview, html, log, html_file, css_file, log_file, zip_file


with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Multi-Agent UI Generator")
    user_prompt = gr.Textbox(label="Describe your UI", lines=4, value="A yoga retreat homepage with schedule and testimonials.")
    run_btn = gr.Button("Run Agents")

    gr.Markdown("### 🔍 Website Preview")
    preview_html = gr.HTML()

    gr.Markdown("### 🖥️ Generated HTML Code")
    html_output = gr.Textbox(lines=15, label="Final HTML")

    gr.Markdown("### 💬 Agent Dialogue")
    log_output = gr.Textbox(lines=15, label="Conversation Log")

    gr.Markdown("### 📥 Download Files")
    html_download = gr.File(label="Download HTML")
    log_download = gr.File(label="Download Agent Log")
    zip_download = gr.File(label="Download All (.zip)")

    html_download = gr.File(label="Download HTML")
    css_download = gr.File(label="Download CSS")
    log_download = gr.File(label="Download Agent Log")
    zip_download = gr.File(label="Download All (.zip)")

    run_btn.click(
        fn=process,
        inputs=[user_prompt],
        outputs=[
        preview_html,
        html_output,
        log_output,
        html_download,
        css_download,
        log_download,
        zip_download
    ]
)

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