import os
from typing import Literal

from fastapi import Depends, FastAPI, Header, HTTPException, status
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer

MODEL_NAME = os.getenv("EMBEDDING_MODEL_NAME", "BAAI/bge-base-en-v1.5")
SERVICE_TOKEN = os.getenv("EMBEDDING_SERVICE_TOKEN", "")
PASSAGE_PREFIX = "Represent this passage for searching relevant passages: "
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "

app = FastAPI(title="CognifyAI Embedding Service")
model = SentenceTransformer(MODEL_NAME)


class EmbedRequest(BaseModel):
    text: str = Field(min_length=1)
    kind: Literal["passage", "query"] = "passage"


class EmbedBatchRequest(BaseModel):
    texts: list[str] = Field(min_length=1)
    kind: Literal["passage", "query"] = "passage"


def require_token(authorization: str = Header(default="")) -> None:
    if not SERVICE_TOKEN:
        return

    expected = f"Bearer {SERVICE_TOKEN}"
    if authorization != expected:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid embedding service token",
        )


def prefix_text(text: str, kind: str) -> str:
    if kind == "query":
        return f"{QUERY_PREFIX}{text}"
    return f"{PASSAGE_PREFIX}{text}"


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


@app.post("/embed", dependencies=[Depends(require_token)])
def embed(payload: EmbedRequest) -> dict:
    embedding = model.encode(
        prefix_text(payload.text, payload.kind),
        normalize_embeddings=True,
    )
    return {"embedding": embedding.tolist(), "model": MODEL_NAME}


@app.post("/embed/batch", dependencies=[Depends(require_token)])
def embed_batch(payload: EmbedBatchRequest) -> dict:
    prefixed_texts = [prefix_text(text, payload.kind) for text in payload.texts]
    embeddings = model.encode(
        prefixed_texts,
        normalize_embeddings=True,
        show_progress_bar=False,
    )
    return {"embeddings": embeddings.tolist(), "model": MODEL_NAME}