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

# Load a small code model
model_name = "Salesforce/codegen-350M-multi"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_website(prompt):
    instruction = f"""
You are an expert web developer.
Generate ONLY three files for the website:
===index.html===
(use semantic HTML, basic structure)
===style.css===
(responsive, clean CSS)
===script.js===
(simple JS interactions)
Do NOT include explanations.
User request: {prompt}
"""
    result = generator(instruction, max_length=500)[0]['generated_text']

    # Split into files
    files = {}
    if "===index.html===" in result:
        parts = result.split("===")
        for i in range(1, len(parts), 2):
            name = parts[i].strip()
            code = parts[i+1].strip()
            files[name] = code
    # Fallback if model output is unexpected
    return (
        files.get("index.html", ""),
        files.get("style.css", ""),
        files.get("script.js", "")
    )

# Gradio UI
iface = gr.Interface(
    fn=generate_website,
    inputs=gr.Textbox(lines=3, placeholder="Describe your website..."),
    outputs=[
        gr.Textbox(label="index.html"),
        gr.Textbox(label="style.css"),
        gr.Textbox(label="script.js")
    ],
    title="AI Website Generator",
    description="Generate HTML/CSS/JS websites directly online. No files are saved locally."
)

iface.launch()