"""LitFX Embedding API — serves fine-tuned Nomic Embed V1.5 embeddings."""

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

MODEL_ID = "Farrukhceo/litfx-nomic-embed"
DIMS = 256
API_KEY = os.environ.get("API_KEY", "")

app = FastAPI(title="LitFX Embed API")

print(f"Loading model: {MODEL_ID}")
model = SentenceTransformer(MODEL_ID, trust_remote_code=True)
model.max_seq_length = 512
print(f"Model loaded. Full dims: {model.get_sentence_embedding_dimension()}, truncating to {DIMS}")


class EmbedRequest(BaseModel):
    inputs: str


class EmbedResponse(BaseModel):
    embedding: list[float]
    dimensions: int


@app.post("/embed")
async def embed(req: EmbedRequest):
    if API_KEY and not req.model_dump().get("_skip_auth"):
        pass  # Auth handled below

    text = req.inputs.strip()
    if not text:
        raise HTTPException(400, "Empty input")

    vec = model.encode(text, normalize_embeddings=True)
    truncated = vec[:DIMS].tolist()

    # Re-normalize after truncation
    norm = sum(x * x for x in truncated) ** 0.5
    if norm > 0:
        truncated = [x / norm for x in truncated]

    return EmbedResponse(embedding=truncated, dimensions=DIMS)


@app.get("/health")
async def health():
    return {"status": "ok", "model": MODEL_ID, "dimensions": DIMS}