app.py

import gradio as gr
from huggingface_hub import InferenceClient, HfHubHTTPError
import os
import re
import traceback

# --- Configuration ---
API_TOKEN = os.getenv("HF_TOKEN", None)
# MODEL = "Qwen/Qwen3-32B" # This is a very large model, might require specific inference endpoint/hardware
# Let's try a smaller, generally available model for testing first, e.g., Mixtral
# You can change this back if you are sure Qwen3-32B is available and configured for your space/token
# MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
# Or uncomment the Qwen model if you are certain it's correctly set up for inference:
MODEL = "Qwen/Qwen3-32B"
# i have used Qwen3 because its quiet compatible 


# --- Hugging Face Client Initialization ---
print("--- App Start ---")
if not API_TOKEN:
    print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
    print("Certain models might require a token for access.")
else:
    print(f"HF_TOKEN found (length={len(API_TOKEN)}).") # Don't print the token itself

try:
    print(f"Initializing Inference Client for model: {MODEL}")
    # Explicitly pass token=None if not found, though InferenceClient handles it.
    client = InferenceClient(model=MODEL, token=API_TOKEN if API_TOKEN else None)
    print("Inference Client Initialized Successfully.")
    # Optional: Add a quick test call if feasible, but be mindful of potential costs/rate limits
    # try:
    #     client.text_generation("test", max_new_tokens=1)
    #     print("Test generation successful.")
    # except Exception as test_e:
    #     print(f"Warning: Test generation failed. Client might be initialized but model access could be problematic. Error: {test_e}")

except HfHubHTTPError as http_err:
    # More specific error handling for HTTP errors (like 401 Unauthorized, 403 Forbidden, 404 Not Found)
    error_message = (
        f"Failed to initialize model client for {MODEL} due to an HTTP error.\n"
        f"Status Code: {http_err.response.status_code}\n"
        f"Error: {http_err}\n"
        f"Check:\n"
        f"1. If '{MODEL}' is a valid model ID on Hugging Face Hub.\n"
        f"2. If the model requires gating or specific permissions.\n"
        f"3. If your HF_TOKEN is correct and has the necessary permissions (set as a Secret in your Space).\n"
        f"4. If the default Inference API supports this model or if a dedicated Inference Endpoint is needed."
    )
    print(f"ERROR: {error_message}")
    raise gr.Error(error_message)
except Exception as e:
    error_message = (
        f"An unexpected error occurred while initializing the model client for {MODEL}.\n"
        f"Error Type: {type(e).__name__}\n"
        f"Error: {e}\n"
        f"Traceback:\n{traceback.format_exc()}\n" # Add traceback
        f"Check HF_TOKEN, model availability, network connection, and Space resources."
    )
    print(f"ERROR: {error_message}")
    raise gr.Error(error_message)

# --- Helper Functions ---

# Parse all ```filename.ext\n<code>``` blocks
def parse_code_blocks(response: str) -> list:
    pattern = r"```([^\n]+)\n(.*?)```"
    blocks = re.findall(pattern, response, re.DOTALL)
    files = []
    for filename, code in blocks:
        filename = filename.strip()
        code = code.strip()
        # Basic language detection (can be expanded)
        lang = None
        if filename.endswith(".py"):
            lang = "python"
        elif filename.endswith(".js"):
            lang = "javascript"
        elif filename.endswith(".html"):
            lang = "html"
        elif filename.endswith(".css"):
            lang = "css"
        elif filename.endswith(".json"):
            lang = "json"
        elif filename.endswith(".md"):
            lang = "markdown"
        elif filename.endswith(".sh") or filename.endswith(".bash"):
            lang = "bash"
        elif filename.endswith(".java"):
            lang = "java"
        # Add more extensions as needed

        files.append({
            "filename": filename,
            "language": lang,
            "code": code
        })
    # Add logging to see what's parsed
    # print(f"Parsed {len(files)} code blocks.")
    # for i, f in enumerate(files):
    #     print(f"  Block {i}: filename='{f['filename']}', lang='{f['language']}', code_len={len(f['code'])}")
    return files

def strip_think_tags(text: str) -> str:
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

def extract_thoughts(text: str) -> str:
    matches = re.findall(r"<think>(.*?)</think>", text, flags=re.DOTALL)
    # Only return the last thought block for cleaner display? Or join all? Let's join.
    return "\n---\n".join(match.strip() for match in matches).strip()

