production/production_server.py

"""
Production-grade API server for ml-intern.

Features:
- FastAPI with async endpoints
- Distributed rate limiting (Redis-backed token bucket)
- Circuit breaker for external API resilience
- Automatic fallback: NIM (primary) -> Cloudflare Workers AI (secondary) -> Gemini (tertiary) -> MLX
- Request/response caching with Redis TTL
- Multi-tenant session isolation
- Health checks and graceful shutdown
- Structured logging with correlation IDs
- Cost tracking and budget enforcement
- Connection pooling for all HTTP clients
- Cloudflare Workers AI support via OpenAI-compatible API
- Google Gemini support via OpenAI-compatible API
- MLX local support for Apple Silicon
"""

import asyncio
import hashlib
import json
import logging
import os
import signal
import sys
import time
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, Optional

import redis.asyncio as aioredis
import asyncpg
from fastapi import FastAPI, HTTPException, Request, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from pydantic import BaseModel, Field
import uvicorn
from prometheus_client import Counter, Histogram, Gauge, generate_latest, CONTENT_TYPE_LATEST
import httpx
from tenacity import retry, stop_after_attempt, wait_exponential, RetryError

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379")
DATABASE_URL = os.environ.get("DATABASE_URL", "")
MAX_CONCURRENT_REQUESTS = int(os.environ.get("MAX_CONCURRENT_REQUESTS", "100"))
DEFAULT_RPM_LIMIT = int(os.environ.get("DEFAULT_RPM_LIMIT", "40"))
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "120"))
CACHE_TTL_SECONDS = int(os.environ.get("CACHE_TTL_SECONDS", "300"))
BUDGET_USD_PER_SESSION = float(os.environ.get("BUDGET_USD_PER_SESSION", "10.0"))
CIRCUIT_BREAKER_FAILURE_THRESHOLD = int(os.environ.get("CIRCUIT_BREAKER_FAILURE_THRESHOLD", "5"))
CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int(os.environ.get("CIRCUIT_BREAKER_RECOVERY_TIMEOUT", "60"))

# Provider-specific endpoints
NIM_API_BASE = os.environ.get("NIM_API_BASE", "https://integrate.api.nvidia.com/v1")
CLOUDFLARE_API_KEY = os.environ.get("CLOUDFLARE_API_KEY", "")
CLOUDFLARE_ACCOUNT_ID = os.environ.get("CLOUDFLARE_ACCOUNT_ID", "")

# Google Gemini / AI Studio
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
GEMINI_API_BASE = os.environ.get("GEMINI_API_BASE", "https://generativelanguage.googleapis.com/v1beta/openai")

# Fallback configuration
FALLBACK_ENABLED = os.environ.get("FALLBACK_ENABLED", "true").lower() == "true"
FALLBACK_PRIMARY = os.environ.get("FALLBACK_PRIMARY", "nim")
FALLBACK_SECONDARY = os.environ.get("FALLBACK_SECONDARY", "cloudflare")
FALLBACK_TERTIARY = os.environ.get("FALLBACK_TERTIARY", "gemini")

# MLX (local Apple Silicon)
MLX_API_BASE = os.environ.get("MLX_API_BASE", "http://localhost:8000/v1")
MLX_ENABLED = os.environ.get("MLX_ENABLED", "false").lower() == "true"

# ---------------------------------------------------------------------------
# Prometheus Metrics
# ---------------------------------------------------------------------------

REQUEST_COUNT = Counter(
    "ml_intern_requests_total",
    "Total requests",
    ["method", "endpoint", "status", "provider"],
)
REQUEST_LATENCY = Histogram(
    "ml_intern_request_duration_seconds",
    "Request duration",
    ["method", "endpoint", "provider"],
    buckets=[0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0],
)
ACTIVE_SESSIONS = Gauge(
    "ml_intern_active_sessions",
    "Number of active sessions",
)
LLM_COST_USD = Counter(
    "ml_intern_llm_cost_usd_total",
    "Total LLM cost in USD",
    ["provider", "model"],
)
CACHE_HIT_COUNT = Counter(
    "ml_intern_cache_hits_total",
    "Cache hits",
    ["cache_type"],
)
CACHE_MISS_COUNT = Counter(
    "ml_intern_cache_misses_total",
    "Cache misses",
    ["cache_type"],
)
CIRCUIT_BREAKER_STATE = Gauge(
    "ml_intern_circuit_breaker_state",
    "Circuit breaker state (0=closed, 1=half-open, 2=open)",
    ["provider"],
)
FALLBACK_COUNT = Counter(
    "ml_intern_fallback_total",
    "Fallback events between providers",
    ["from_provider", "to_provider", "reason"],
)

