# apps/embedding-runtime/server.py
"""SmartMLOps Tier-2 embedding-runtime.

A small OpenAI-compatible embeddings/rerank server for the Vector Inference line.
KServe owns the lifecycle (this is a ClusterServingRuntime container, same pattern
as llama-cpp-runtime); we own only the serving code. The base model is mounted by
KServe at /mnt/models (overridable via EMBEDDING_MODEL_DIR). Two families are covered:

  * instruction/prefix class (E5, Qwen3-Embedding, Instructor): a per-`task`
    instruction string is prepended to each input before encoding.
  * adapter-switching class (Jina v5 base): a LoRA adapter is selected per `task`
    from EMBEDDING_TASK_ADAPTERS (JSON: {"retrieval": "<adapter-subdir-or-repo>"}).

Endpoints:
  POST /v1/embeddings   OpenAI-compatible. Body: {input: str|[str], model?, task?, instruction?}
  POST /v1/rerank       {query: str, documents: [str], top_n?: int, task?, instruction?}
  GET  /health          readiness/liveness
"""
import json
import os
from typing import List, Optional, Union

import numpy as np
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

MODEL_DIR = os.environ.get("EMBEDDING_MODEL_DIR", "/mnt/models")
MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME", MODEL_DIR)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# task -> instruction prefix (E5/Qwen3/Instructor style). Override with EMBEDDING_TASK_INSTRUCTIONS (JSON).
_DEFAULT_INSTRUCTIONS = {
    "retrieval.query": "Represent this query for retrieving relevant documents: ",
    "retrieval.passage": "Represent this document for retrieval: ",
    "classification": "Represent this text for classification: ",
    "clustering": "Represent this text for clustering: ",
}
TASK_INSTRUCTIONS = {
    **_DEFAULT_INSTRUCTIONS,
    **json.loads(os.environ.get("EMBEDDING_TASK_INSTRUCTIONS", "{}")),
}
# task -> LoRA adapter path/repo (Jina v5 base style). Empty by default. JSON.
TASK_ADAPTERS = json.loads(os.environ.get("EMBEDDING_TASK_ADAPTERS", "{}"))

app = FastAPI(title="smartmlops-embedding-runtime")

_model: Optional[SentenceTransformer] = None
_active_adapter: Optional[str] = None


def _load_model() -> SentenceTransformer:
    global _model
    if _model is None:
        _model = SentenceTransformer(MODEL_NAME, device=DEVICE, trust_remote_code=False)
    return _model


def _apply_adapter(task: Optional[str]) -> None:
    """Switch the active LoRA adapter for `task`, if one is configured.

    Adapter switching uses PEFT under the hood when the underlying transformer is a
    PeftModel; for non-adapter base models this is a no-op. Loading is lazy + cached.
    """
    global _active_adapter
    if not task or task not in TASK_ADAPTERS:
        return
    adapter = TASK_ADAPTERS[task]
    if adapter == _active_adapter:
        return
    model = _load_model()
    inner = model[0].auto_model  # transformers model inside the ST wrapper
    load_adapter = getattr(inner, "load_adapter", None)
    if callable(load_adapter):
        try:
            load_adapter(adapter, adapter_name=task)
            set_adapter = getattr(inner, "set_adapter", None)
            if callable(set_adapter):
                set_adapter(task)
            _active_adapter = adapter
        except Exception as exc:  # noqa: BLE001 — surface a clean 400 to the caller
            raise HTTPException(status_code=400, detail=f"adapter load failed for task={task!r}: {exc}")


def _instruction_for(task: Optional[str], explicit: Optional[str]) -> str:
    if explicit:
        return explicit
    if task and task in TASK_INSTRUCTIONS:
        return TASK_INSTRUCTIONS[task]
    return ""


def _encode(texts: List[str], task: Optional[str], instruction: Optional[str]) -> np.ndarray:
    _apply_adapter(task)
    prefix = _instruction_for(task, instruction)
    prepared = [f"{prefix}{t}" for t in texts] if prefix else texts
    model = _load_model()
    return model.encode(prepared, normalize_embeddings=True, convert_to_numpy=True)


class EmbeddingsRequest(BaseModel):
    input: Union[str, List[str]]
    model: Optional[str] = None
    task: Optional[str] = None
    instruction: Optional[str] = None


class RerankRequest(BaseModel):
    query: str
    documents: List[str]
    top_n: Optional[int] = None
    model: Optional[str] = None
    task: Optional[str] = None
    instruction: Optional[str] = None


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "device": DEVICE, "model": MODEL_NAME}


@app.post("/v1/embeddings")
def embeddings(req: EmbeddingsRequest) -> dict:
    texts = [req.input] if isinstance(req.input, str) else list(req.input)
    if not texts:
        raise HTTPException(status_code=400, detail="input must be a non-empty string or list")
    vecs = _encode(texts, req.task, req.instruction)
    data = [
        {"object": "embedding", "index": i, "embedding": vec.tolist()}
        for i, vec in enumerate(vecs)
    ]
    tokens = sum(len(t.split()) for t in texts)  # coarse, input-only (no streaming/completion)
    return {
        "object": "list",
        "data": data,
        "model": req.model or MODEL_NAME,
        "usage": {"prompt_tokens": tokens, "total_tokens": tokens},
    }


@app.post("/v1/rerank")
def rerank(req: RerankRequest) -> dict:
    if not req.documents:
        raise HTTPException(status_code=400, detail="documents must be a non-empty list")
    # Cosine similarity over normalized embeddings (cross-encoder-free; works for any
    # embedding base). The query gets the retrieval.query instruction by default.
    qtask = req.task or "retrieval.query"
    dtask = req.task or "retrieval.passage"
    qvec = _encode([req.query], qtask, req.instruction)[0]
    dvecs = _encode(list(req.documents), dtask, req.instruction)
    scores = dvecs @ qvec  # normalized => dot product == cosine similarity
    order = np.argsort(-scores)
    if req.top_n is not None:
        order = order[: req.top_n]
    results = [
        {"index": int(i), "relevance_score": float(scores[i]), "document": req.documents[i]}
        for i in order
    ]
    return {"model": req.model or MODEL_NAME, "results": results}