import gradio as gr
from huggingface_hub import InferenceClient
import os
import re

API_TOKEN = os.getenv("HF_TOKEN", None)
MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"

try:
    client = InferenceClient(model=MODEL, token=API_TOKEN) if API_TOKEN else InferenceClient(model=MODEL)
except Exception as e:
    raise gr.Error(f"Failed to initialize model client. Error: {e}")

def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature: float, top_p: float):
    system_message = (
        "You are an AI assistant programmed to generate website codes only. "
        "You must not use triple backticks (```html, ```python, etc.). "
        "If multiple files are needed, separate them clearly using:\n"
        "TAB.NAME={filename}\n"
        "Only generate code. No explanations, no phrases like 'Here is the code'. "
        "If user asks non-website code, reply:\n"
        "'hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-('."
    )
    user_prompt = f"USER_PROMPT = {prompt}\nUSER_BACKEND = {backend_choice}"
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_prompt}
    ]
    full_response = ""
    try:
        stream = client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        )
        for message in stream:
            token = message.choices[0].delta.content
            if isinstance(token, str):
                full_response += token
        cleaned_response = full_response.strip()
        cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response)
        cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response)
        cleaned_response = re.sub(r"<\s*\|?\s*(user|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE)
        common_phrases = [
            "Here is the code:", "Okay, here is the code:", "Here's the code:",
            "Sure, here is the code you requested:", "Let me know if you need anything else."
        ]
        for phrase in common_phrases:
            if cleaned_response.lower().startswith(phrase.lower()):
                cleaned_response = cleaned_response[len(phrase):].lstrip()
        return cleaned_response
    except Exception as e:
        return f"## Error\n\nFailed to generate code.\n**Reason:** {e}"

def split_files(full_code_text):
    file_blocks = []
    splits = re.split(r'TAB\.NAME=\{(.+?)\}', full_code_text)
    if not splits[0].strip():
        splits = splits[1:]
    for i in range(0, len(splits), 2):
        if i+1 >= len(splits):
            continue
        filename = splits[i].strip()
        content = splits[i+1].strip()
        file_blocks.append((filename, content))
    return file_blocks

with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
    gr.Markdown("# ✨ Website Code Generator ✨")
    gr.Markdown(
        "Describe the website you want. The AI will generate website code.\n\n"
        "**Rules:**\n"
        "- Backend hint (Static / Flask / Node.js).\n"
        "- Always fully SFW and minimal errors.\n"
        "- Only websites. No other codes.\n"
        "- Multiple files use TAB.NAME={filename}."
    )
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Website Description",
                placeholder="e.g., A simple landing page with a hero section and contact form.",
                lines=6,
            )
            backend_radio = gr.Radio(
                ["Static", "Flask", "Node.js"],
                label="Backend Context",
                value="Static",
            )
            generate_button = gr.Button("✨ Generate Website Code", variant="primary")
        with gr.Column(scale=3):
            main_output = gr.Code(
                label="Full Generated Code",
                language="html",
                lines=25,
                interactive=False,
            )
            extra_outputs = gr.Column()

    with gr.Accordion("Advanced Settings", open=False):
        max_tokens_slider = gr.Slider(
            minimum=512,
            maximum=4096,
            value=3072,
            step=256,
            label="Max New Tokens"
        )
        temperature_slider = gr.Slider(
            minimum=0.1,
            maximum=1.2,
            value=0.7,
            step=0.1,
            label="Temperature"
        )
        top_p_slider = gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.9,
            step=0.05,
            label="Top-P"
        )

    def generate_and_display(prompt, backend, max_tokens, temperature, top_p):
        full_code = generate_code(prompt, backend, max_tokens, temperature, top_p)
        files = split_files(full_code)
        file_outputs = []
        for filename, content in files:
            file_outputs.append(gr.Markdown.update(value=f"### {filename}"))
            file_outputs.append(gr.Code.update(value=content, language="html" if filename.endswith(".html") else "python" if filename.endswith(".py") else "javascript" if filename.endswith(".js") else "text"))
        return full_code, file_outputs

    generate_button.click(
        fn=generate_and_display,
        inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
        outputs=[main_output, extra_outputs]
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()