# ---------------------------------------------------------------------------
# Structured Logging
# ---------------------------------------------------------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | correlation_id=%(correlation_id)s | %(name)s | %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("ml_intern.production")

class CorrelationIdFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        record.correlation_id = getattr(record, "correlation_id", "none")
        return True

logger.addFilter(CorrelationIdFilter())

# ---------------------------------------------------------------------------
# Database Layer
# ---------------------------------------------------------------------------

class DatabasePool:
    def __init__(self, dsn: str):
        self.dsn = dsn
        self._pool: Optional[asyncpg.Pool] = None
    
    async def connect(self):
        if not self.dsn:
            logger.warning("No DATABASE_URL set — skipping database connection")
            return
        self._pool = await asyncpg.create_pool(
            self.dsn,
            min_size=2,
            max_size=10,
            command_timeout=60,
        )
        logger.info("Database pool connected")
    
    async def disconnect(self):
        if self._pool:
            await self._pool.close()
            logger.info("Database pool disconnected")
    
    async def execute(self, query: str, *args):
        if not self._pool:
            return
        async with self._pool.acquire() as conn:
            return await conn.execute(query, *args)
    
    async def fetch(self, query: str, *args):
        if not self._pool:
            return []
        async with self._pool.acquire() as conn:
            return await conn.fetch(query, *args)
    
    async def fetchval(self, query: str, *args):
        if not self._pool:
            return None
        async with self._pool.acquire() as conn:
            return await conn.fetchval(query, *args)

# ---------------------------------------------------------------------------
# Redis Layer
# ---------------------------------------------------------------------------

class RedisManager:
    def __init__(self, url: str):
        self.url = url
        self._redis: Optional[aioredis.Redis] = None
    
    async def connect(self):
        self._redis = aioredis.from_url(self.url, decode_responses=True)
        await self._redis.ping()
        logger.info("Redis connected")
    
    async def disconnect(self):
        if self._redis:
            await self._redis.close()
            logger.info("Redis disconnected")
    
    async def get_cache(self, key: str) -> Optional[str]:
        val = await self._redis.get(key)
        if val:
            CACHE_HIT_COUNT.labels(cache_type="llm_response").inc()
        else:
            CACHE_MISS_COUNT.labels(cache_type="llm_response").inc()
        return val
    
    async def set_cache(self, key: str, value: str, ttl: int = CACHE_TTL_SECONDS):
        await self._redis.setex(key, ttl, value)
    
    async def check_rate_limit(self, key: str, rpm: int) -> tuple[bool, float]:
        now = time.time()
        bucket_key = f"ratelimit:{key}"
        script = """
        local key = KEYS[1]
        local now = tonumber(ARGV[1])
        local rpm = tonumber(ARGV[2])
        local interval = 60.0 / rpm
        local last = redis.call('hget', key, 'last')
        local tokens = redis.call('hget', key, 'tokens')
        if not last then
            last = 0
            tokens = 1
        else
            last = tonumber(last)
            tokens = tonumber(tokens)
        end
        local elapsed = now - last
        tokens = math.min(1, tokens + elapsed / interval)
        if tokens >= 1 then
            tokens = tokens - 1
            redis.call('hmset', key, 'last', now, 'tokens', tokens)
            redis.call('expire', key, 120)
            return {1, 0}
        else
            local retry_after = interval - (elapsed % interval)
            redis.call('hmset', key, 'last', last, 'tokens', tokens)
            redis.call('expire', key, 120)
            return {0, retry_after}
        end
        """
        result = await self._redis.eval(script, 1, bucket_key, now, rpm)
        allowed = bool(result[0])
        retry_after = float(result[1]) if not allowed else 0.0
        return allowed, retry_after
    
    async def get_circuit_state(self, provider: str) -> dict:
        key = f"circuit:{provider}"
        val = await self._redis.get(key)
        if val:
            return json.loads(val)
        return {"state": "closed", "failures": 0, "last_failure": 0}
    
    async def set_circuit_state(self, provider: str, state: dict):
        key = f"circuit:{provider}"
        await self._redis.setex(key, 3600, json.dumps(state))

