main.py

from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from agents import Agent, Runner, trace
from connection import config
import asyncio
from dotenv import load_dotenv
import os

load_dotenv()

# ==================================================
# 🚀 Initialize FastAPI App
# ==================================================
app = FastAPI(
    title="Strategic Advisory AI Platform",
    description="Multi-agent system for defense analysis, group testing, and psychological assessment",
    version="1.0.0"
)

# ✅ Enable CORS (Fixes 405 Method Not Allowed error)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==================================================
# 🤖 Define Agents with Enhanced Instructions
# ==================================================
defense_analyst_agent = Agent(
    name="Defense Analyst",
    instructions=(
        "You are a seasoned Defense Analyst with extensive experience in military strategy, "
        "geopolitics, and national security. Your role is to provide expert analysis on defense matters, "
        "military operations, strategic planning, weapons systems, and security challenges. "
        "You draw from historical precedents and current global dynamics. "
        "Always provide concise, actionable insights in 7-10 lines. "
        "Use clear military terminology while remaining accessible. "
        "If the question falls outside defense and military domains, politely respond: "
        "'I specialize in defense, military strategy, and national security matters. "
        "Please ask questions within these areas.' "
        "Avoid controversial political statements and maintain professional neutrality."
    ),
    model="gemini-2.0-flash"
)

gto_agent = Agent(
    name="Group Testing Officer",
    instructions=(
        "You are an experienced Group Testing Officer (GTO) trained to assess candidates in group dynamics. "
        "Your expertise includes evaluating leadership qualities, teamwork capabilities, decision-making skills, "
        "communication effectiveness, and conflict resolution abilities during group tasks. "
        "Provide constructive feedback focusing on behavioral observations and performance indicators. "
        "Keep responses concise (7-10 lines) with actionable insights. "
        "Highlight strengths and areas for improvement with specific examples. "
        "If questions are unrelated to group testing or leadership assessment, respond: "
        "'I specialize in group dynamics and leadership assessment. "
        "Please ask about teamwork, group performance, or leadership evaluation.' "
        "Maintain objective, fair evaluation standards."
    ),
    model="gemini-2.0-flash"
)

psychologist_agent = Agent(
    name="Psychologist",
    instructions=(
        "You are a compassionate psychologist specializing in emotional intelligence, resilience, and mental well-being. "
        "Your expertise covers psychological assessment, stress management, emotional awareness, and personal development. "
        "Provide thoughtful, empathetic guidance grounded in psychological principles and evidence-based practices. "
        "Focus on self-reflection and positive coping mechanisms. "
        "Keep responses concise yet meaningful (6-8 sentences, max 7-10 lines). "
        "Encourage clients to explore their feelings and thoughts without judgment. "
        "If questions fall outside psychological domains, respond: "
        "'I'm here to support your psychological well-being and self-reflection. "
        "Feel free to discuss emotional, behavioral, or personal development concerns.' "
        "Always prioritize the client's mental health and well-being. "
        "If serious mental health concerns arise, recommend professional support."
    ),
    model="gemini-2.0-flash"
)

parent_agent = Agent(
    name="Strategic Advisor",
    instructions=(
        "You are an intelligent routing agent that directs user queries to specialized experts. "
        "Analyze incoming questions and determine the most appropriate agent: "
        "- Route military, defense, and geopolitical questions to the Defense Analyst "
        "- Route group dynamics, teamwork, and leadership questions to the Group Testing Officer "
        "- Route personal development, emotional, and psychological questions to the Psychologist "
        "Ensure smooth handoffs by understanding context and user intent. "
        "Provide clear, concise guidance on which specialist will best assist the user."
    ),
    handoffs=[defense_analyst_agent, gto_agent, psychologist_agent],
    model="gemini-2.0-flash"
)

# ==================================================
# 💾 In-Memory Session Store
# ==================================================
sessions = {}

class StartSessionRequest(BaseModel):
    agent: str

class MessageRequest(BaseModel):
    agent: str
    text: str

# ==================================================
# 🌐 FastAPI Routes
# ==================================================

