backend/api/agent_api.py

"""
SAAP Agent Management API
FastAPI endpoints for modular agent management
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict, Any, Optional
import asyncio
from datetime import datetime
import uvicorn

# Import SAAP models
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))

from models.agent import SaapAgent, AgentTemplates, AgentStatus, AgentType
from api.colossus_client import ColossusClient

# FastAPI App
app = FastAPI(
    title="SAAP Agent Management API",
    description="Modular AI Agent Management for SAAP Platform",
    version="1.0.0"
)

# CORS for Vue.js frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173", "http://localhost:8080", "*"],  # Vue.js dev server
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# In-memory agent storage (later: replace with database)
agents_db: Dict[str, SaapAgent] = {}

@app.on_event("startup")
async def startup_event():
    """Initialize SAAP with default agents"""
    print("🚀 Initializing SAAP Agent Management API...")
    
    # Load default Alesi agents
    default_agents = [
        AgentTemplates.jane_alesi(),
        AgentTemplates.john_alesi(), 
        AgentTemplates.lara_alesi()
    ]
    
    for agent in default_agents:
        agents_db[agent.id] = agent
        print(f"✅ Loaded: {agent.name} ({agent.type.value})")
    
    print(f"🎉 SAAP initialized with {len(agents_db)} agents")

# Agent Management Endpoints
@app.get("/")
async def root():
    """API root with status info"""
    return {
        "message": "SAAP Agent Management API",
        "version": "1.0.0",
        "agents_count": len(agents_db),
        "status": "operational"
    }

@app.get("/agents", response_model=List[Dict[str, Any]])
async def list_agents():
    """List all registered agents"""
    return [agent.to_dict() for agent in agents_db.values()]

@app.get("/agents/{agent_id}")
async def get_agent(agent_id: str):
    """Get specific agent details"""
    if agent_id not in agents_db:
        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
    
    agent = agents_db[agent_id]
    return {
        "agent": agent.to_dict(),
        "is_active": agent.is_active(),
        "capabilities_display": agent.get_capabilities_display()
    }

@app.post("/agents")
async def create_agent(agent_data: Dict[str, Any]):
    """Create new agent from JSON data"""
    try:
        # Validate and create agent
        agent = SaapAgent.from_dict(agent_data)
        
        # Check if agent already exists
        if agent.id in agents_db:
            raise HTTPException(status_code=400, detail=f"Agent '{agent.id}' already exists")
        
        # Store agent
        agents_db[agent.id] = agent
        
        return {
            "message": f"Agent '{agent.name}' created successfully",
            "agent": agent.to_dict()
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid agent data: {str(e)}")

@app.put("/agents/{agent_id}")
async def update_agent(agent_id: str, agent_data: Dict[str, Any]):
    """Update existing agent"""
    if agent_id not in agents_db:
        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
    
    try:
        # Create updated agent
        updated_agent = SaapAgent.from_dict(agent_data)
        updated_agent.updated_at = datetime.utcnow()
        
        # Store updated agent
        agents_db[agent_id] = updated_agent
        
        return {
            "message": f"Agent '{agent_id}' updated successfully",
            "agent": updated_agent.to_dict()
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid agent data: {str(e)}")

@app.delete("/agents/{agent_id}")
async def delete_agent(agent_id: str):
    """Delete agent"""
    if agent_id not in agents_db:
        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
    
    # Stop agent if running
    agent = agents_db[agent_id]
    if agent.is_active():
        agent.update_status(AgentStatus.INACTIVE)
    
    # Remove from database
    del agents_db[agent_id]
    
    return {"message": f"Agent '{agent_id}' deleted successfully"}

# Agent Control Endpoints
@app.post("/agents/{agent_id}/start")
async def start_agent(agent_id: str, background_tasks: BackgroundTasks):
    """Start agent (set status to active)"""
    if agent_id not in agents_db:
        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
    
    agent = agents_db[agent_id]
    
    if agent.is_active():
        raise HTTPException(status_code=400, detail=f"Agent '{agent_id}' is already active")
    
    # Update status to starting
    agent.update_status(AgentStatus.STARTING)
    
    # Background task to set to active (simulate startup time)
    def activate_agent():
        asyncio.run(finalize_agent_start(agent_id))
    
    background_tasks.add_task(activate_agent)
    
    return {
        "message": f"Agent '{agent.name}' is starting...",
        "agent_id": agent_id,
        "status": agent.status.value
    }

async def finalize_agent_start(agent_id: str):
    """Finalize agent startup"""
    await asyncio.sleep(2)  # Simulate startup time
    
    if agent_id in agents_db:
        agent = agents_db[agent_id]
        agent.update_status(AgentStatus.ACTIVE)
        agent.update_metrics(messages_processed=0, average_response_time=0.0)
        print(f"✅ Agent '{agent.name}' activated")

@app.post("/agents/{agent_id}/stop") 
async def stop_agent(agent_id: str):
    """Stop agent (set status to inactive)"""
    if agent_id not in agents_db:
        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
    
    agent = agents_db[agent_id]
    
    if not agent.is_active():
        raise HTTPException(status_code=400, detail=f"Agent '{agent_id}' is not active")
    
    # Update status
    agent.update_status(AgentStatus.INACTIVE)
    
    return {
        "message": f"Agent '{agent.name}' stopped",
        "agent_id": agent_id,
        "status": agent.status.value
    }

# Agent Communication Endpoints
@app.post("/agents/{agent_id}/chat")
async def chat_with_agent(agent_id: str, message_data: Dict[str, Any]):
    """Send message to agent and get response"""
    if agent_id not in agents_db:
        raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
    
    agent = agents_db[agent_id]
    
    if not agent.is_active():
        raise HTTPException(status_code=400, detail=f"Agent '{agent_id}' is not active")
    
    try:
        user_message = message_data.get("message", "")
        if not user_message:
            raise HTTPException(status_code=400, detail="Message content required")
        
        # Prepare messages for colossus
        messages = []
        
        # Add system prompt if available
        if agent.personality and agent.personality.system_prompt:
            messages.append({
                "role": "system", 
                "content": agent.personality.system_prompt
            })
        
        # Add user message
        messages.append({
            "role": "user",
            "content": user_message
        })
        
        # Send to colossus
        async with ColossusClient() as client:
            result = await client.chat_completion(
                messages=messages,
                agent_id=agent_id,
                temperature=agent.llm_config.temperature,
                max_tokens=agent.llm_config.max_tokens
            )
            
            if result["success"]:
                response_text = result["response"]["choices"][0]["message"]["content"]
                response_time = result["response_time"]
                
                # Update agent metrics
                agent.update_metrics(
                    messages_processed=agent.metrics.messages_processed + 1 if agent.metrics else 1,
                    average_response_time=response_time
                )
                
                return {
                    "agent_id": agent_id,
                    "agent_name": agent.name,
                    "user_message": user_message,
                    "agent_response": response_text,
                    "response_time": response_time,
                    "model": agent.llm_config.model
                }
            else:
                raise HTTPException(status_code=500, detail=f"Agent communication failed: {result['error']}")
                
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Chat error: {str(e)}")

# System Status Endpoints
@app.get("/system/status")
async def system_status():
    """Get system status and metrics"""
    active_agents = [agent for agent in agents_db.values() if agent.is_active()]
    inactive_agents = [agent for agent in agents_db.values() if not agent.is_active()]
    
    return {
        "system": "SAAP Agent Management",
        "status": "operational",
        "agents": {
            "total": len(agents_db),
            "active": len(active_agents),
            "inactive": len(inactive_agents)
        },
        "active_agents": [{"id": agent.id, "name": agent.name, "type": agent.type.value} 
                         for agent in active_agents],
        "timestamp": datetime.utcnow().isoformat()
    }

# Template Endpoints
@app.get("/templates")
async def list_agent_templates():
    """List available agent templates"""
    templates = [
        {
            "id": "jane_alesi",
            "name": "Jane Alesi", 
            "type": "coordinator",
            "description": "Lead AI Architect Template"
        },
        {
            "id": "john_alesi", 
            "name": "John Alesi",
            "type": "developer", 
            "description": "Software Developer Template"
        },
        {
            "id": "lara_alesi",
            "name": "Lara Alesi",
            "type": "specialist",
            "description": "Medical Specialist Template"
        }
    ]
    
    return {
        "templates": templates,
        "count": len(templates)
    }

@app.post("/templates/{template_id}/create")
async def create_agent_from_template(template_id: str, customization: Optional[Dict[str, Any]] = None):
    """Create agent from template with optional customization"""
    
    template_functions = {
        "jane_alesi": AgentTemplates.jane_alesi,
        "john_alesi": AgentTemplates.john_alesi,
        "lara_alesi": AgentTemplates.lara_alesi
    }
    
    if template_id not in template_functions:
        raise HTTPException(status_code=404, detail=f"Template '{template_id}' not found")
    
    # Create agent from template
    agent = template_functions[template_id]()
    
    # Apply customization if provided
    if customization:
        if "name" in customization:
            agent.name = customization["name"]
        if "color" in customization:
            agent.color = customization["color"]
        if "description" in customization:
            agent.description = customization["description"]
    
    # Generate unique ID if agent exists
    original_id = agent.id
    counter = 1
    while agent.id in agents_db:
        agent.id = f"{original_id}_{counter}"
        counter += 1
    
    # Store agent
    agents_db[agent.id] = agent
    
    return {
        "message": f"Agent '{agent.name}' created from template",
        "agent": agent.to_dict()
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)