import gradio as gr
from huggingface_hub import InferenceClient

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

def respond(message, history, system_message, max_tokens, temperature, top_p, backend):
    forced_system = 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()

    system_message = forced_system + "\n\n" + system_message

    messages = [{"role": "system", "content": system_message}]
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

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

with gr.Blocks() as demo:
    gr.Markdown("# WebGen AI\nGenerate a complete website with your selected backend.")

    with gr.Row():
        system_msg = gr.Textbox(value="You're a helpful assistant.", label="System Message", lines=2)
        backend = gr.Dropdown(["Flask", "Static", "Node.js"], value="Static", label="Backend")
    
    with gr.Row():
        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")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Ask to generate a website...", label="Your Message")

    def user_submit(message, chat_history):
        return "", chat_history + [[message, None]]

    submit = gr.Button("Send")

    state = gr.State([])

    submit.click(
        user_submit,
        [msg, state],
        [msg, state]
    ).then(
        respond,
        [msg, state, system_msg, max_tokens, temperature, top_p, backend],
        chatbot
    )

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