app/api/v1/endpoints.py

"""
API v1 endpoints for Sema Chat API
"""

import time
from typing import List, Optional, Dict, Any
from datetime import datetime

from fastapi import APIRouter, HTTPException, Request, Query, WebSocket, WebSocketDisconnect
from sse_starlette.sse import EventSourceResponse
from slowapi import Limiter
from slowapi.util import get_remote_address
from prometheus_client import generate_latest, CONTENT_TYPE_LATEST
from fastapi.responses import Response

from ...models.schemas import (
    ChatRequest, ChatResponse, ConversationHistory,
    HealthResponse, ModelInfo
)
from ...services.chat_manager import get_chat_manager
from ...services.model_manager import get_model_manager
from ...services.model_backends.base import ModelNotLoadedError, GenerationError
from ...core.config import settings
from ...core.logging import get_logger

# Initialize router and rate limiter
router = APIRouter()
limiter = Limiter(key_func=get_remote_address)
logger = get_logger()

# WebSocket connection manager
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

manager = ConnectionManager()


@router.post("/chat", response_model=ChatResponse)
@limiter.limit(f"{settings.rate_limit}/minute")
async def chat(chat_request: ChatRequest, request: Request):
    """
    Send a chat message and get a complete response
    """
    start_time = time.time()

    try:
        chat_manager = await get_chat_manager()
        response = await chat_manager.process_chat_request(chat_request)

        # Add timing information
        total_time = time.time() - start_time
        response.generation_time = getattr(response, 'generation_time', total_time)

        logger.info("chat_request_completed",
                   session_id=chat_request.session_id,
                   message_length=len(chat_request.message),
                   response_length=len(response.message),
                   total_time=total_time)

        return response

    except ModelNotLoadedError as e:
        logger.error("model_not_loaded", error=str(e), session_id=chat_request.session_id)
        raise HTTPException(status_code=503, detail="Model not available")

    except GenerationError as e:
        logger.error("generation_error", error=str(e), session_id=chat_request.session_id)
        raise HTTPException(status_code=500, detail="Failed to generate response")

    except Exception as e:
        logger.error("chat_request_failed", error=str(e), session_id=chat_request.session_id)
        raise HTTPException(status_code=500, detail="Internal server error")


@router.get("/chat/stream")
@limiter.limit(f"{settings.rate_limit}/minute")
async def chat_stream(
    request: Request,
    message: str = Query(..., description="Chat message"),
    session_id: str = Query(..., description="Session ID"),
    system_prompt: Optional[str] = Query(None, description="Custom system prompt"),
    temperature: Optional[float] = Query(None, ge=0.0, le=1.0, description="Temperature"),
    max_tokens: Optional[int] = Query(None, ge=1, le=2048, description="Max tokens")
):
    """
    Send a chat message and get a streaming response via Server-Sent Events
    """
    try:
        # Create chat request
        chat_request = ChatRequest(
            message=message,
            session_id=session_id,
            system_prompt=system_prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True
        )

        chat_manager = await get_chat_manager()

        async def event_generator():
            try:
                async for chunk in chat_manager.process_streaming_chat_request(chat_request):
                    # Format as SSE event
                    chunk_data = {
                        "content": chunk.content,
                        "session_id": chunk.session_id,
                        "message_id": chunk.message_id,
                        "chunk_id": chunk.chunk_id,
                        "is_final": chunk.is_final,
                        "timestamp": chunk.timestamp.isoformat()
                    }

                    yield {
                        "event": "chunk",
                        "data": chunk_data
                    }

                    if chunk.is_final:
                        yield {
                            "event": "done",
                            "data": {"message": "Stream completed"}
                        }
                        break

            except Exception as e:
                logger.error("streaming_error", error=str(e), session_id=session_id)
                yield {
                    "event": "error",
                    "data": {"error": str(e)}
                }

        return EventSourceResponse(event_generator())

    except Exception as e:
        logger.error("stream_setup_failed", error=str(e), session_id=session_id)
        raise HTTPException(status_code=500, detail="Failed to setup stream")