# ---------------------------------------------------------------------------
# Circuit Breaker
# ---------------------------------------------------------------------------

class CircuitBreaker:
    def __init__(self, redis: RedisManager, provider: str):
        self.redis = redis
        self.provider = provider
        self.failure_threshold = CIRCUIT_BREAKER_FAILURE_THRESHOLD
        self.recovery_timeout = CIRCUIT_BREAKER_RECOVERY_TIMEOUT
    
    async def can_execute(self) -> bool:
        state = await self.redis.get_circuit_state(self.provider)
        if state["state"] == "open":
            if time.time() - state["last_failure"] > self.recovery_timeout:
                state["state"] = "half-open"
                state["failures"] = 0
                await self.redis.set_circuit_state(self.provider, state)
                CIRCUIT_BREAKER_STATE.labels(provider=self.provider).set(1)
                logger.info(f"Circuit breaker {self.provider} half-open")
                return True
            CIRCUIT_BREAKER_STATE.labels(provider=self.provider).set(2)
            return False
        CIRCUIT_BREAKER_STATE.labels(provider=self.provider).set(0 if state["state"] == "closed" else 1)
        return True
    
    async def record_success(self):
        state = await self.redis.get_circuit_state(self.provider)
        if state["state"] == "half-open":
            state["state"] = "closed"
            state["failures"] = 0
            await self.redis.set_circuit_state(self.provider, state)
            CIRCUIT_BREAKER_STATE.labels(provider=self.provider).set(0)
            logger.info(f"Circuit breaker {self.provider} closed")
    
    async def record_failure(self):
        state = await self.redis.get_circuit_state(self.provider)
        state["failures"] += 1
        state["last_failure"] = time.time()
        if state["failures"] >= self.failure_threshold:
            state["state"] = "open"
            CIRCUIT_BREAKER_STATE.labels(provider=self.provider).set(2)
            logger.warning(f"Circuit breaker {self.provider} OPENED after {state['failures']} failures")
        await self.redis.set_circuit_state(self.provider, state)

# ---------------------------------------------------------------------------
# Fallback Manager
# ---------------------------------------------------------------------------

@dataclass
class FallbackConfig:
    primary: str = "nim"
    secondary: str = "cloudflare"
    tertiary: str = "gemini"
    enabled: bool = True