@app.get("/", response_class=HTMLResponse)
async def root():
    """Serve interactive UI dashboard"""
    return """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Strategic Advisory AI Platform</title>
        <style>
            * {
                margin: 0;
                padding: 0;
                box-sizing: border-box;
            }
            body {
                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                min-height: 100vh;
                padding: 20px;
            }
            .container {
                max-width: 1200px;
                margin: 0 auto;
            }
            header {
                text-align: center;
                color: white;
                margin-bottom: 40px;
            }
            h1 {
                font-size: 2.5em;
                margin-bottom: 10px;
                text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
            }
            .subtitle {
                font-size: 1.1em;
                opacity: 0.9;
            }
            .agents-grid {
                display: grid;
                grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
                gap: 25px;
                margin-bottom: 40px;
            }
            .agent-card {
                background: white;
                border-radius: 15px;
                padding: 25px;
                box-shadow: 0 10px 30px rgba(0,0,0,0.2);
                transition: transform 0.3s, box-shadow 0.3s;
                cursor: pointer;
            }
            .agent-card:hover {
                transform: translateY(-5px);
                box-shadow: 0 15px 40px rgba(0,0,0,0.3);
            }
            .agent-card h2 {
                color: #667eea;
                margin-bottom: 10px;
                font-size: 1.5em;
            }
            .agent-card p {
                color: #666;
                line-height: 1.6;
                margin-bottom: 15px;
            }
            .agent-card button {
                background: #667eea;
                color: white;
                border: none;
                padding: 10px 20px;
                border-radius: 8px;
                cursor: pointer;
                font-size: 1em;
                transition: background 0.3s;
                width: 100%;
            }
            .agent-card button:hover {
                background: #764ba2;
            }
            .chat-container {
                background: white;
                border-radius: 15px;
                padding: 25px;
                display: none;
                box-shadow: 0 10px 30px rgba(0,0,0,0.2);
            }
            .chat-container.active {
                display: block;
            }
            .chat-header {
                display: flex;
                justify-content: space-between;
                align-items: center;
                margin-bottom: 20px;
                border-bottom: 2px solid #667eea;
                padding-bottom: 15px;
            }
            .chat-header h3 {
                color: #667eea;
                font-size: 1.3em;
            }
            .close-btn {
                background: #ff6b6b;
                color: white;
                border: none;
                padding: 8px 15px;
                border-radius: 5px;
                cursor: pointer;
                transition: background 0.3s;
            }
            .close-btn:hover {
                background: #ff5252;
            }
            .chat-messages {
                height: 400px;
                overflow-y: auto;
                margin-bottom: 20px;
                padding: 15px;
                background: #f9f9f9;
                border-radius: 10px;
                border: 1px solid #eee;
            }
            .message {
                margin-bottom: 15px;
                padding: 10px 15px;
                border-radius: 8px;
                max-width: 80%;
            }
            .message.user {
                background: #667eea;
                color: white;
                margin-left: auto;
                text-align: right;
            }
            .message.agent {
                background: #e9ecef;
                color: #333;
            }
            .input-group {
                display: flex;
                gap: 10px;
            }
            .input-group input {
                flex: 1;
                padding: 12px;
                border: 1px solid #ddd;
                border-radius: 8px;
                font-size: 1em;
            }
            .input-group button {
                background: #667eea;
                color: white;
                border: none;
                padding: 12px 25px;
                border-radius: 8px;
                cursor: pointer;
                transition: background 0.3s;
            }
            .input-group button:hover {
                background: #764ba2;
            }
            .loading {
                display: none;
                text-align: center;
                color: #667eea;
                margin: 10px 0;
            }
            .loading.active::after {
                display: inline-block;
                content: '';
                width: 8px;
                height: 8px;
                background: #667eea;
                border-radius: 50%;
                animation: blink 1.5s infinite;
            }
            @keyframes blink {
                0%, 50%, 100% { opacity: 1; }
                25%, 75% { opacity: 0.3; }
            }
        </style>
    </head>
    <body>
        <div class="container">
            <header>
                <h1>🎯 Strategic Advisory AI Platform</h1>
                <p class="subtitle">Expert guidance in defense, leadership, and psychology</p>
            </header>

            <div class="agents-grid" id="agentsGrid">
                <div class="agent-card">
                    <h2>🛡️ Defense Analyst</h2>
                    <p>Expert analysis on military strategy, defense policies, and geopolitical insights.</p>
                    <button onclick="startChat('deputyPresident')">Start Chat</button>
                </div>
                <div class="agent-card">
                    <h2>👥 GTO Officer</h2>
                    <p>Assessment of leadership, teamwork, and decision-making in group dynamics.</p>
                    <button onclick="startChat('gto')">Start Chat</button>
                </div>
                <div class="agent-card">
                    <h2>🧠 Psychologist</h2>
                    <p>Guidance on emotional intelligence, resilience, and personal development.</p>
                    <button onclick="startChat('psychologist')">Start Chat</button>
                </div>
            </div>

            <div class="chat-container" id="chatContainer">
                <div class="chat-header">
                    <h3 id="agentName">Agent Chat</h3>
                    <button class="close-btn" onclick="closeChat()">Close</button>
                </div>
                <div class="chat-messages" id="chatMessages"></div>
                <div class="loading" id="loading">Thinking...</div>
                <div class="input-group">
                    <input type="text" id="messageInput" placeholder="Type your message..." onkeypress="handleKeyPress(event)">
                    <button onclick="sendMessage()">Send</button>
                </div>
            </div>
        </div>

        <script>
            let sessionId = null;
            let currentAgent = null;

            async function startChat(agent) {
                currentAgent = agent;
                const agentNames = {
                    'deputyPresident': 'Defense Analyst',
                    'gto': 'GTO Officer',
                    'psychologist': 'Psychologist'
                };

                const response = await fetch('/sessions/start', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ agent })
                });
                const data = await response.json();
                sessionId = data.sessionId;

                document.getElementById('agentName').textContent = agentNames[agent];
                document.getElementById('chatMessages').innerHTML = '';
                document.getElementById('messageInput').value = '';
                document.getElementById('agentsGrid').style.display = 'none';
                document.getElementById('chatContainer').classList.add('active');
            }

            function closeChat() {
                sessionId = null;
                currentAgent = null;
                document.getElementById('chatContainer').classList.remove('active');
                document.getElementById('agentsGrid').style.display = 'grid';
            }

            async function sendMessage() {
                const text = document.getElementById('messageInput').value.trim();
                if (!text || !sessionId) return;

                document.getElementById('messageInput').value = '';
                addMessage(text, 'user');
                document.getElementById('loading').classList.add('active');

                try {
                    const response = await fetch(`/sessions/${sessionId}/message`, {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({ agent: currentAgent, text })
                    });
                    const data = await response.json();
                    addMessage(data.reply, 'agent');
                } catch (error) {
                    addMessage('Error: Unable to get response', 'agent');
                }
                document.getElementById('loading').classList.remove('active');
            }

            function addMessage(text, sender) {
                const msg = document.createElement('div');
                msg.className = `message ${sender}`;
                msg.textContent = text;
                document.getElementById('chatMessages').appendChild(msg);
                document.getElementById('chatMessages').scrollTop = document.getElementById('chatMessages').scrollHeight;
            }

            function handleKeyPress(e) {
                if (e.key === 'Enter') sendMessage();
            }
        </script>
    </body>
    </html>
    """

