tab_code.py

import gradio as gr
import re
import os
import sys
import tempfile
import time
import uuid
import logging
import subprocess
from models import generate_code_for_tab
from utils import find_free_port

# Global variable to track the Gradio subprocess
gradio_process = None

# Configure logging
logger = logging.getLogger(__name__)

def get_gradio_sys_prompt():
    """Get the System Prompt for generating Gradio applications"""
    return """
You are an expert Gradio developer. Create a complete, runnable, single-file Gradio application based on the user's request.
The code must be self-contained in a single Python script.
IMPORTANT: The main Gradio app instance MUST be assigned to a variable named `demo`.
The script must end with the app launch command: `demo.launch()`.
Do not include any explanations, just the raw Python code inside a ```python block.
"""

def run_gradio_app(code, port):
    """
    Write the Gradio code to a temporary file and run it in a subprocess on the specified port.
    """
    temp_dir = tempfile.mkdtemp()
    file_path = os.path.join(temp_dir, "app_to_run.py")
    with open(file_path, "w") as f:
        f.write(code)

    env = os.environ.copy()
    env["GRADIO_SERVER_PORT"] = str(port)
    
    # Use subprocess.Popen to run the script in a completely isolated process
    process = subprocess.Popen([sys.executable, file_path], env=env)
    
    # Return the process object for management
    return process, temp_dir, file_path


def get_html_sys_prompt():
    """Get the System Prompt for generating static pages"""
    return """
You are an expert front-end developer. Create a complete, modern, and responsive single HTML file based on the user's request.
The file must be self-contained, including all necessary HTML, CSS, and JavaScript.
Do not include any explanations, just the raw HTML code.
"""

def run_gradio_in_thread(code, url_queue, session_id):
    """Run the Gradio application in a separate thread to avoid blocking the main application"""
    temp_dir = tempfile.mkdtemp()
    file_path = os.path.join(temp_dir, "app.py")
    with open(file_path, "w") as f:
        f.write(code)

    python_executable = sys.executable
    process = subprocess.Popen(
        [python_executable, file_path],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
        universal_newlines=True
    )
    running_processes[session_id] = process

    for line in process.stdout:
        print(f"Gradio App stdout: {line.strip()}")
        if "Running on local URL:" in line:
            url = line.split("Running on local URL:")[1].strip()
            url_queue.put(url)
            break
    
    process.wait()
    try:
        os.remove(file_path)
        os.rmdir(temp_dir)
    except OSError as e:
        print(f"Error cleaning up temp files: {e}")

def get_spinner_html():
    """Return HTML with a CSS spinner animation"""
    return """
    <div style="width: 100%; height: 600px; display: flex; justify-content: center; align-items: center; border: 1px solid #ddd; background-color: #f9f9f9;">
        <div class="spinner"></div>
    </div>
    <style>
        .spinner {
            border: 4px solid rgba(0, 0, 0, 0.1);
            width: 36px;
            height: 36px;
            border-radius: 50%;
            border-left-color: #09f;
            animation: spin 1s ease infinite;
        }
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }
    </style>
    """

