app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
import re
# import traceback # Optional for debugging

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

try:
    print(f"Initializing Inference Client for model: {MODEL}")
    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 for {MODEL}. Error: {e}. Check HF_TOKEN and model availability.")

# --- Helper Function to Parse Code during Streaming ---
def parse_streaming_code(current_response: str) -> dict:
    """
    Parses potentially incomplete AI output stream.
    Identifies if .TAB separator is present and splits code accordingly.
    Returns dict with html_code, backend_code, filename, language, and visibility flag.
    """
    files = {
        'html_code': '',
        'backend_code': '',
        'backend_filename': 'Backend', # Default label
        'backend_language': None,
        'backend_visible': False       # Default visibility
    }
    separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?'
    match = re.search(separator_pattern, current_response)

    if match:
        # Separator found in the stream so far
        html_part = current_response[:match.start()].strip()
        backend_part = current_response[match.end():].strip() # Code after separator
        backend_filename = match.group(1).strip()

        files['html_code'] = html_part
        files['backend_code'] = backend_part
        files['backend_filename'] = backend_filename
        files['backend_visible'] = True # Make backend visible

        # Determine language
        if backend_filename.endswith(".py"): files['backend_language'] = 'python'
        elif backend_filename.endswith(".js"): files['backend_language'] = 'javascript'
        elif backend_filename.endswith(".css"): files['backend_language'] = 'css'
        else: files['backend_language'] = None
    else:
        # No separator found yet, assume all content is HTML
        files['html_code'] = current_response.strip()
        # Keep backend_visible as False

    return files

# --- Minimal Cleaning for Intermediate Stream Chunks ---
def clean_intermediate_stream(text: str) -> str:
    """ Basic cleaning for streaming chunks (e.g., remove chat markers). """
    cleaned = re.sub(r"<\s*\|?\s*(user|system|assistant)\s*\|?\s*>", "", text, flags=re.IGNORECASE)
    # Avoid stripping aggressively during stream as it might remove partial code
    return cleaned