# --- System Message ---
system_message = (
    "You are a helpful AI assistant specialized in generating website code. "
    "Generate all the necessary files based on the user's request. "
    "Output each file within a separate markdown code block formatted exactly like this:\n"
    "```filename.ext\n"
    "<code>\n"
    "```\n"
    "Do not add any explanatory text outside the code blocks. Ensure the filenames have appropriate extensions. "
    "If you need to think step-by-step, use <think>...</think> tags. These tags will be hidden from the final user output but help guide your generation process."
)

# --- Code Generation Function ---
def generate_code(prompt, backend_choice, max_tokens, temperature, top_p):
    if not prompt:
        # Handle empty prompt case
        yield [], gr.update(value="Please enter a description for the website.", visible=True)
        return

    # Use f-string formatting for clarity
    user_prompt = f"USER_PROMPT: {prompt}\nUSER_BACKEND_PREFERENCE: {backend_choice}"

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_prompt}
    ]

    full_response = ""
    current_thoughts = ""
    accumulated_error = "" # Accumulate errors during stream

    # Reset outputs: Clear previous code blocks and show/clear thinking box
    # Yield an empty list to the gr.Column to clear it.
    # Make thinking box visible but empty.
    yield [], gr.update(visible=True, value="Generating code...")

    print(f"\n--- Generating Code ---")
    print(f"Prompt: {prompt[:100]}...") # Log truncated prompt
    print(f"Backend: {backend_choice}, Max Tokens: {max_tokens}, Temp: {temperature}, Top-P: {top_p}")

    try:
        stream = client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature if temperature > 0 else 0.01, # Ensure temp is positive
            top_p=top_p,
            # Consider adding stop sequences if the model tends to run on
            # stop=["```\n\n", "\n\nHuman:", "\n\nUSER:"] # Example stop sequences
        )

        code_updates = [] # Store the gr.Code components to yield

        for i, message in enumerate(stream):
            # Check for errors in the stream message (some providers might include error info)
            if hasattr(message, 'error') and message.error:
                 accumulated_error += f"Error in stream chunk {i}: {message.error}\n"
                 print(f"ERROR in stream chunk {i}: {message.error}")
                 continue # Skip this chunk if it's an error indicator

            # Ensure the path to content is correct
            try:
                 # Common path: message.choices[0].delta.content
                 token = message.choices[0].delta.content
                 # Handle potential None token at the end of the stream or in error cases
                 if token is None:
                     token = ""
                 # print(f"Token {i}: '{token}'") # DEBUG: print each token
            except (AttributeError, IndexError, TypeError) as e:
                 # Handle unexpected message structure
                 print(f"Warning: Could not extract token from stream message {i}. Structure: {message}. Error: {e}")
                 token = "" # Assign empty string to avoid breaking accumulation

            if isinstance(token, str):
                full_response += token

                # Update thinking box periodically (e.g., every 10 tokens or if thoughts change)
                if i % 10 == 0 or "<think>" in token or "</think>" in token:
                    thoughts = extract_thoughts(full_response)
                    if thoughts != current_thoughts:
                        current_thoughts = thoughts
                        # Don't yield code_updates here yet, only update thoughts
                        yield code_updates, gr.update(value=current_thoughts if current_thoughts else "Thinking...", visible=True)


                # Update code blocks less frequently or when a block seems complete
                # Heuristic: update if the response ends with ```
                if token.strip().endswith("```") or i % 20 == 0: # Adjust frequency as needed
                    cleaned_response = strip_think_tags(full_response)
                    parsed_files = parse_code_blocks(cleaned_response)

                    # Create gr.Code components for the parsed files
                    # Compare with existing code_updates to avoid redundant updates if content hasn't changed significantly
                    new_code_updates = []
                    changed = False
                    if len(parsed_files) != len(code_updates):
                        changed = True
                    else:
                         # Quick check if filenames/code lengths differ significantly
                         for idx, f in enumerate(parsed_files):
                             if (idx >= len(code_updates) or
                                 f["filename"] != code_updates[idx].label or
                                 len(f["code"]) != len(code_updates[idx].value)): # Simple length check
                                 changed = True
                                 break

                    if changed or not code_updates: # Update if changed or first time
                         code_updates = []
                         for f in parsed_files:
                             code_updates.append(
                                 gr.Code(
                                     value=f["code"],
                                     label=f["filename"],
                                     language=f["language"]
                                 )
                             )
                         # Yield the list of gr.Code components to the gr.Column
                         # Also update thoughts (might be slightly out of sync, but acceptable)
                         yield code_updates, gr.update(value=current_thoughts if current_thoughts else "Thinking...", visible=True)


        # --- Final Update after Stream Ends ---
        print("Stream finished.")
        if accumulated_error:
            print(f"Errors occurred during stream:\n{accumulated_error}")
            # Decide how to show this to the user, e.g., append to thoughts or show separately
            current_thoughts += f"\n\n**Streaming Errors:**\n{accumulated_error}"

        cleaned_response = strip_think_tags(full_response)
        final_files = parse_code_blocks(cleaned_response)
        print(f"Final parsed files: {len(final_files)}")

        final_code_updates = []
        if not final_files and not accumulated_error:
             # Handle case where no code blocks were generated
             final_code_updates.append(gr.Markdown("No code blocks were generated. The model might have responded with text instead, or the format was incorrect."))
             print("Warning: No code blocks found in the final response.")
             # Optionally show the raw response for debugging
             # final_code_updates.append(gr.Code(label="Raw Response", value=cleaned_response, language="text"))

        elif not final_files and accumulated_error:
             final_code_updates.append(gr.Markdown(f"**Error during generation:**\n{accumulated_error}"))

        else:
            for f in final_files:
                final_code_updates.append(
                    gr.Code(
                        value=f["code"],
                        label=f["filename"],
                        language=f["language"]
                    )
                )

        # Yield final code blocks and hide thinking box (or show final thoughts/errors)
        final_thought_update = gr.update(visible=True if current_thoughts else False, value=current_thoughts)
        yield final_code_updates, final_thought_update

    except HfHubHTTPError as http_err:
        # Handle errors during the streaming call itself
        error_message = (
            f"**Error during code generation (HTTP Error):**\n"
            f"Status Code: {http_err.response.status_code}\n"
            f"Error: {http_err}\n"
            f"This could be due to rate limits, invalid input, model errors, or token issues.\n"
            f"Check the Hugging Face Space logs for more details."
        )
        print(f"ERROR: {error_message}")
        print(traceback.format_exc())
        # Yield error message in the output area
        yield [gr.Markdown(error_message)], gr.update(visible=False) # Hide thinking box on error

    except Exception as e:
        error_message = (
             f"**An unexpected error occurred during code generation:**\n"
             f"Error Type: {type(e).__name__}\n"
             f"Error: {e}\n\n"
             f"**Traceback:**\n```\n{traceback.format_exc()}\n```\n"
             f"Check the Hugging Face Space logs for more details."
        )
        print(f"ERROR: {error_message}")
        # Yield error message in the output area
        yield [gr.Markdown(error_message)], gr.update(visible=False) # Hide thinking box on error


