import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def generate_code(prompt, system_message, max_tokens, temperature, top_p, backend):
    full_system_prompt = f"""
You are a code-generation AI. You MUST generate a full website including an index.html file.
Use ONLY the {backend} backend structure. 
Respond ONLY with raw code and file/folder structure. Do NOT explain or add commentary.
""".strip() + "\n\n" + system_message

    messages = [
        {"role": "system", "content": full_system_prompt},
        {"role": "user", "content": prompt}
    ]

    response = ""
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p
    ):
        token = chunk.choices[0].delta.content
        if token:
            response += token
            yield response

with gr.Blocks() as demo:
    gr.Markdown("## WebGen AI — One Prompt → Full Website Generator")

    with gr.Row():
        prompt = gr.Textbox(label="Enter Prompt", placeholder="Describe the website you want...")
        backend = gr.Dropdown(["Flask", "Static", "Node.js"], value="Static", label="Select Backend")

    system_msg = gr.Textbox(value="You are a helpful assistant.", label="System Message")
    max_tokens = gr.Slider(1, 2048, value=512, label="Max Tokens")
    temperature = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
    top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")

    code_output = gr.Code(label="AI Output (Live)", language="text")

    run_btn = gr.Button("Generate Code")

    run_btn.click(
        generate_code,
        inputs=[prompt, system_msg, max_tokens, temperature, top_p, backend],
        outputs=code_output
    )

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