import gradio as gr
import base64
import io
import json
from PIL import Image
import httpx
import html
from lzstring import LZString
import os
from openai import OpenAI

# --- 默认API配置 ---
DEFAULT_NEBIUS_API_KEY = "eyJhbGciOiJIUzI1NiIsImtpZCI6IlV6SXJWd1h0dnprLVRvdzlLZWstc0M1akptWXBvX1VaVkxUZlpnMDRlOFUiLCJ0eXAiOiJKV1QifQ.eyJzdWIiOiJnb29nbGUtb2F1dGgyfDEwNTA1MTQzMDg2MDMwMzIxNDEwMiIsInNjb3BlIjoib3BlbmlkIG9mZmxpbmVfYWNjZXNzIiwiaXNzIjoiYXBpX2tleV9pc3N1ZXIiLCJhdWQiOlsiaHR0cHM6Ly9uZWJpdXMtaW5mZXJlbmNlLmV1LmF1dGgwLmNvbS9hcGkvdjIvIl0sImV4cCI6MTkwNjU5ODA0NCwidXVpZCI6ImNkOGFiMWZlLTIxN2QtNDJlMy04OWUwLWM1YTg4MjcwMGVhNyIsIm5hbWUiOiJodW5nZ2luZyIsImV4cGlyZXNfYXQiOiIyMDMwLTA2LTAyVDAyOjM0OjA0KzAwMDAifQ.MA52QuIiNruK7_lX688RXAEI2TkcCOjcf_02XrpnhI8"
NEBIUS_BASE_URL = "https://api.studio.nebius.com/v1/"

# --- 核心工具函数 ---
def analyze_image(image: Image.Image, nebius_api_key: str = "") -> str:
    """
    Analyze an uploaded image and provide a detailed description of its content and layout.
    
    Args:
        image: The PIL Image object to analyze
        nebius_api_key: Nebius API key for image analysis
        
    Returns:
        A detailed description of the image content, layout, and website type
    """
    if image is None:
        return "Error: No image provided"
    
    # 使用提供的API密钥或默认密钥
    api_key = nebius_api_key.strip() if nebius_api_key.strip() else DEFAULT_NEBIUS_API_KEY
    
    if not api_key:
        return "Error: Nebius API key not provided"
    
    try:
        # 配置Nebius OpenAI客户端
        client = OpenAI(
            base_url=NEBIUS_BASE_URL,
            api_key=api_key,
        )
        
        # 转换图片为base64
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        # 创建提示
        prompt = """
        Analyze this image and provide a concise description.
        Describe the main elements, colors, layout, and UI components.
        Identify what type of website or application this resembles.
        Focus on structural and visual elements that would be important for recreating the design.
        """
        
        # 使用Qwen2.5-VL-72B-Instruct模型进行图像分析
        response = client.chat.completions.create(
            model="Qwen/Qwen2.5-VL-72B-Instruct",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{img_str}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=1000,
            temperature=0.7
        )
        
        return response.choices[0].message.content
        
    except Exception as e:
        return f"Error analyzing image: {str(e)}"