class FallbackManager:
    def __init__(self, redis: RedisManager, config: FallbackConfig = None):
        self.redis = redis
        self.config = config or FallbackConfig()
        self._http_client: Optional[httpx.AsyncClient] = None
    
    async def init_client(self):
        if not self._http_client:
            self._http_client = httpx.AsyncClient(
                limits=httpx.Limits(max_connections=50, max_keepalive_connections=20),
                timeout=httpx.Timeout(REQUEST_TIMEOUT),
            )
    
    async def close_client(self):
        if self._http_client:
            await self._http_client.aclose()
    
    async def get_active_provider(self) -> tuple[str, dict]:
        if not self.config.enabled:
            return self.config.primary, self._get_provider_config(self.config.primary)
        
        primary_breaker = CircuitBreaker(self.redis, self.config.primary)
        if await primary_breaker.can_execute():
            return self.config.primary, self._get_provider_config(self.config.primary)
        
        secondary_breaker = CircuitBreaker(self.redis, self.config.secondary)
        if await secondary_breaker.can_execute():
            FALLBACK_COUNT.labels(
                from_provider=self.config.primary,
                to_provider=self.config.secondary,
                reason="circuit_open",
            ).inc()
            logger.warning(
                f"Fallback: {self.config.primary} unavailable, switching to {self.config.secondary}"
            )
            return self.config.secondary, self._get_provider_config(self.config.secondary)
        
        tertiary_breaker = CircuitBreaker(self.redis, self.config.tertiary)
        if await tertiary_breaker.can_execute():
            FALLBACK_COUNT.labels(
                from_provider=self.config.secondary,
                to_provider=self.config.tertiary,
                reason="secondary_down",
            ).inc()
            logger.warning(
                f"Fallback: both {self.config.primary} and {self.config.secondary} down, "
                f"switching to {self.config.tertiary}"
            )
            return self.config.tertiary, self._get_provider_config(self.config.tertiary)
        
        if MLX_ENABLED:
            mlx_breaker = CircuitBreaker(self.redis, "mlx")
            if await mlx_breaker.can_execute():
                FALLBACK_COUNT.labels(
                    from_provider=self.config.tertiary,
                    to_provider="mlx",
                    reason="all_cloud_down",
                ).inc()
                logger.warning("All cloud providers down — falling back to MLX local")
                return "mlx", self._get_provider_config("mlx")
        
        raise HTTPException(status_code=503, detail="All LLM providers unavailable.")
    
    def _get_provider_config(self, provider: str) -> dict:
        configs = {
            "nim": {
                "api_base": NIM_API_BASE,
                "api_key": os.environ.get("NVIDIA_API_KEY", "no-key"),
                "rpm_limit": 40,
                "cost_per_1m_input": 0.0,
                "cost_per_1m_output": 0.0,
            },
            "cloudflare": {
                "api_base": f"https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/v1",
                "api_key": CLOUDFLARE_API_KEY,
                "rpm_limit": 100,
                "cost_per_1m_input": 0.0,
                "cost_per_1m_output": 0.0,
            },
            "gemini": {
                "api_base": GEMINI_API_BASE,
                "api_key": GEMINI_API_KEY,
                "rpm_limit": 60,
                "cost_per_1m_input": 0.075,
                "cost_per_1m_output": 0.30,
            },
            "mlx": {
                "api_base": MLX_API_BASE,
                "api_key": "no-key",
                "rpm_limit": 1000,
                "cost_per_1m_input": 0.0,
                "cost_per_1m_output": 0.0,
            },
        }
        return configs.get(provider, configs["nim"])

# ---------------------------------------------------------------------------
# Cost Tracking
# ---------------------------------------------------------------------------

@dataclass
class CostTracker:
    session_id: str
    budget_usd: float = BUDGET_USD_PER_SESSION
    spent_usd: float = 0.0
    provider: str = "unknown"
    model: str = "unknown"
    
    def can_spend(self, estimated_cost: float) -> bool:
        return (self.spent_usd + estimated_cost) <= self.budget_usd
    
    def record_spend(self, cost_usd: float):
        self.spent_usd += cost_usd
        LLM_COST_USD.labels(provider=self.provider, model=self.model).inc(cost_usd)
        logger.info(f"Session {self.session_id}: spent ${cost_usd:.4f}, total ${self.spent_usd:.4f} / ${self.budget_usd:.2f}")

# ---------------------------------------------------------------------------
# Concurrency Limiter
# ---------------------------------------------------------------------------

class ConcurrencyLimiter:
    def __init__(self, max_concurrent: int):
        self.semaphore = asyncio.Semaphore(max_concurrent)
    
    async def acquire(self):
        await self.semaphore.acquire()
    
    def release(self):
        self.semaphore.release()

# ---------------------------------------------------------------------------
# Pydantic Models
# ---------------------------------------------------------------------------

class ChatRequest(BaseModel):
    model: str = Field(..., description="Model ID (e.g., gemma-4-31b-bf16)")
    messages: list[dict] = Field(..., description="OpenAI-compatible messages")
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 4096
    stream: bool = False
    tools: Optional[list[dict]] = None
    tool_choice: Optional[str] = "auto"
    session_id: Optional[str] = None
    provider_override: Optional[str] = None

