code_kit/agent_code_generator.py

import gradio as gr
import time
import logging
import os
from model_handler import ModelHandler
from tab_code_prompts.html_system_prompt import get_html_system_prompt

# Configure logging
logger = logging.getLogger(__name__)

# Read the content of the JavaScript file for error catching
# Assuming the script is run from the project root or ling-space root.
try:
    with open("static/catch-error.js", "r", encoding="utf-8") as f:
        CATCH_ERROR_JS_SCRIPT = f.read()
except FileNotFoundError:
    logger.error("Error: static/catch-error.js not found. The error catching overlay will not work.")
    CATCH_ERROR_JS_SCRIPT = ""

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 code_generation_agent(code_type, model_choice, user_prompt, color_palette, decoration_style, overall_style, chatbot_history):
    """Generate code and provide a preview, updating a log stream chatbot."""
    logger.info(f"--- [Code Generation] Start ---")
    logger.info(f"Code Type: {code_type}, Model: {model_choice}, Prompt: '{user_prompt}'")

    if not user_prompt:
        chatbot_history.append({"role": "assistant", "content": "🚨 **错误**: 请输入提示词。"})
        yield "", gr.update(value="<p>预览将在此处显示。</p>"), chatbot_history, gr.update()
        return

    chatbot_history.append({"role": "assistant", "content": "⏳ 开始生成代码..."})
    yield "", gr.HTML(get_spinner_html()), chatbot_history, gr.update()

    if user_prompt == "create an error" or user_prompt == "创建一个报错示例":
        error_code = f"""<h1>This will create an error</h1><script>{CATCH_ERROR_JS_SCRIPT}</script><script>nonExistentFunction();</script>"""
        escaped_code = error_code.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>
        """
        chatbot_history.append({"role": "assistant", "content": "✅ **成功**: 已生成一个用于测试的错误页面。"})
        yield error_code, gr.update(value=final_preview_html), chatbot_history, gr.Tabs(selected=0)
        return

    # --- Append Style Prompt ---
    full_user_prompt = user_prompt
    
    if color_palette or decoration_style or overall_style:
        full_user_prompt += "\n\n--- Visual Style Requirements (Strictly Follow These) ---\n"
        if overall_style:
            full_user_prompt += f"Overall Theme/Vibe: {overall_style}\n"
        if decoration_style:
            full_user_prompt += f"UI Decoration Style: {decoration_style}\n"
        if color_palette:
            full_user_prompt += f"Color Palette Mapping (Use these colors for their assigned roles): {color_palette}\n"
            
    logger.info(f"Full Prompt with Style: {full_user_prompt}")
    # ---------------------------

    start_time = time.time()
    model_handler = ModelHandler()

    if code_type == "静态页面":
        system_prompt = get_html_system_prompt()
        full_code_with_think = ""
        full_code_for_preview = ""
        buffer = ""
        is_thinking = False

        for code_chunk in model_handler.generate_code(system_prompt, full_user_prompt, model_choice):
            full_code_with_think += code_chunk
            buffer += code_chunk
            
            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
            
            log_message = f"""
            **⏳ 正在生成中...**
            - **时间:** {elapsed_time:.2f}s
            - **长度:** {generated_length} chars
            - **速度:** {speed:.2f} char/s
            """
            
            if len(chatbot_history) > 0 and "正在生成中" in chatbot_history[-1]["content"]:
                chatbot_history[-1] = {"role": "assistant", "content": log_message}
            else:
                chatbot_history.append({"role": "assistant", "content": log_message})

            yield full_code_with_think, gr.update(), chatbot_history, gr.update()
        
        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>
        """
        chatbot_history.append({"role": "assistant", "content": "✅ **成功**: 代码生成完成！"})
        yield full_code_with_think, gr.HTML(final_preview_html), chatbot_history, gr.Tabs(selected=0)
        logger.info("Static page streaming finished.")