app.py

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

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

# --- Initialize Inference Client ---
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 and Format Code ---
def parse_and_format_code(raw_response: str) -> str:
    """
    Parses raw AI output containing .TAB separators
    and formats it for display in a single code block.
    """
    # Default filename for the first block if no TAB is present or before the first TAB
    default_filename = "index.html"
    separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?'

    filenames = re.findall(r'\.TAB\[NAME=([^\]]+)\]', raw_response)
    code_blocks = re.split(separator_pattern, raw_response)

    formatted_output = []

    # Handle the first block (before any potential separator)
    first_block = code_blocks[0].strip()
    if first_block:
        formatted_output.append(f"--- START FILE: {default_filename} ---\n\n{first_block}\n\n--- END FILE: {default_filename} ---")

    # Handle subsequent blocks associated with filenames found by the separator
    # re.split with capturing groups includes the separator AND the filename capture group
    # in the results. We need to iterate carefully.
    # Example: ['code1', 'app.py', 'code2', 'style.css', 'code3']
    idx = 1
    while idx < len(code_blocks) -1 :
        filename = code_blocks[idx] # This should be the filename captured by the pattern
        code = code_blocks[idx + 1].strip() # This should be the code after the separator
        if code : # Only add if there's actual code content
             formatted_output.append(f"--- START FILE: {filename} ---\n\n{code}\n\n--- END FILE: {filename} ---")
        idx += 2 # Move past the filename and the code block

    # If no separators were found, filenames list will be empty, and only the first block runs.
    # If separators were found, the loop processes the rest.

    if not formatted_output: # Handle case where input was empty or only whitespace
        return raw_response # Return the original if parsing yields nothing

    # Join the formatted blocks with clear separation
    return "\n\n\n".join(formatted_output)


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

    system_message = (
        "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}
    ]

    response_stream = ""
    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):
                response_stream += token
                full_response += token
                # Yield intermediate stream for responsiveness during generation
                yield response_stream # Show raw output as it comes

        # --- Post-processing (After Stream Ends) ---
        cleaned_response = full_response.strip()

        # Fallback fence removal
        cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response)
        cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response)
        # Remove potential chat markers
        cleaned_response = re.sub(r"<\s*\|?\s*(user|system|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE).strip()
        # Remove common conversational phrases (if they slip through)
        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.",
            "```html", "```python", "```javascript", "```",
        ]
        temp_response_lower = cleaned_response.lower()
        for phrase in common_phrases:
            if temp_response_lower.startswith(phrase.lower()):
                cleaned_response = cleaned_response[len(phrase):].lstrip()
                temp_response_lower = cleaned_response.lower()

        # Check for refusal message
        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 full_response:
             yield refusal_message # Yield the exact refusal message
             return # Stop processing

        # --- PARSE and FORMAT the final cleaned response ---
        formatted_final_code = parse_and_format_code(cleaned_response)

        # Yield the final, formatted version replacing the streamed content
        yield formatted_final_code

    except Exception as e:
        print(f"ERROR during code generation: {e}") # Log detailed error
        # traceback.print_exc() # Uncomment for full traceback in logs
        yield f"## Error\n\nFailed to generate or process code.\n**Reason:** {e}"


# --- 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. The AI will generate the necessary code.\n"
        "If multiple files are generated (e.g., for Flask/Node.js), they will be shown below separated by `--- START FILE: filename ---` markers.\n" # Updated description
        "**Output Format:**\n"
        "- No explanations, just code.\n"
        "- Multiple files separated by file markers.\n" # Updated description
        "- Minimal necessary comments only.\n\n"
        "**Rules:**\n"
        "- Backend choice guides the AI on whether to include server-side code.\n"
        "- Always SFW and aims for minimal errors.\n"
        "- Only generates website-related code."
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Website Description",
                placeholder="e.g., A Flask app with a form that stores data in a variable.",
                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):
            code_output = gr.Code(
                label="Generated Code (Scroll for multiple files)", # Updated label
                language=None, # Keep as None for mixed/plain text display
                lines=30,
                interactive=False, # Keep non-interactive
            )

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

    # The click function now yields the final formatted string AFTER streaming
    generate_button.click(
        fn=generate_code,
        inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
        outputs=code_output,
    )

if __name__ == "__main__":
    if not API_TOKEN:
        print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
    # Increased max_size slightly for potentially larger combined outputs
    demo.queue(max_size=15).launch()