class ChatResponse(BaseModel):
    id: str
    session_id: str
    model: str
    provider: str
    content: Optional[str] = None
    tool_calls: Optional[list[dict]] = None
    usage: dict = Field(default_factory=dict)
    cost_usd: float = 0.0
    cached: bool = False
    finish_reason: Optional[str] = None
    fallback_used: bool = False

class HealthResponse(BaseModel):
    status: str
    version: str = "1.0.0"
    uptime_seconds: float
    active_sessions: int
    redis_connected: bool
    db_connected: bool
    circuit_breakers: dict[str, str]
    fallback_status: dict[str, str]

# ---------------------------------------------------------------------------
# Global State
# ---------------------------------------------------------------------------

db_pool: Optional[DatabasePool] = None
redis_manager: Optional[RedisManager] = None
concurrency_limiter: Optional[ConcurrencyLimiter] = None
fallback_manager: Optional[FallbackManager] = None
start_time: float = 0.0
shutdown_event: asyncio.Event = asyncio.Event()

# ---------------------------------------------------------------------------
# FastAPI App
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    global db_pool, redis_manager, concurrency_limiter, fallback_manager, start_time
    
    start_time = time.time()
    
    db_pool = DatabasePool(DATABASE_URL)
    await db_pool.connect()
    
    redis_manager = RedisManager(REDIS_URL)
    await redis_manager.connect()
    
    concurrency_limiter = ConcurrencyLimiter(MAX_CONCURRENT_REQUESTS)
    fallback_manager = FallbackManager(redis_manager)
    await fallback_manager.init_client()
    
    loop = asyncio.get_event_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, lambda: asyncio.create_task(_shutdown()))
    
    if DATABASE_URL:
        await _init_schema()
    
    logger.info("ml-intern production server started")
    
    yield
    
    logger.info("Shutting down...")
    shutdown_event.set()
    
    if fallback_manager:
        await fallback_manager.close_client()
    if redis_manager:
        await redis_manager.disconnect()
    if db_pool:
        await db_pool.disconnect()
    
    logger.info("ml-intern production server stopped")

async def _shutdown():
    logger.info("Shutdown signal received")
    shutdown_event.set()

async def _init_schema():
    await db_pool.execute("""
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            created_at TIMESTAMP DEFAULT NOW(),
            last_active_at TIMESTAMP DEFAULT NOW(),
            budget_usd NUMERIC DEFAULT 10.0,
            spent_usd NUMERIC DEFAULT 0.0,
            metadata JSONB DEFAULT '{}'
        )
    """)
    await db_pool.execute("""
        CREATE TABLE IF NOT EXISTS requests (
            id TEXT PRIMARY KEY,
            session_id TEXT,
            model TEXT,
            provider TEXT,
            input_tokens INTEGER,
            output_tokens INTEGER,
            cost_usd NUMERIC,
            latency_ms INTEGER,
            cached BOOLEAN DEFAULT FALSE,
            fallback_used BOOLEAN DEFAULT FALSE,
            created_at TIMESTAMP DEFAULT NOW()
        )
    """)
    logger.info("Database schema initialized")