@router.websocket("/chat/ws")
async def websocket_chat(websocket: WebSocket):
    """
    WebSocket endpoint for real-time chat
    """
    await manager.connect(websocket)
    session_id = None

    try:
        while True:
            # Receive message from client
            data = await websocket.receive_json()

            # Extract request data
            message = data.get("message")
            session_id = data.get("session_id")
            system_prompt = data.get("system_prompt")
            temperature = data.get("temperature")
            max_tokens = data.get("max_tokens")

            if not message or not session_id:
                await websocket.send_json({
                    "error": "Message and session_id are required"
                })
                continue

            # Create chat request
            chat_request = ChatRequest(
                message=message,
                session_id=session_id,
                system_prompt=system_prompt,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=True
            )

            # Process streaming request
            chat_manager = await get_chat_manager()

            try:
                async for chunk in chat_manager.process_streaming_chat_request(chat_request):
                    await websocket.send_json({
                        "type": "chunk",
                        "content": chunk.content,
                        "session_id": chunk.session_id,
                        "message_id": chunk.message_id,
                        "chunk_id": chunk.chunk_id,
                        "is_final": chunk.is_final,
                        "timestamp": chunk.timestamp.isoformat()
                    })

                    if chunk.is_final:
                        break

            except Exception as e:
                logger.error("websocket_generation_error", error=str(e), session_id=session_id)
                await websocket.send_json({
                    "type": "error",
                    "error": str(e)
                })

    except WebSocketDisconnect:
        manager.disconnect(websocket)
        logger.info("websocket_disconnected", session_id=session_id)
    except Exception as e:
        logger.error("websocket_error", error=str(e), session_id=session_id)
        manager.disconnect(websocket)


@router.get("/sessions/{session_id}", response_model=ConversationHistory)
async def get_session(session_id: str):
    """
    Get conversation history for a session
    """
    try:
        chat_manager = await get_chat_manager()
        history = await chat_manager.get_conversation_history(session_id)

        if not history:
            raise HTTPException(status_code=404, detail="Session not found")

        return history

    except HTTPException:
        raise
    except Exception as e:
        logger.error("get_session_failed", error=str(e), session_id=session_id)
        raise HTTPException(status_code=500, detail="Failed to get session")


@router.delete("/sessions/{session_id}")
async def clear_session(session_id: str):
    """
    Clear conversation history for a session
    """
    try:
        chat_manager = await get_chat_manager()
        success = await chat_manager.clear_conversation(session_id)

        if not success:
            raise HTTPException(status_code=404, detail="Session not found")

        return {"message": "Session cleared successfully"}

    except HTTPException:
        raise
    except Exception as e:
        logger.error("clear_session_failed", error=str(e), session_id=session_id)
        raise HTTPException(status_code=500, detail="Failed to clear session")


@router.get("/sessions", response_model=List[Dict[str, Any]])
async def get_active_sessions():
    """
    Get list of active chat sessions
    """
    try:
        chat_manager = await get_chat_manager()
        sessions = await chat_manager.get_active_sessions()
        return sessions

    except Exception as e:
        logger.error("get_active_sessions_failed", error=str(e))
        raise HTTPException(status_code=500, detail="Failed to get active sessions")


@router.get("/model/info", response_model=ModelInfo)
async def get_model_info():
    """
    Get information about the current model
    """
    try:
        model_manager = await get_model_manager()
        info = model_manager.get_model_info()

        return ModelInfo(
            name=info["name"],
            type=info["type"],
            loaded=info["loaded"],
            parameters=info.get("parameters"),
            capabilities=info.get("capabilities", [])
        )

    except Exception as e:
        logger.error("get_model_info_failed", error=str(e))
        raise HTTPException(status_code=500, detail="Failed to get model info")


@router.get("/health", response_model=HealthResponse)
async def health_check():
    """
    Comprehensive health check endpoint
    """
    try:
        chat_manager = await get_chat_manager()
        health_data = await chat_manager.health_check()

        # Extract key information
        overall_status = health_data.get("overall", {})
        model_info = health_data.get("model_manager", {})
        session_info = health_data.get("session_manager", {})

        return HealthResponse(
            status=overall_status.get("status", "unknown"),
            version=settings.app_version,
            model_type=settings.model_type,
            model_name=settings.model_name,
            model_loaded=model_info.get("status") == "healthy",
            uptime=time.time(),  # Simplified uptime
            active_sessions=session_info.get("active_sessions", 0),
            timestamp=datetime.utcnow()
        )

    except Exception as e:
        logger.error("health_check_failed", error=str(e))
        return HealthResponse(
            status="unhealthy",
            version=settings.app_version,
            model_type=settings.model_type,
            model_name=settings.model_name,
            model_loaded=False,
            uptime=time.time(),
            active_sessions=0,
            timestamp=datetime.utcnow()
        )


@router.get("/status")
async def status():
    """
    Simple status endpoint
    """
    return {
        "status": "ok",
        "service": "sema-chat-api",
        "version": settings.app_version,
        "timestamp": datetime.utcnow().isoformat()
    }


@router.get("/metrics")
async def metrics():
    """
    Prometheus metrics endpoint
    """
    if not settings.enable_metrics:
        raise HTTPException(status_code=404, detail="Metrics not enabled")

    return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)