def generate_code(code_type, model_choice, user_prompt):
    """Generate code and provide a preview based on the user's selected type (static page or Gradio application)."""
    global gradio_process
    logger.info(f"--- [Code Generation] Start ---")
    logger.info(f"Code Type: {code_type}, Model: {model_choice}, Prompt: '{user_prompt}'")

    # If a Gradio subprocess is running, terminate it first
    if gradio_process and gradio_process.poll() is None:
        gradio_process.terminate()
        gradio_process.wait(timeout=5)
        logger.info("Terminated existing Gradio process.")

    if not user_prompt:
        yield "Please enter a prompt.", gr.HTML("Preview will appear here."), "Status: Error", gr.update()
        return

    yield "", gr.HTML(get_spinner_html()), "Status: Initializing...", gr.update()

    start_time = time.time()

    if code_type == "Gradio App":
        port = find_free_port()
        system_prompt = get_gradio_sys_prompt()
        
        full_code = ""
        for chunk in generate_code_for_tab(system_prompt, user_prompt, code_type, model_choice):
            full_code += chunk
            elapsed_time = time.time() - start_time
            generated_length = len(full_code)
            speed = generated_length / elapsed_time if elapsed_time > 0 else 0
            status = f"""
            **Status:** Generating...
            - **Generation Time:** {elapsed_time:.2f}s
            - **Generated Length:** {generated_length} chars
            - **Average Speed:** {speed:.2f} char/s
            """
            # Display the generated raw code in real-time
            yield full_code, gr.update(), status, gr.update()

        # Extract the Python code block from the model output
        python_code_match = re.search(r"```python\n(.*?)```", full_code, re.DOTALL)
        if python_code_match:
            python_code = python_code_match.group(1).strip()
        else:
            # If no code block is found, assume the entire output is code
            logger.warning("Could not find Python code block, assuming the entire output is code.")
            python_code = full_code.strip()

        try:
            # Verify that the code can at least be compiled
            compile(python_code, '<string>', 'exec')
            
            # Use subprocess to launch the Gradio application
            process, temp_dir, temp_file_path = run_gradio_app(python_code, port)
            gradio_process = process
            
            # Give the application some time to start
            time.sleep(10) 
            
            preview_html = f'<iframe src="http://127.0.0.1:{port}" width="100%" height="600px" frameborder="0"></iframe>'
            status = "**Status:** Generation Complete"
            yield python_code, gr.HTML(preview_html), status, gr.Tabs(selected=0)
            logger.info(f"Gradio app started on port {port}.")
        except SyntaxError as e:
            error_message = f"Generated code is not valid Python: {e}"
            logger.error(f"Raw output from the model:\n{full_code}")
            status = f"**Status:** Error - {error_message}"
            yield full_code, gr.HTML(f"<p style='color: red;'>{error_message}</p>"), status, gr.update()
            logger.error(error_message)

    elif code_type == "Static Page":
        system_prompt = get_html_sys_prompt()
        full_code_with_think = ""
        full_code_for_preview = ""
        buffer = ""
        is_thinking = False
        
        # The model's raw output is streamed here
        for code_chunk in generate_code_for_tab(system_prompt, user_prompt, code_type, model_choice):
            full_code_with_think += code_chunk
            buffer += code_chunk
            
            # Process the buffer to filter out think tags for the preview
            while True:
                if is_thinking:
                    end_index = buffer.find("</think>")
                    if end_index != -1:
                        is_thinking = False
                        buffer = buffer[end_index + len("</think>"):]
                    else:
                        break 
                else:
                    start_index = buffer.find("<think>")
                    if start_index != -1:
                        part_to_add = buffer[:start_index]
                        full_code_for_preview += part_to_add
                        is_thinking = True
                        buffer = buffer[start_index:]
                    else:
                        full_code_for_preview += buffer
                        buffer = ""
                        break
            
            elapsed_time = time.time() - start_time
            generated_length = len(full_code_with_think)
            speed = generated_length / elapsed_time if elapsed_time > 0 else 0
            status = f"""
            **Status:** Generating...
            - **Generation Time:** {elapsed_time:.2f}s
            - **Generated Length:** {generated_length} chars
            - **Average Speed:** {speed:.2f} char/s
            """

            # In the generation process, we only update the code and status, not the preview.
            # The preview will be updated once at the very end.
            yield full_code_with_think, gr.update(), status, gr.update()
        
        # Final update for the preview without the spinner
        escaped_code = full_code_for_preview.replace("'", "&apos;").replace('"', '&quot;')
        final_preview_html = f"""
        <div style="width: 100%; height: 600px; border: 1px solid #ddd; overflow: hidden; position: relative; background-color: #f9f9f9;">
            <iframe srcdoc='{escaped_code}'
                    style="position: absolute; top: 0; left: 0; width: 200%; height: 200%; transform: scale(0.5); transform-origin: 0 0; border: none;">
            </iframe>
        </div>
        """
        status = "**Status:** Generation Complete"
        yield full_code_with_think, gr.HTML(final_preview_html), status, gr.Tabs(selected=0)
        logger.info("Static page streaming finished.")


