"""FastAPI application with WebSocket chat support."""

from __future__ import annotations

import asyncio
import json
import uuid
from contextlib import asynccontextmanager
from typing import Any

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from pathlib import Path

from agenai.agent import AgentConfig, AgentContext
from agenai.agent.llm_ollama import OllamaProvider
from agenai.agents.teaching_assistant import TeachingAssistantAgent


# Store active connections
class ConnectionManager:
    """Manage WebSocket connections."""

    def __init__(self):
        self.active_connections: dict[str, WebSocket] = {}
        self.agents: dict[str, TeachingAssistantAgent] = {}

    async def connect(self, websocket: WebSocket, client_id: str):
        await websocket.accept()
        self.active_connections[client_id] = websocket

        # Create agent for this client
        try:
            ollama = OllamaProvider(model="llama3.1")
            if not ollama.is_available():
                ollama = OllamaProvider(model="mistral")

            self.agents[client_id] = TeachingAssistantAgent(
                config=AgentConfig(max_steps=10, temperature=0.7),
                llm=ollama,
            )
        except Exception as e:
            # Will handle in message processing
            pass

    def disconnect(self, client_id: str):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
        if client_id in self.agents:
            del self.agents[client_id]

    async def send_message(self, client_id: str, message: dict):
        if client_id in self.active_connections:
            await self.active_connections[client_id].send_json(message)

    def get_agent(self, client_id: str) -> TeachingAssistantAgent | None:
        return self.agents.get(client_id)


manager = ConnectionManager()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler."""
    # Startup
    print("Starting Agenai Server...")
    yield
    # Shutdown
    print("Shutting down Agenai Server...")


def create_app() -> FastAPI:
    """Create and configure the FastAPI application."""
    app = FastAPI(
        title="Agenai Server",
        description="AI Agent API with WebSocket chat support",
        version="0.1.0",
        lifespan=lifespan,
    )

    # CORS middleware
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    return app


app = create_app()


# Request/Response models
class ChatRequest(BaseModel):
    message: str
    session_id: str | None = None


class ChatResponse(BaseModel):
    response: str
    session_id: str
    success: bool
    error: str | None = None


class HealthResponse(BaseModel):
    status: str
    ollama_available: bool
    models: list[str]


# REST Endpoints
@app.get("/")
async def root():
    """Serve the chat frontend."""
    frontend_path = Path(__file__).parent / "static" / "index.html"
    if frontend_path.exists():
        return FileResponse(frontend_path)
    return {"message": "Agenai Server", "docs": "/docs", "websocket": "/ws/{client_id}"}


@app.get("/health", response_model=HealthResponse)
async def health():
    """Health check endpoint."""
    try:
        ollama = OllamaProvider()
        available = ollama.is_available()
        models = ollama.list_models() if available else []
    except Exception:
        available = False
        models = []

    return HealthResponse(
        status="healthy",
        ollama_available=available,
        models=models,
    )


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """REST endpoint for chat (alternative to WebSocket)."""
    session_id = request.session_id or str(uuid.uuid4())

    try:
        # Check Ollama availability
        ollama = OllamaProvider(model="llama3.1")
        if not ollama.is_available():
            raise HTTPException(
                status_code=503,
                detail="Ollama is not running. Please start Ollama first.",
            )

        # Create or get agent
        agent = TeachingAssistantAgent(
            config=AgentConfig(max_steps=10),
            llm=ollama,
        )

        # Run agent
        result = agent.run(request.message)

        return ChatResponse(
            response=result.output if result.success else f"Error: {result.error}",
            session_id=session_id,
            success=result.success,
            error=result.error,
        )

    except Exception as e:
        return ChatResponse(
            response="",
            session_id=session_id,
            success=False,
            error=str(e),
        )


# WebSocket Endpoint
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    """WebSocket endpoint for real-time chat."""
    await manager.connect(websocket, client_id)

    # Send welcome message
    await manager.send_message(client_id, {
        "type": "connected",
        "message": "Connected to Teaching Assistant. Ask me anything about Python programming!",
        "client_id": client_id,
    })

    try:
        while True:
            # Receive message
            data = await websocket.receive_json()
            message = data.get("message", "")

            if not message:
                continue

            # Send thinking indicator
            await manager.send_message(client_id, {
                "type": "thinking",
                "message": "Thinking...",
            })

            # Get agent and process
            agent = manager.get_agent(client_id)

            if agent is None:
                # Try to create agent
                try:
                    ollama = OllamaProvider(model="llama3.1")
                    if not ollama.is_available():
                        await manager.send_message(client_id, {
                            "type": "error",
                            "message": "Ollama is not running. Please install and start Ollama: https://ollama.ai",
                        })
                        continue

                    agent = TeachingAssistantAgent(
                        config=AgentConfig(max_steps=10),
                        llm=ollama,
                    )
                    manager.agents[client_id] = agent
                except Exception as e:
                    await manager.send_message(client_id, {
                        "type": "error",
                        "message": f"Failed to initialize agent: {e}",
                    })
                    continue

            # Run agent in thread pool (blocking operation)
            try:
                loop = asyncio.get_event_loop()
                result = await loop.run_in_executor(
                    None,
                    lambda: agent.run(message),
                )

                # Send response
                await manager.send_message(client_id, {
                    "type": "response",
                    "message": result.output if result.success else f"Error: {result.error}",
                    "success": result.success,
                    "steps": len(result.steps),
                    "tokens": result.total_tokens,
                })

            except Exception as e:
                await manager.send_message(client_id, {
                    "type": "error",
                    "message": f"Error processing request: {e}",
                })

    except WebSocketDisconnect:
        manager.disconnect(client_id)