def generate_html_code(description: str, nebius_api_key: str = "") -> str:
    """
    Generate HTML/CSS/JavaScript code based on a website description.
    
    Args:
        description: Detailed description of the website to generate
        nebius_api_key: Nebius API key for code generation
        
    Returns:
        Complete HTML code with embedded CSS and JavaScript
    """
    if not description or description.startswith("Error"):
        return "Error: Invalid or missing description"
    
    # 使用提供的API密钥或默认密钥
    api_key = nebius_api_key.strip() if nebius_api_key.strip() else DEFAULT_NEBIUS_API_KEY
    
    if not api_key:
        return "Error: Nebius API key not provided"
    
    prompt = f"""
    Generate a complete, responsive webpage based on this description:
    
    {description}
    
    Requirements:
    - Use modern HTML5, CSS3, and vanilla JavaScript only
    - Include TailwindCSS via CDN for styling
    - Make it responsive and visually appealing
    - Use placeholder images from https://unsplash.com/ if needed
    - Include proper semantic HTML structure
    - Add interactive elements where appropriate
    - Ensure the design matches the described layout and style
    
    Return only the complete HTML code starting with <!DOCTYPE html> and ending with </html>.
    """
    
    try:
        # 配置Nebius OpenAI客户端
        client = OpenAI(
            base_url=NEBIUS_BASE_URL,
            api_key=api_key,
        )
        
        # 使用DeepSeek-V3-0324模型进行代码生成
        response = client.chat.completions.create(
            model="deepseek-ai/DeepSeek-V3-0324",
            messages=[
                {"role": "user", "content": prompt}
            ],
            max_tokens=8000,
            temperature=0.7
        )
        
        html_code = response.choices[0].message.content
        
        # 清理代码格式
        if html_code.strip().startswith("```html"):
            html_code = html_code.split("```html", 1)[1].strip()
        if html_code.strip().endswith("```"):
            html_code = html_code.rsplit("```", 1)[0].strip()
        
        # 确保代码完整性
        if "<!DOCTYPE html>" in html_code and "</html>" in html_code:
            start = html_code.find("<!DOCTYPE html>")
            end = html_code.rfind("</html>") + 7
            return html_code[start:end]
        else:
            return html_code
            
    except Exception as e:
        return f"Error generating HTML code: {str(e)}"

def create_codesandbox(html_code: str) -> str:
    """
    Create a CodeSandbox project from HTML code.
    
    Args:
        html_code: Complete HTML code to upload to CodeSandbox
        
    Returns:
        CodeSandbox URL or error message
    """
    if not html_code or html_code.startswith("Error"):
        return "Error: No valid HTML code provided"
    
    try:
        # 准备文件结构
        files = {
            "index.html": {
                "content": html_code,
                "isBinary": False
            },
            "package.json": {
                "content": json.dumps({
                    "name": "ai-generated-website",
                    "version": "1.0.0",
                    "description": "Website generated from image analysis",
                    "main": "index.html",
                    "scripts": {
                        "start": "serve .",
                        "build": "echo 'No build required'"
                    },
                    "devDependencies": {
                        "serve": "^14.0.0"
                    }
                }, indent=2),
                "isBinary": False
            }
        }
        
        # 准备参数
        parameters = {
            "files": files,
            "template": "static"
        }
        
        # 压缩数据
        json_str = json.dumps(parameters, separators=(',', ':'))
        lz = LZString()
        compressed = lz.compressToBase64(json_str)
        compressed = compressed.replace('+', '-').replace('/', '_').rstrip('=')
        
        # 生成URL
        codesandbox_url = f"https://codesandbox.io/api/v1/sandboxes/define?parameters={compressed}"
        
        # 尝试创建sandbox
        import requests
        
        # 尝试POST请求创建sandbox
        try:
            response = requests.post(
                "https://codesandbox.io/api/v1/sandboxes/define",
                json=parameters,
                timeout=10
            )
        except:
            # 如果POST失败，返回GET URL
            return codesandbox_url
        
        if response.status_code == 200:
            result = response.json()
            sandbox_id = result.get("sandbox_id")
            if sandbox_id:
                return f"https://codesandbox.io/s/{sandbox_id}"
        
        # 如果POST失败，返回GET URL
        return codesandbox_url
        
    except Exception as e:
        return f"Error creating CodeSandbox: {str(e)}"

def screenshot_to_code(image: Image.Image, nebius_api_key: str = "") -> tuple:
    """
    Complete pipeline: analyze image and generate corresponding HTML code.
    
    Args:
        image: Screenshot image to analyze
        nebius_api_key: Nebius API key for both image analysis and code generation
        
    Returns:
        Tuple of (description, html_code)
    """
    # 分析图片
    description = analyze_image(image, nebius_api_key)
    
    if description.startswith("Error"):
        return description, "Error: Cannot generate code due to image analysis failure"
    
    # 生成代码
    html_code = generate_html_code(description, nebius_api_key)
    
    return description, html_code