app = FastAPI(
    title="ml-intern Production API",
    description="Production-grade API with NIM/Cloudflare/Gemini fallback and MLX local support",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(GZipMiddleware, minimum_size=1000)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------

@app.middleware("http")
async def correlation_id_middleware(request: Request, call_next):
    correlation_id = request.headers.get("X-Correlation-ID", str(uuid.uuid4()))
    request.state.correlation_id = correlation_id
    
    old_factory = logging.getLogRecordFactory()
    def record_factory(*args, **kwargs):
        record = old_factory(*args, **kwargs)
        record.correlation_id = correlation_id
        return record
    logging.setLogRecordFactory(record_factory)
    
    start = time.time()
    response = await call_next(request)
    latency = time.time() - start
    
    REQUEST_COUNT.labels(
        method=request.method,
        endpoint=request.url.path,
        status=response.status_code,
        provider=getattr(request.state, "provider", "unknown"),
    ).inc()
    
    REQUEST_LATENCY.labels(
        method=request.method,
        endpoint=request.url.path,
        provider=getattr(request.state, "provider", "unknown"),
    ).observe(latency)
    
    response.headers["X-Correlation-ID"] = correlation_id
    return response

# ---------------------------------------------------------------------------
# Helper Functions
# ---------------------------------------------------------------------------

def estimate_cost(provider_config: dict, input_tokens: int, output_tokens: int) -> float:
    cost = (input_tokens / 1_000_000) * provider_config.get("cost_per_1m_input", 0.0)
    cost += (output_tokens / 1_000_000) * provider_config.get("cost_per_1m_output", 0.0)
    return cost

def generate_cache_key(request: ChatRequest) -> str:
    content = json.dumps({
        "model": request.model,
        "messages": request.messages,
        "temperature": request.temperature,
        "max_tokens": request.max_tokens,
        "tools": request.tools,
    }, sort_keys=True)
    return f"cache:llm:{hashlib.sha256(content.encode()).hexdigest()}"

# ---------------------------------------------------------------------------
# LLM Call Implementation
# ---------------------------------------------------------------------------

async def call_llm(
    provider: str,
    provider_config: dict,
    request: ChatRequest,
    session_id: str,
) -> ChatResponse:
    if not fallback_manager or not fallback_manager._http_client:
        raise HTTPException(status_code=503, detail="HTTP client not initialized")
    
    api_base = provider_config["api_base"]
    api_key = provider_config["api_key"]
    
    payload = {
        "model": request.model,
        "messages": request.messages,
        "temperature": request.temperature,
        "max_tokens": request.max_tokens,
        "stream": False,
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    
    if provider == "cloudflare":
        headers["Authorization"] = f"Bearer {api_key}"
    
    start_time = time.time()
    
    try:
        response = await fallback_manager._http_client.post(
            f"{api_base}/chat/completions",
            json=payload,
            headers=headers,
            timeout=REQUEST_TIMEOUT,
        )
        response.raise_for_status()
        data = response.json()
        
        latency_ms = int((time.time() - start_time) * 1000)
        
        usage = data.get("usage", {})
        input_tokens = usage.get("prompt_tokens", 0)
        output_tokens = usage.get("completion_tokens", 0)
        
        cost = estimate_cost(provider_config, input_tokens, output_tokens)
        
        content = None
        tool_calls = None
        if "choices" in data and len(data["choices"]) > 0:
            choice = data["choices"][0]
            message = choice.get("message", {})
            content = message.get("content")
            tool_calls = message.get("tool_calls")
        
        return ChatResponse(
            id=data.get("id", str(uuid.uuid4())),
            session_id=session_id,
            model=request.model,
            provider=provider,
            content=content,
            tool_calls=tool_calls,
            usage=usage,
            cost_usd=cost,
            cached=False,
            finish_reason=data.get("choices", [{}])[0].get("finish_reason"),
            fallback_used=False,
        )
        
    except httpx.HTTPStatusError as e:
        logger.error(f"HTTP error from {provider}: {e.response.status_code} - {e.response.text[:200]}")
        raise HTTPException(status_code=502, detail=f"Provider {provider} returned HTTP {e.response.status_code}")
    except httpx.RequestError as e:
        logger.error(f"Network error calling {provider}: {e}")
        raise HTTPException(status_code=503, detail=f"Cannot reach provider {provider}: {str(e)}")

# ---------------------------------------------------------------------------
# API Endpoints
# ---------------------------------------------------------------------------

@app.get("/health", response_model=HealthResponse)
async def health_check():
    uptime = time.time() - start_time
    
    redis_ok = False
    try:
        await redis_manager._redis.ping()
        redis_ok = True
    except Exception:
        pass
    
    db_ok = False
    if DATABASE_URL:
        try:
            await db_pool.fetchval("SELECT 1")
            db_ok = True
        except Exception:
            pass
    
    circuits = {}
    fallback_status = {}
    for provider in ["nim", "cloudflare", "gemini", "mlx"]:
        try:
            state = await redis_manager.get_circuit_state(provider)
            circuits[provider] = state["state"]
            fallback_status[provider] = "up" if state["state"] == "closed" else "down"
        except Exception:
            circuits[provider] = "unknown"
            fallback_status[provider] = "unknown"
    
    return HealthResponse(
        status="healthy" if redis_ok else "degraded",
        version="1.0.0",
        uptime_seconds=uptime,
        active_sessions=0,
        redis_connected=redis_ok,
        db_connected=db_ok,
        circuit_breakers=circuits,
        fallback_status=fallback_status,
    )

@app.get("/metrics")
async def metrics():
    from starlette.responses import Response
    return Response(content=generate_latest(), media_type=CONTENT_TYPE_LATEST)

@app.post("/v1/chat/completions", response_model=ChatResponse)
async def chat_completions(request: ChatRequest, background_tasks: BackgroundTasks):
    correlation_id = getattr(request.state, "correlation_id", str(uuid.uuid4()))
    session_id = request.session_id or str(uuid.uuid4())
    
    logger.info(f"Chat request: model={request.model}, stream={request.stream}, session={session_id}")
    
    await concurrency_limiter.acquire()
    try:
        if request.provider_override:
            provider = request.provider_override
            provider_config = fallback_manager._get_provider_config(provider)
            breaker = CircuitBreaker(redis_manager, provider)
            if not await breaker.can_execute():
                raise HTTPException(status_code=503, detail=f"Provider {provider} circuit breaker is open")
        else:
            provider, provider_config = await fallback_manager.get_active_provider()
        
        request.state.provider = provider
        
        rpm = provider_config.get("rpm_limit", DEFAULT_RPM_LIMIT)
        rate_limit_key = f"{provider}:{session_id}"
        allowed, retry_after = await redis_manager.check_rate_limit(rate_limit_key, rpm)
        if not allowed:
            logger.warning(f"Rate limit exceeded for {rate_limit_key}")
            raise HTTPException(status_code=429, detail=f"Rate limit exceeded. Retry after {retry_after:.1f}s", headers={"Retry-After": str(int(retry_after))})
        
        if not request.stream:
            cache_key = generate_cache_key(request)
            cached = await redis_manager.get_cache(cache_key)
            if cached:
                logger.info(f"Cache hit for {cache_key}")
                data = json.loads(cached)
                return ChatResponse(
                    id=str(uuid.uuid4()),
                    session_id=session_id,
                    model=request.model,
                    provider=provider,
                    content=data.get("content"),
                    tool_calls=data.get("tool_calls"),
                    usage=data.get("usage", {}),
                    cost_usd=0.0,
                    cached=True,
                    finish_reason=data.get("finish_reason"),
                    fallback_used=False,
                )
        
        cost_tracker = CostTracker(session_id, provider=provider, model=request.model)
        
        response = await call_llm(provider, provider_config, request, session_id)
        
        breaker = CircuitBreaker(redis_manager, provider)
        await breaker.record_success()
        
        cost_tracker.record_spend(response.cost_usd)
        
        if not request.stream:
            cache_key = generate_cache_key(request)
            await redis_manager.set_cache(
                cache_key,
                json.dumps({
                    "content": response.content,
                    "tool_calls": response.tool_calls,
                    "usage": response.usage,
                    "finish_reason": response.finish_reason,
                }),
            )
        
        if DATABASE_URL:
            background_tasks.add_task(_persist_request, session_id, request, response, provider)
        
        if provider != FALLBACK_PRIMARY and FALLBACK_ENABLED:
            response.fallback_used = True
        
        return response
        
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"Error processing request: {e}")
        breaker = CircuitBreaker(redis_manager, provider if 'provider' in locals() else "unknown")
        await breaker.record_failure()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        concurrency_limiter.release()

async def _persist_request(session_id: str, request: ChatRequest, response: ChatResponse, provider: str):
    try:
        await db_pool.execute(
            """
            INSERT INTO requests (id, session_id, model, provider, input_tokens, 
                                  output_tokens, cost_usd, latency_ms, cached, fallback_used)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
            """,
            response.id,
            session_id,
            request.model,
            provider,
            response.usage.get("prompt_tokens", 0),
            response.usage.get("completion_tokens", 0),
            response.cost_usd,
            0,
            response.cached,
            response.fallback_used,
        )
    except Exception as e:
        logger.error(f"Failed to persist request: {e}")

@app.get("/v1/models")
async def list_models():
    models = []
    
    if os.environ.get("NVIDIA_API_KEY"):
        models.extend([
            {"id": "nim/llama-3.1-405b-instruct", "object": "model", "owned_by": "nvidia"},
            {"id": "nim/llama-3.1-70b-instruct", "object": "model", "owned_by": "nvidia"},
            {"id": "nim/llama-3.1-8b-instruct", "object": "model", "owned_by": "nvidia"},
            {"id": "nim/mistral-7b-instruct", "object": "model", "owned_by": "nvidia"},
        ])
    
    if CLOUDFLARE_API_KEY and CLOUDFLARE_ACCOUNT_ID:
        models.extend([
            {"id": "cloudflare/@cf/meta/llama-3.1-8b-instruct", "object": "model", "owned_by": "cloudflare"},
            {"id": "cloudflare/@cf/meta/llama-3.1-70b-instruct", "object": "model", "owned_by": "cloudflare"},
            {"id": "cloudflare/@cf/mistral/mistral-7b-instruct", "object": "model", "owned_by": "cloudflare"},
            {"id": "cloudflare/@cf/qwen/qwen1.5-14b-chat-awq", "object": "model", "owned_by": "cloudflare"},
            {"id": "cloudflare/@cf/google/gemma-4-26b-a4b-it", "object": "model", "owned_by": "cloudflare"},
        ])
    
    if GEMINI_API_KEY:
        models.extend([
            {"id": "gemini/gemini-2.5-pro-preview", "object": "model", "owned_by": "google"},
            {"id": "gemini/gemini-2.5-flash-preview", "object": "model", "owned_by": "google"},
            {"id": "gemini/gemma-4-26b", "object": "model", "owned_by": "google"},
            {"id": "gemini/gemma-4-9b", "object": "model", "owned_by": "google"},
        ])
    
    if MLX_ENABLED:
        models.extend([
            {"id": "mlx/llama-3.1-8b", "object": "model", "owned_by": "mlx"},
            {"id": "mlx/llama-3.1-70b", "object": "model", "owned_by": "mlx"},
            {"id": "mlx/gemma-4-26b-a4b-it", "object": "model", "owned_by": "mlx"},
            {"id": "mlx/gemma-4-31b-bf16", "object": "model", "owned_by": "mlx"},
            {"id": "mlx/gemma-4-e4b-it", "object": "model", "owned_by": "mlx"},
        ])
    
    return {"object": "list", "data": models}

@app.get("/v1/fallback/status")
async def fallback_status():
    status = {}
    for provider in ["nim", "cloudflare", "gemini", "mlx"]:
        breaker = CircuitBreaker(redis_manager, provider)
        can_execute = await breaker.can_execute()
        state = await redis_manager.get_circuit_state(provider)
        status[provider] = {
            "circuit_state": state["state"],
            "failures": state["failures"],
            "available": can_execute,
            "last_failure": state["last_failure"],
        }
    
    return {
        "fallback_enabled": FALLBACK_ENABLED,
        "primary": FALLBACK_PRIMARY,
        "secondary": FALLBACK_SECONDARY,
        "tertiary": FALLBACK_TERTIARY,
        "providers": status,
        "active_provider": await _get_active_provider_name(),
    }

async def _get_active_provider_name() -> str:
    try:
        provider, _ = await fallback_manager.get_active_provider()
        return provider
    except HTTPException:
        return "none_available"

# ---------------------------------------------------------------------------
# Main Entry Point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    port = int(os.environ.get("PORT", "8000"))
    workers = int(os.environ.get("WORKERS", "1"))
    
    uvicorn.run(
        "production_server:app",
        host="0.0.0.0",
        port=port,
        workers=workers,
        log_level="info",
        access_log=True,
    )