"""Embedding microservice — exposes sentence-transformers models via REST API.

POST /embed  {"texts": ["...", "..."]}  →  {"embeddings": [[...], [...]], "model": "...", "dimension": N}
GET  /health                            →  {"status": "ok", "model": "...", "dimension": N}
"""

from __future__ import annotations

import logging
import os
import time

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)

MODEL_NAME = os.environ.get("EMBEDDING_MODEL", "all-MiniLM-L6-v2")
BATCH_SIZE = int(os.environ.get("EMBEDDING_BATCH_SIZE", "64"))

app = FastAPI(title="Docling Studio — Embedding Service", version="0.4.0")

# Load model at startup (downloaded / cached in HF cache dir)
model: SentenceTransformer | None = None


@app.on_event("startup")
async def _load_model() -> None:
    global model  # noqa: PLW0603
    logger.info("Loading sentence-transformers model '%s' …", MODEL_NAME)
    t0 = time.monotonic()
    model = SentenceTransformer(MODEL_NAME)
    elapsed = time.monotonic() - t0
    dim = model.get_sentence_embedding_dimension()
    logger.info("Model loaded in %.1fs — dimension=%d", elapsed, dim)


# -- Schemas -------------------------------------------------------------------


class EmbedRequest(BaseModel):
    texts: list[str] = Field(..., min_length=1, description="Texts to embed")


class EmbedResponse(BaseModel):
    embeddings: list[list[float]]
    model: str
    dimension: int


class HealthResponse(BaseModel):
    status: str
    model: str
    dimension: int


# -- Endpoints -----------------------------------------------------------------


@app.post("/embed", response_model=EmbedResponse)
async def embed(request: EmbedRequest) -> EmbedResponse:
    """Generate embeddings for a batch of texts."""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    vectors = model.encode(
        request.texts,
        batch_size=BATCH_SIZE,
        show_progress_bar=False,
        normalize_embeddings=True,
    )
    return EmbedResponse(
        embeddings=vectors.tolist(),
        model=MODEL_NAME,
        dimension=model.get_sentence_embedding_dimension(),
    )


@app.get("/health", response_model=HealthResponse)
async def health() -> HealthResponse:
    """Health check — verifies the model is loaded."""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    return HealthResponse(
        status="ok",
        model=MODEL_NAME,
        dimension=model.get_sentence_embedding_dimension(),
    )