import os
import re
import time
import json
import subprocess
import gradio as gr
from openai import OpenAI
from duckduckgo_search import DDGS

# --- OpenRouter Setup ---
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "your-openrouter-key")
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
)

# --- Model Definitions ---
MODEL_RESEARCHER = "z-ai/glm-4.5-air"
MODEL_PLANNER = "arcee-ai/trinity-large-preview"
MODEL_FRONTEND = "qwen/qwen3-coder"
MODEL_BACKEND = "minimax/minimax-m2.5"

WORKSPACE_DIR = "/app/workspace"
preview_process = None

# --- Helper Functions ---
def extract_code(text, lang="javascript"):
    """Extracts code blocks from markdown responses."""
    pattern = rf"```{lang}\n(.*?)\n```"
    matches = re.findall(pattern, text, re.DOTALL)
    if matches:
        return matches[0]
    
    # Fallback to any code block
    pattern_any = r"```.*?\n(.*?)\n```"
    matches_any = re.findall(pattern_any, text, re.DOTALL)
    return matches_any[0] if matches_any else text

def web_search(query):
    """Real-time web search for the Researcher AI"""
    try:
        results = DDGS().text(query, max_results=3)
        return "\n".join([f"- {r['title']}: {r['body']}" for r in results])
    except Exception as e:
        return f"Web search failed: {e}"

def run_openrouter(model, system_prompt, user_prompt):
    """Wrapper for OpenRouter API calls"""
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    )
    return response.choices[0].message.content