# --- Core Code Generation Function - Modified for Streaming UI Updates ---
def generate_code(
    prompt: str,
    backend_choice: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    print(f"Streaming code generation for: {prompt[:100]}... | Backend: {backend_choice}")

    system_message = ( # System message remains the same
         "You are an AI that generates website code. You MUST ONLY output the raw code, without any conversational text like 'Here is the code' or explanations before or after the code blocks. "
         "You MUST NOT wrap the code in markdown fences like ```html, ```python, or ```js. "
         "If the user requests 'Static' or the prompt clearly implies only frontend code, generate ONLY the content for the `index.html` file. "
         "If the user requests 'Flask' or 'Node.js' and the prompt requires backend logic, you MUST generate both the `index.html` content AND the corresponding main backend file content (e.g., `app.py` for Flask, `server.js` or `app.js` for Node.js). "
         "When generating multiple files, you MUST separate them EXACTLY as follows: "
         "1. Output the complete code for the first file (e.g., `index.html`). "
         "2. On a new line immediately after the first file's code, add the separator '.TAB[NAME=filename.ext]' (e.g., '.TAB[NAME=app.py]' or '.TAB[NAME=server.js]'). "
         "3. On the next line, immediately start the code for the second file. "
         "Generate only the necessary files (usually index.html and potentially one backend file). "
         "The generated website code must be SFW and have minimal errors. "
         "Only include comments where user modification is strictly required. Avoid explanatory comments. "
         "If the user asks you to create code that is NOT for a website, you MUST respond ONLY with the exact phrase: "
         "'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 = ""
    # Initialize state for UI updates
    current_html = ""
    current_backend = ""
    current_backend_label = "Backend"
    current_backend_lang = None
    is_backend_visible = False

    # Initial clear of outputs
    yield gr.update(value="", visible=True), gr.update(visible=False), gr.update(value="", visible=False)

    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
                # Clean intermediate response minimally
                cleaned_response_chunk = clean_intermediate_stream(full_response)

                # Parse the *entire accumulated* cleaned response on each iteration
                parsed_state = parse_streaming_code(cleaned_response_chunk)

                # Update state variables
                current_html = parsed_state['html_code']
                current_backend = parsed_state['backend_code']
                current_backend_label = parsed_state['backend_filename']
                current_backend_lang = parsed_state['backend_language']
                is_backend_visible = parsed_state['backend_visible'] # This determines visibility

                # Prepare Gradio updates based on the *current* parsed state
                html_update = gr.update(value=current_html)
                # Update the backend tab's visibility
                tab_update = gr.update(visible=is_backend_visible)
                # Update the backend code block's content, label, language, and visibility
                backend_code_update = gr.update(
                    value=current_backend,
                    label=current_backend_label,
                    language=current_backend_lang,
                    visible=is_backend_visible # Make code block visible *if* tab is visible
                )

                # Yield updates for html_code, backend_tab, backend_code
                yield html_update, tab_update, backend_code_update

        # --- Final Clean and Update after Stream ---
        # Ensure the final state is clean and fully parsed
        final_cleaned_response = full_response.strip()
        # Remove fences/phrases missed during stream (optional but good practice)
        final_cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", final_cleaned_response)
        final_cleaned_response = re.sub(r"\n?\s*```\s*$", "", final_cleaned_response)
        common_phrases = ["Here is the code:", "Okay, here is the code:", "Here's the code:", "Sure, here is the code you requested:"]
        temp_lower = final_cleaned_response.lower()
        for phrase in common_phrases:
            if temp_lower.startswith(phrase.lower()):
                final_cleaned_response = final_cleaned_response[len(phrase):].lstrip()
                temp_lower = final_cleaned_response.lower()

        # Check for refusal message in the final response
        refusal_message = "hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-("
        if refusal_message in final_cleaned_response:
             yield gr.update(value=refusal_message), gr.update(visible=False), gr.update(value="", visible=False)
             return # Stop processing

        # Final parse
        final_parsed_state = parse_streaming_code(final_cleaned_response)

        # Final updates to ensure everything is correct
        final_html_update = gr.update(value=final_parsed_state['html_code'])
        final_tab_update = gr.update(visible=final_parsed_state['backend_visible'])
        final_backend_code_update = gr.update(
            value=final_parsed_state['backend_code'],
            label=final_parsed_state['backend_filename'],
            language=final_parsed_state['backend_language'],
            visible=final_parsed_state['backend_visible']
        )
        yield final_html_update, final_tab_update, final_backend_code_update


    except Exception as e:
        print(f"ERROR during code generation stream: {e}")
        # traceback.print_exc() # Uncomment for detailed traceback
        error_message = f"## Error\n\nFailed during streaming.\n**Reason:** {e}"
        # Show error in HTML block, hide backend tab and code
        yield gr.update(value=error_message), gr.update(visible=False), gr.update(value="", visible=False)


# --- Build Gradio Interface ---
with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
    gr.Markdown("# ✨ Website Code Generator ✨")
    gr.Markdown(
        "Describe the website you want. See code generated live.\n"
        "If backend code is generated, a second tab will appear."
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Website Description",
                placeholder="e.g., A Flask app with a simple chat using Socket.IO",
                lines=6,
            )
            backend_radio = gr.Radio(
                ["Static", "Flask", "Node.js"],
                label="Backend Context",
                value="Static",
                info="Guides AI if backend code (like Python/JS) is needed alongside HTML."
            )
            generate_button = gr.Button("✨ Generate Website Code", variant="primary")

        with gr.Column(scale=3):
            # Define Tabs structure
            with gr.Tabs(elem_id="code-tabs"):
                # Tab 1: Always present for HTML
                with gr.Tab("index.html", elem_id="html-tab") as html_tab:
                     html_code_output = gr.Code(
                         label="index.html",
                         language="html",
                         lines=25,
                         interactive=False,
                         elem_id="html_code",
                     )
                # Tab 2: Backend - defined but starts hidden
                with gr.Tab("Backend", elem_id="backend-tab", visible=False) as backend_tab:
                     backend_code_output = gr.Code(
                         label="Backend", # Label updated dynamically if tab becomes visible
                         language=None,   # Language updated dynamically
                         lines=25,
                         interactive=False,
                         elem_id="backend_code",
                         visible=False # Code block also starts hidden
                     )

    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"
        )

    # Click function now targets html_code_output, backend_tab, and backend_code_output
    generate_button.click(
        fn=generate_code,
        inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
        # Outputs MUST match the number and order of updates yielded by the function
        outputs=[html_code_output, backend_tab, backend_code_output],
    )

if __name__ == "__main__":
    if not API_TOKEN:
        print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
    demo.queue(max_size=10).launch()