from __future__ import annotations

from dotenv import load_dotenv

load_dotenv()  # Load .env before any other imports that may use env vars

from fastapi import FastAPI, WebSocket
from fastapi.middleware.cors import CORSMiddleware

from prompt_auto_lab.api.routes import sessions, versions, prompts, badcases, auth, import_session, copilot, local
from prompt_auto_lab.api.websocket.handler import websocket_handler

app = FastAPI(
    title="AutoPromptLab2 API",
    description="API for the multi-agent prompt optimization system",
    version="0.1.0",
)

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

# Include routers
app.include_router(sessions.router)
app.include_router(versions.router)
app.include_router(prompts.router)
app.include_router(badcases.router)
app.include_router(auth.router)
app.include_router(import_session.router)
app.include_router(copilot.router)
app.include_router(local.router)


@app.get("/")
async def root():
    """Root endpoint."""
    return {
        "name": "AutoPromptLab2 API",
        "version": "0.1.0",
        "docs": "/docs",
    }


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {"status": "healthy"}


@app.websocket("/ws/sessions/{session_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str, user_id: str | None = None):
    """WebSocket endpoint for real-time session updates.

    Args:
        websocket: The WebSocket connection
        session_id: The session to connect to
        user_id: Optional user ID for ownership verification (query param)
    """
    await websocket_handler(websocket, session_id, user_id)


if __name__ == "__main__":
    import uvicorn
    import os

    host = os.getenv("AUTOPROMPTLAB_HOST", "localhost")
    port_str = os.getenv("AUTOPROMPTLAB_PORT", "8237")
    try:
        port = int(port_str)
    except ValueError:
        port = 8237

    uvicorn.run(app, host=host, port=port)
