import os
import torch
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from typing import List, Union
from sentence_transformers import SentenceTransformer

MODEL_ID = "google/embeddinggemma-2"
TRUNC_DIM = int(os.getenv("TRUNCATE_DIM", "256"))
API_KEY = os.getenv("EMBED_API_KEY", "")

print("loading text-only 270M...", flush=True)
model = SentenceTransformer(
  MODEL_ID,
  config_kwargs={"vision_config": None, "audio_config": None},
  model_kwargs={"torch_dtype": torch.float32},
  truncate_dim=TRUNC_DIM,
)
print("ready", flush=True)
app = FastAPI()

def auth_ok(auth: str | None):
  if API_KEY and auth != f"Bearer {API_KEY}":
    raise HTTPException(401, "unauthorized")

class EmbedIn(BaseModel):
  input: Union[str, List[str]]
  model: str | None = None
  prompt_name: str | None = None  # SearchQuery / Document / QuestionAnswering etc.

class RerankIn(BaseModel):
  query: str
  documents: List[str]
  top_n: int | None = 8
  query_prompt: str | None = "SearchQuery"
  doc_prompt: str | None = "Document"

@app.get("/health")
def health(): return {"ok": True, "dim": TRUNC_DIM}

@app.post("/embed")
def embed(body: EmbedIn, authorization: str | None = Header(None)):
  auth_ok(authorization)
  texts = [body.input] if isinstance(body.input, str) else body.input
  texts = [(t or ".")[:4000] for t in texts][:32]
  kwargs = {"truncate_dim": TRUNC_DIM, "normalize_embeddings": True, "show_progress_bar": False}
  if body.prompt_name:
    kwargs["prompt_name"] = body.prompt_name
  vecs = model.encode(texts, **kwargs)
  return {"data": [{"embedding": v.tolist(), "index": i} for i, v in enumerate(vecs)], "model": "embeddinggemma-2"}

@app.post("/rerank")
def rerank(body: RerankIn, authorization: str | None = Header(None)):
  auth_ok(authorization)
  q = model.encode([body.query[:2000]], prompt_name=body.query_prompt, truncate_dim=TRUNC_DIM, normalize_embeddings=True, show_progress_bar=False)[0]
  d = model.encode([x[:2000] for x in body.documents[:20]], prompt_name=body.doc_prompt, truncate_dim=TRUNC_DIM, normalize_embeddings=True, show_progress_bar=False)
  scores = sorted([(float(q @ v), i) for i, v in enumerate(d)], reverse=True)
  return {"order": [i for _, i in scores[:body.top_n or 8]]}