"""FastAPI backend for Companion AI."""

from __future__ import annotations

import json
import uuid
from contextlib import asynccontextmanager
from typing import AsyncGenerator

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse

from companion.config import Config, load_config
from companion.memory import SessionMemory
from companion.orchestrator import ChatOrchestrator
from companion.rag.search import SearchEngine
from companion.rag.vector_store import VectorStore


class ChatRequest(BaseModel):
    """Chat request model."""

    message: str
    session_id: str | None = None
    stream: bool = True
    use_rag: bool = True


class ChatResponse(BaseModel):
    """Chat response model (non-streaming)."""

    response: str
    session_id: str
    sources: list[dict] | None = None


# Global instances
config: Config
vector_store: VectorStore
search_engine: SearchEngine
memory: SessionMemory
orchestrator: ChatOrchestrator
http_client: httpx.AsyncClient


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    """Manage application lifespan."""
    global config, vector_store, search_engine, memory, orchestrator, http_client

    # Startup
    config = load_config("config.json")
    vector_store = VectorStore(
        uri=config.rag.vector_store.path, dimensions=config.rag.embedding.dimensions
    )
    search_engine = SearchEngine(
        vector_store=vector_store,
        embedder_base_url=config.rag.embedding.base_url,
        embedder_model=config.rag.embedding.model,
        embedder_batch_size=config.rag.embedding.batch_size,
        default_top_k=config.rag.search.default_top_k,
        similarity_threshold=config.rag.search.similarity_threshold,
        hybrid_search_enabled=config.rag.search.hybrid_search.enabled,
    )
    memory = SessionMemory(db_path=config.companion.memory.persistent_store)
    http_client = httpx.AsyncClient(timeout=300.0)
    orchestrator = ChatOrchestrator(
        config=config,
        search_engine=search_engine,
        memory=memory,
        http_client=http_client,
    )

    yield

    # Shutdown
    await http_client.aclose()
    memory.close()


app = FastAPI(
    title="Companion AI",
    description="Personal AI companion with RAG",
    version="0.1.0",
    lifespan=lifespan,
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=config.api.cors_origins
    if "config" in globals()
    else ["http://localhost:5173"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
async def health_check() -> dict:
    """Health check endpoint."""
    return {
        "status": "healthy",
        "version": "0.1.0",
        "indexed_chunks": vector_store.count() if "vector_store" in globals() else 0,
    }


@app.post("/chat")
async def chat(request: ChatRequest) -> EventSourceResponse:
    """Chat endpoint with SSE streaming."""
    if not request.message.strip():
        raise HTTPException(status_code=400, detail="Message cannot be empty")

    # Generate or use existing session ID
    session_id = request.session_id or str(uuid.uuid4())

    async def event_generator() -> AsyncGenerator[str, None]:
        """Generate SSE events."""
        try:
            # Stream response from orchestrator
            async for chunk in orchestrator.chat(
                session_id=session_id,
                user_message=request.message,
                use_rag=request.use_rag,
            ):
                # Parse the SSE data format
                if chunk.startswith("data: "):
                    data = chunk[6:]
                    if data == "[DONE]":
                        yield json.dumps({"type": "done", "session_id": session_id})
                    else:
                        try:
                            parsed = json.loads(data)
                            if "content" in parsed:
                                yield json.dumps(
                                    {"type": "chunk", "content": parsed["content"]}
                                )
                            elif "citations" in parsed:
                                yield json.dumps(
                                    {
                                        "type": "citations",
                                        "citations": parsed["citations"],
                                    }
                                )
                            elif "error" in parsed:
                                yield json.dumps(
                                    {"type": "error", "message": parsed["error"]}
                                )
                            else:
                                yield data
                        except json.JSONDecodeError:
                            # Pass through raw data
                            yield data
                else:
                    yield chunk

        except Exception as e:
            yield json.dumps({"type": "error", "message": str(e)})

    return EventSourceResponse(
        event_generator(),
        headers={"X-Session-ID": session_id},
    )


@app.get("/sessions/{session_id}/history")
async def get_session_history(session_id: str) -> dict:
    """Get conversation history for a session."""
    history = memory.get_history(session_id)
    return {
        "session_id": session_id,
        "messages": [
            {
                "role": msg.role,
                "content": msg.content,
                "timestamp": msg.timestamp.isoformat(),
            }
            for msg in history
        ],
    }


class ReloadModelRequest(BaseModel):
    """Model reload request."""

    model_path: str


@app.post("/admin/reload-model")
async def reload_model_endpoint(request: ReloadModelRequest) -> dict:
    """Reload the model with a new fine-tuned version (admin only)."""
    from pathlib import Path

    from companion.forge.reload import reload_model

    new_path = Path(request.model_path).expanduser()

    if not new_path.exists():
        raise HTTPException(status_code=404, detail=f"Model not found: {new_path}")

    try:
        active_path = reload_model(config, new_path, backup=True)
        return {
            "status": "success",
            "message": f"Model reloaded successfully",
            "active_model": str(active_path),
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to reload model: {e}")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=7373)
