from fastmcp import FastMCP
import os
import sys

# Define path to MCP servers
sys.path.append("/mnt/e/genesis-system/mcp-servers")

# Try to import call_model, but handle the case where fastmcp isn't fully set up for direct import
try:
    from openrouter_swarm import call_model
except ImportError:
    # If direct import fails (e.g. if fastmcp wraps it weirdly), 
    # we might need to rely on the running MCP or duplicate the logic.
    # Duplicating logic here for robustness in this script context.
    import requests
    import logging
    
    OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
    OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"

    MODELS = {
        "deepseek-r1": "deepseek/deepseek-r1",
        "kimi": "moonshot/kimi-k2.5",
        "minimax": "minimax/minimax-01"
    }

    def call_model(model_key: str, prompt: str, system_prompt: str = "") -> str:
        if not OPENROUTER_API_KEY:
            return "Error: OPENROUTER_API_KEY not found."
        
        model_id = MODELS.get(model_key, model_key)
        
        headers = {
            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": model_id,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ]
        }
        
        try:
            response = requests.post(OPENROUTER_URL, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            return result['choices'][0]['message']['content']
        except Exception as e:
            print(f"Error calling {model_key}: {e}")
            return f"Error: {str(e)}"

def execute_swarm_plan():
    print("--- 🚀 SWARM ARCHITECTURE INITIATED ---")
    
    # 1. DEEPSEEK R1: Architecture Plan
    print("\n🧠 Agent DeepSeek R1 (Thinking)...")
    objective = """
    Design a production-ready MVP for the 'Genesis Patent OS' (Powered by Sunaiva) in 1 hour.
    
    Core Components:
    1. 'patent_os_db' (Postgres): Schema for Immutable Audit Logs (Patent 4) and Risk Events (Patent 3).
    2. 'Sunaiva Shield Widget': A JavaScript widget for client websites that shows live validation status.
    3. API Integration: How to connect the existing FastAPI to Postgres.
    
    Constraints: 
    - Use SQLAlchemy for ORM.
    - Widget must be lightweight (vanilla JS).
    - API must be robust.
    """
    plan = call_model("deepseek-r1", objective)
    with open("docs/GENESIS PATENTS/SWARM_ARCHITECTURE_PLAN.md", "w") as f:
        f.write(plan)
    print("✅ Architecture Plan Generated.")

    # 2. KIMI K1.5: Backend Code
    print("\n💻 Agent Kimi (Coding Backend)...")
    backend_task = """
    Based on the architecture plan (assume standard rigorous design):
    Write the `database.py` and `models.py` for `patent-os` using SQLAlchemy and Pydantic.
    
    Models needed:
    - AuditLog (id, timestamp, event_type, details, user_id, prev_hash, current_hash)
    - RiskAssessment (id, proposal_hash, risk_score, details)
    
    Output ONLY the python code for `models.py` and `database.py`.
    """
    backend_code = call_model("kimi", backend_task)
    with open("mcp-servers/patent-os/models.py", "w") as f:
        f.write(backend_code)
    print("✅ Backend Models Generated.")
    
    # 3. MINIMAX: Frontend Widget
    print("\n🎨 Agent MiniMax (Creative Frontend)...")
    frontend_task = """
    Create a beautiful, trustworthy 'Powered by Sunaiva' badge widget.
    
    Features:
    - HTML/CSS/JS in one file (or separate).
    - Shows 'Shield Active' status with a green pulsing dot.
    - Hover effect: '9-Layer Protection Active'.
    - Link to 'https://sunaiva.com/compliance'.
    
    Style: Professional, clean, banking-grade trust.
    Output the HTML/CSS/JS code.
    """
    frontend_code = call_model("minimax", frontend_task)
    with open("mcp-servers/patent-os/widget/shield_badge.html", "w") as f:
        f.write(frontend_code)
    print("✅ Frontend Widget Generated.")

if __name__ == "__main__":
    execute_swarm_plan()