# --- Multi-Agent Pipeline ---
def generate_website(prompt):
    global preview_process
    logs = []
    
    yield "Step 1: 🌐 Researcher (GLM-4.5-Air) gathering real-time data...", "", ""
    
    # 1. RESEARCHER
    search_context = web_search(prompt)
    research_prompt = f"User Request: {prompt}\nReal-time Web Context: {search_context}\nSummarize the requirements, target audience, and modern tech trends for this app."
    research_data = run_openrouter(
        MODEL_RESEARCHER, 
        "You are an expert tech researcher. Analyze the user request and web data.", 
        research_prompt
    )
    logs.append("✅ Research completed.")
    yield "\n".join(logs), "Step 2: 🏗️ Planner (Trinity-Large) designing architecture...", ""
    
    # 2. PLANNER
    plan_prompt = f"Based on this research: {research_data}\nCreate a full stack architecture plan. Define the React frontend components and Express backend endpoints."
    architecture_plan = run_openrouter(
        MODEL_PLANNER,
        "You are a Lead Solutions Architect. Output a clear, structured technical design.",
        plan_prompt
    )
    logs.append("✅ Architecture planned.")
    yield "\n".join(logs), "Step 3: 💻 Frontend Coder (Qwen3-Coder) building UI...", ""
    
    # 3. FRONTEND CODER
    fe_prompt = f"Architecture: {architecture_plan}\nWrite a complete React App using standard React and Tailwind. Output two things: 1. A package.json file wrapped in ```json ... ``` containing required dependencies. 2. A complete App.jsx wrapped in ```javascript ... ```."
    frontend_code_raw = run_openrouter(
        MODEL_FRONTEND,
        "You are an Expert React Developer. Generate functional, modern React code.",
        fe_prompt
    )
    app_jsx = extract_code(frontend_code_raw, "javascript")
    package_json = extract_code(frontend_code_raw, "json")
    
    # Save Frontend files
    os.makedirs(os.path.join(WORKSPACE_DIR, "frontend", "src"), exist_ok=True)
    with open(os.path.join(WORKSPACE_DIR, "frontend", "package.json"), "w") as f:
        f.write(package_json)
    with open(os.path.join(WORKSPACE_DIR, "frontend", "src", "App.jsx"), "w") as f:
        f.write(app_jsx)
        
    logs.append("✅ Frontend generated and saved.")
    yield "\n".join(logs), "Step 4: 🔐 Backend Coder (MiniMax-M2.5) building secure API...", ""
    
    # 4. BACKEND CODER & INSTRUCTION GUIDE
    be_prompt = f"Architecture: {architecture_plan}\nWrite a secure Express.js backend. Provide the server code wrapped in ```javascript ... ```. Then, provide a detailed Markdown guide wrapped in ```markdown ... ``` explaining exactly where the user needs to add their Postgres/MongoDB strings, Stripe Secret Keys, and Clerk/Auth0 tokens."
    backend_code_raw = run_openrouter(
        MODEL_BACKEND,
        "You are a Senior Security & Backend Engineer.",
        be_prompt
    )
    server_js = extract_code(backend_code_raw, "javascript")
    setup_guide = extract_code(backend_code_raw, "markdown")
    
    # Save Backend files
    os.makedirs(os.path.join(WORKSPACE_DIR, "backend"), exist_ok=True)
    with open(os.path.join(WORKSPACE_DIR, "backend", "server.js"), "w") as f:
        f.write(server_js)
    with open(os.path.join(WORKSPACE_DIR, "SETUP_GUIDE.md"), "w") as f:
        f.write(setup_guide)
        
    logs.append("✅ Backend and Setup Guide generated.")
    yield "\n".join(logs), "Step 5: ⚙️ Installing Dependencies & Starting Live Preview...", setup_guide
    
    # 5. EXECUTION & DEPENDENCY DOWNLOAD
    try:
        # Install dependencies
        subprocess.run(["npm", "install"], cwd=os.path.join(WORKSPACE_DIR, "frontend"), check=True)
        logs.append("✅ AI successfully downloaded dependencies.")
        yield "\n".join(logs), "Starting Server...", setup_guide
        
        # Kill previous process if it exists
        if preview_process:
            preview_process.kill()
            
        # Normally you would use a bundler like Vite. To keep it robust for this example, 
        # we start a simple HTTP server or a lightweight bundler command if defined in package.json.
        # Assuming AI creates standard start scripts:
        preview_process = subprocess.Popen(
            ["npm", "start"], 
            cwd=os.path.join(WORKSPACE_DIR, "frontend"),
            env=dict(os.environ, PORT="3000")
        )
        
        # Give server time to boot
        time.sleep(5)
        
        logs.append("🚀 Live preview is running!")
        iframe_html = f'<iframe src="https://{os.getenv("SPACE_HOST", "localhost")}:3000" width="100%" height="600px" style="border:1px solid #ccc; border-radius: 8px;"></iframe>'
        
        yield "\n".join(logs), iframe_html, setup_guide

    except Exception as e:
        logs.append(f"❌ Execution Error: {str(e)}")
        yield "\n".join(logs), "Failed to start live preview.", setup_guide


# --- UI Setup (Gradio) ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Multi-Agent AI Full-Stack Builder")
    gr.Markdown("Powered by `z-ai/glm-4.5-air` (Researcher), `arcee-ai/trinity-large` (Architect), `qwen/qwen3-coder` (Frontend), and `minimax-m2.5` (Backend).")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(lines=4, placeholder="E.g., Build a SaaS dashboard for an AI writing tool with Stripe subscriptions...", label="App Idea")
            build_btn = gr.Button("Build Application", variant="primary")
            status_box = gr.Textbox(label="Agent Status Logs", lines=10, interactive=False)
            
        with gr.Column(scale=2):
            gr.Markdown("### Live Website Preview")
            preview_box = gr.HTML("<div style='text-align:center; padding: 50px; border:1px dashed gray;'>Preview will appear here once built.</div>")
            
    with gr.Row():
        gr.Markdown("### Setup Guide (DB, Auth & Payments)")
        guide_box = gr.Markdown("Setup instructions will appear here.")
        
    build_btn.click(
        fn=generate_website,
        inputs=[prompt_input],
        outputs=[status_box, preview_box, guide_box]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)