# --- Gradio界面 ---
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="AI Website Generator - MCP Compatible (Nebius)",
    css="""
    .api-section { background: #f8f9fa; padding: 15px; border-radius: 8px; margin: 10px 0; }
    .tool-section { border: 1px solid #e0e0e0; padding: 15px; border-radius: 8px; margin: 10px 0; }
    """
) as app:
    
    gr.Markdown("""
    # 🚀 AI Website Generator (MCP Compatible - Nebius)
    
    Transform website screenshots into functional HTML code using Nebius AI.
    
    **Features:**
    - 📸 Image analysis with Qwen2.5-VL-72B-Instruct
    - 💻 HTML/CSS/JS code generation with DeepSeek-V3-0324
    - 🌐 Direct CodeSandbox deployment
    - 🔧 MCP (Model Context Protocol) compatible
    
    **Tools Available:**
    - `analyze_image`: Analyze website screenshots
    - `generate_html_code`: Generate HTML from descriptions
    - `create_codesandbox`: Deploy to CodeSandbox
    - `screenshot_to_code`: Complete pipeline
    """)
    
    with gr.Tab("🎯 Quick Generate"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 Input", elem_classes=["tool-section"])
                
                # API配置
                with gr.Group():
                    gr.Markdown("**Nebius API Key (Required)**")
                    nebius_key = gr.Textbox(
                        label="Nebius API Key",
                        type="password",
                        placeholder="Enter your Nebius API key",
                        value=DEFAULT_NEBIUS_API_KEY
                    )
                
                # 图片上传
                image_input = gr.Image(
                    type="pil",
                    label="Upload Website Screenshot",
                    sources=["upload", "clipboard"]
                )
                
                generate_btn = gr.Button(
                    "🎨 Generate Website",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=2):
                gr.Markdown("### 📋 Results", elem_classes=["tool-section"])
                
                description_output = gr.Textbox(
                    label="📝 Image Analysis",
                    lines=6,
                    interactive=False
                )
                
                html_output = gr.Code(
                    label="💻 Generated HTML Code",
                    language="html",
                    lines=15
                )
                
                with gr.Row():
                    codesandbox_btn = gr.Button("🚀 Deploy to CodeSandbox")
                    codesandbox_output = gr.Textbox(
                        label="CodeSandbox URL",
                        interactive=False
                    )
    
    with gr.Tab("🔧 Individual Tools"):
        gr.Markdown("### Use individual MCP tools")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("#### 📸 Image Analysis Tool")
                img_tool = gr.Image(type="pil", label="Image")
                nebius_key_tool = gr.Textbox(label="Nebius API Key", type="password", value=DEFAULT_NEBIUS_API_KEY)
                analyze_btn = gr.Button("Analyze Image")
                analysis_result = gr.Textbox(label="Analysis Result", lines=5)
            
            with gr.Column():
                gr.Markdown("#### 💻 Code Generation Tool")
                desc_input = gr.Textbox(label="Description", lines=3)
                nebius_key_tool2 = gr.Textbox(label="Nebius API Key", type="password", value=DEFAULT_NEBIUS_API_KEY)
                code_btn = gr.Button("Generate Code")
                code_result = gr.Code(label="Generated Code", language="html")
    
    # 事件绑定
    generate_btn.click(
        fn=screenshot_to_code,
        inputs=[image_input, nebius_key],
        outputs=[description_output, html_output]
    )
    
    codesandbox_btn.click(
        fn=create_codesandbox,
        inputs=[html_output],
        outputs=[codesandbox_output]
    )
    
    analyze_btn.click(
        fn=analyze_image,
        inputs=[img_tool, nebius_key_tool],
        outputs=[analysis_result]
    )
    
    code_btn.click(
        fn=generate_html_code,
        inputs=[desc_input, nebius_key_tool2],
        outputs=[code_result]
    )
    
    # 示例
    gr.Examples(
        examples=[["1.jpg"]],
        inputs=[image_input],
        label="📷 Example Screenshots"
    )

if __name__ == "__main__":
    app.launch(mcp_server=True, share=False)