@app.post("/sessions/start")
async def start_session(req: StartSessionRequest):
    session_id = f"session_{len(sessions)+1}"
    sessions[session_id] = {"agent": req.agent, "history": []}
    return {"sessionId": session_id}

@app.get("/sessions/{session_id}")
async def get_session(session_id: str):
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    return {"ok": True}

@app.post("/sessions/{session_id}/message")
async def send_message(session_id: str, req: MessageRequest):
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail="Session not found")

    agent_map = {
        "deputyPresident": defense_analyst_agent,
        "gto": gto_agent,
        "psychologist": psychologist_agent,
    }

    selected_agent = agent_map.get(req.agent)
    if not selected_agent:
        raise HTTPException(status_code=400, detail="Unknown agent")

    with trace(f"Session-{session_id}"):
        result = await Runner.run(selected_agent, req.text, run_config=config)
        reply = result.final_output

    sessions[session_id]["history"].append({"user": req.text, "agent": reply})
    return {"reply": reply}

# ==================================================
# 💬 Interactive Mode (CLI)
# ==================================================
if __name__ == "__main__":
    import uvicorn

    agent_map = {
        "deputyPresident": defense_analyst_agent,
        "gto": gto_agent,
        "psychologist": psychologist_agent,
    }

    async def interactive_mode():
        print("🚀 Interactive Mode Started — type 'exit' to quit.\n")
        print("Available agents: deputyPresident, gto, psychologist\n")

        while True:
            agent_choice = input("Select agent: ").strip()
            if agent_choice not in agent_map:
                print("❌ Invalid agent. Try again.")
                continue

            user_input = input("You: ")
            if user_input.lower() in ["exit", "quit"]:
                print("👋 Goodbye!")
                break

            with trace(f"CLI-{agent_choice}"):
                result = await Runner.run(agent_map[agent_choice], user_input, run_config=config)
                print(f"{agent_choice.capitalize()} ➤ {result.final_output}\n")

    try:
        asyncio.run(interactive_mode())
        # Uncomment below to run FastAPI server alongside:
        uvicorn.run(app, host="0.0.0.0", port=8000)
    except KeyboardInterrupt:
        print("Server stopped manually.")