"""FastAPI application for the specialist inference service."""

from __future__ import annotations

import logging
import os
import time
from contextlib import asynccontextmanager

from fastapi import FastAPI

from services.specialist.engine import SpecialistEngine
from services.specialist.router import router, set_engine

logger = logging.getLogger(__name__)

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

SPECIALIST_MODEL = os.environ.get("SPECIALIST_MODEL", "urchade/gliner_large-v2.1")
MAX_BATCH_SIZE = int(os.environ.get("SPECIALIST_MAX_BATCH_SIZE", "32"))
MAX_WAIT_MS = float(os.environ.get("SPECIALIST_MAX_WAIT_MS", "50.0"))
MAX_QUEUE_SIZE = int(os.environ.get("SPECIALIST_MAX_QUEUE_SIZE", "256"))


# ---------------------------------------------------------------------------
# Lifespan: load model and warm up on startup
# ---------------------------------------------------------------------------


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan — loads model and warms up on startup."""
    engine = SpecialistEngine(
        model_name=SPECIALIST_MODEL,
        max_batch_size=MAX_BATCH_SIZE,
        max_wait_ms=MAX_WAIT_MS,
        max_queue_size=MAX_QUEUE_SIZE,
    )
    engine.load_model()
    engine.warm_up()
    app.state.engine = engine
    set_engine(engine)
    logger.info(
        "Specialist service ready: model=%s, max_batch=%d, max_wait_ms=%.1f, max_queue=%d",
        SPECIALIST_MODEL,
        MAX_BATCH_SIZE,
        MAX_WAIT_MS,
        MAX_QUEUE_SIZE,
    )
    yield
    logger.info("Specialist service shutting down")


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Specialist Inference Service",
    description="CPU-first NER, classification, relation, and structured extraction",
    version="1.0.0",
    lifespan=lifespan,
)

app.include_router(router)

# Track startup time for metrics
_start_time = time.time()


# ---------------------------------------------------------------------------
# Health / Readiness endpoints
# ---------------------------------------------------------------------------


@app.get("/health")
async def health():
    """Liveness probe — returns 200 if the process is running."""
    return {"status": "ok"}


@app.get("/ready")
async def ready():
    """Readiness probe — returns 200 only when the engine is loaded."""
    engine = getattr(app.state, "engine", None)
    if engine is None or not engine.is_ready:
        return {"status": "not_ready"}, 503
    return {
        "status": "ready",
        "model": engine.model_name,
        "uptime_seconds": round(time.time() - _start_time, 1),
    }


@app.get("/metrics")
async def metrics():
    """Basic Prometheus-style metrics endpoint."""
    engine = getattr(app.state, "engine", None)
    model_loaded = 1 if (engine and engine.is_ready) else 0
    uptime = round(time.time() - _start_time, 1)

    batcher_metrics = engine.batcher_metrics if engine else {}

    return {
        "specialist_model_loaded": model_loaded,
        "specialist_uptime_seconds": uptime,
        "specialist_max_batch_size": MAX_BATCH_SIZE,
        "specialist_max_wait_ms": MAX_WAIT_MS,
        "specialist_max_queue_size": MAX_QUEUE_SIZE,
        "specialist_model_name": SPECIALIST_MODEL,
        "specialist_total_batches": batcher_metrics.get("total_batches", 0),
        "specialist_total_items": batcher_metrics.get("total_items", 0),
        "specialist_total_rejections": batcher_metrics.get("total_rejections", 0),
        "specialist_queue_depth": batcher_metrics.get("queue_depth", 0),
    }