def toggle_fullscreen(is_fullscreen):
    """Toggle fullscreen mode visibility"""
    is_fullscreen = not is_fullscreen
    new_button_text = "Exit Fullscreen" if is_fullscreen else "Fullscreen Preview"
    panel_visibility = not is_fullscreen
    return is_fullscreen, gr.update(value=new_button_text), gr.update(visible=panel_visibility)

def cleanup_process():
    """Terminate the subprocess when the application exits"""
    global gradio_process
    if gradio_process and gradio_process.poll() is None:
        gradio_process.terminate()
        gradio_process.wait()
        logger.info("Cleaned up Gradio process on exit.")

import atexit
atexit.register(cleanup_process)

def create_code_tab():
    """Create the UI Tab for the code generation feature"""
    fullscreen_state = gr.State(False)

    html_examples = [
        "Write a hello world alert",
        "Create a Canvas animation of continuous colorful fireworks blooming on a black background.",
        "Generate a Canvas special effect with iridescent light streams.",
        "Design a particle system Canvas animation that interacts with the mouse.",
        "Implement a classic snake game using HTML Canvas.",
        "Create a Canvas animation simulating the 'Elephant\\\'s Toothpaste' chemical experiment: in a container, colored foam continuously and rapidly gushes out, expands, and overflows, filling the entire screen. Use a suitable drawing js library, such as three.js.",
        "Create a dreamy low-poly floating island scene with dynamic lighting and soft animations in a single HTML file. Use a suitable drawing js library, such as three.js.",
    ]

    gradio_examples = [
        "Write a hello world alert",
        "Create a simple calculator that supports addition, subtraction, multiplication, and division",
        "Create a simple Gradio application that displays 'Hello, Gradio!'",
        "Create an image classification application that can classify uploaded images",
        "Create a text summarization application that can summarize long input texts",
        "Create a simple chatbot",
    ]

    def update_examples(code_type):
        if code_type == "Static Page":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column(scale=1) as left_panel:
                gr.Markdown("### 1. Select Code Type")
                code_type_radio = gr.Radio(["Static Page", "Gradio App"], value="Static Page", label="Code Type")
                
                gr.Markdown("### 2. Select Model")
                model_choice_radio = gr.Radio(
                    ["Better (using Ling-1T)", "Faster (using Ring-flash-2.0)"],
                    value="Faster (using Ring-flash-2.0)",
                    label="Model Selection"
                )

                gr.Markdown("### 3. Enter Your Requirements")
                prompt_input = gr.Textbox(lines=5, placeholder="e.g., Create a simple page with a title and a button", label="Prompt")
                with gr.Column() as html_examples_column:
                    html_examples_component = gr.Examples(
                        examples=html_examples,
                        inputs=prompt_input,
                        label="✨ Why not try these cool examples",
                        examples_per_page=12,
                    )
                with gr.Column(visible=False) as gradio_examples_column:
                    gradio_examples_component = gr.Examples(
                        examples=gradio_examples,
                        inputs=prompt_input,
                        label="✨ Why not try these cool examples",
                        examples_per_page=12,
                    )
                generate_button = gr.Button("Generate Code", variant="primary")
                status_output = gr.Markdown(value="Status: Waiting", visible=True)
            
            with gr.Column(scale=2):
                with gr.Tabs(elem_id="result_tabs") as result_tabs:
                    with gr.TabItem("Live Preview", id=0):
                        with gr.Row():
                            gr.Markdown("### 3. Live Preview")
                            fullscreen_button = gr.Button("Fullscreen Preview", scale=0)
                        preview_output = gr.HTML(value="<p>Preview will appear here.</p>")
                    with gr.TabItem("Generated Source Code", id=1):
                        gr.Markdown("### 4. Generated Source Code")
                        code_output = gr.Code(language="html", label="Generated Code")

        code_type_radio.change(fn=update_examples, inputs=code_type_radio, outputs=[html_examples_column, gradio_examples_column])

        generate_button.click(
            fn=generate_code,
            inputs=[code_type_radio, model_choice_radio, prompt_input],
            outputs=[code_output, preview_output, status_output, result_tabs]
        )
        
        fullscreen_button.click(
            fn=toggle_fullscreen,
            inputs=[fullscreen_state],
            outputs=[fullscreen_state, fullscreen_button, left_panel]
        )

    return demo