# --- 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. Code files will appear below. Uses `mistralai/Mixtral-8x7B-Instruct-v0.1` by default (check code to change).") # Update description

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(label="Website Description", lines=6, placeholder="e.g., A simple landing page with a title, a paragraph, and a button linking to example.com")
            backend_radio = gr.Radio(["Static (HTML/CSS/JS)", "Flask", "Node.js"], label="Backend Preference (Influences AI)", value="Static (HTML/CSS/JS)")
            generate_button = gr.Button("✨ Generate Website Code", variant="primary")

            with gr.Accordion("Advanced Settings", open=False):
                 max_tokens_slider = gr.Slider(512, 8192, value=4096, step=256, label="Max New Tokens") # Increased max potential tokens
                 temperature_slider = gr.Slider(0.0, 1.2, value=0.6, step=0.05, label="Temperature (0=deterministic, >1=more creative)") # Allow 0
                 top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P (Nucleus Sampling)")

        with gr.Column(scale=3):
            thinking_box = gr.Textbox(label="Model Activity / Thoughts", visible=False, interactive=False, lines=2)
            # Use gr.Column to hold the dynamic code blocks
            # Remove the update lambda, it's not needed for Column
            file_outputs = gr.Column(elem_id="code-output-area")


    generate_button.click(
        fn=generate_code,
        inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
        # Output to the Column and the Textbox
        outputs=[file_outputs, thinking_box],
        # api_name="generate_code" # Optional: for API access
    )

# --- Launch ---
if __name__ == "__main__":
    print("Starting Gradio App...")
    # Use queue() for handling multiple users and streaming
    # Set share=False unless you specifically want a public link from local execution
    # Set debug=True for more detailed Gradio errors locally (remove/set False for production)
    demo.queue().launch(debug=False, share=False)
    print("Gradio App Launched.")