import uvicorn
from fastapi import FastAPI, Request
from pydantic import BaseModel
from utils.logger import get_logger

logger = get_logger("sail.api")

class ChatRequest(BaseModel):
    message: str

class APIServer:
    def __init__(self, engine, port=8080):
        self.engine = engine
        self.port = port
        self.app = FastAPI(title="SAIL v2 API")
        
        self.app.post("/chat")(self.chat)
        self.app.get("/status")(self.status)

    async def chat(self, request: ChatRequest):
        logger.info(f"API request: {request.message}")
        result = self.engine.generate(request.message)
        return result

    async def status(self):
        return {
            "status": "online",
            "model": self.engine.settings.model_id,
            "vram": self.engine.vram_usage()
        }

    def start(self):
        logger.info(f"Starting API server on port {self.port}")
        uvicorn.run(self.app, host="0.0.0.0", port=self.port)