"""Hugging Face ZeroGPU Space hosting Qwen3-VL-Embedding-2B.

This is a drop-in replacement for the Modal deployment of `colab_server/`. It
speaks the **same wire contract** (`GET /health`, `POST /embed`,
`POST /embed_text`), so `app/embed_client.py` talks to it unchanged - only
`EMBEDDING_SERVER_URL` moves.

Two things are specific to ZeroGPU and deliberate:

1. **The model is placed on `cuda` at module level**, not lazily inside the
   GPU function. ZeroGPU runs a CUDA emulation layer outside `@spaces.GPU`
   so this works, and the docs are explicit that startup placement is far more
   efficient than transferring inside the decorated call.
2. **Only the forward pass is decorated.** A real GPU is attached for the
   duration of a `@spaces.GPU` call and released after, and daily quota is
   consumed by that time - so request parsing, base64 decoding and response
   building all stay outside it.

Embeddings are produced **directly from the page/tile image**. There is no
caption step anywhere in this file; that path is the OpenAI baseline and lives
in `app/embed_openai.py`.
"""

from __future__ import annotations

import base64
import binascii
import io
import logging
import os
import time
from typing import Annotated, Literal

import spaces
import torch
from fastapi import Header, HTTPException, Request
from fastapi.responses import JSONResponse
from gradio import Server
from PIL import Image, UnidentifiedImageError
from pydantic import BaseModel, Field, field_validator
from sentence_transformers import SentenceTransformer

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)-7s | %(message)s")
log = logging.getLogger("visualops.space")

# ======================================================================
# Configuration
# ======================================================================

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen3-VL-Embedding-2B")

# The model card publishes up to 2048 dimensions. This is asserted rather than
# assumed: the local FAISS manifest pins the width, and a silent change would
# invalidate every stored vector while still returning plausible numbers.
EXPECTED_DIMENSION = int(os.environ.get("EXPECTED_DIMENSION", 2048))

MAX_BATCH_ITEMS = int(os.environ.get("MAX_BATCH_ITEMS", 64))
AUTH_TOKEN = os.environ.get("EMBEDDING_SERVER_TOKEN", "").strip()

# Seconds of GPU time requested per call. Quota is consumed by *effective*
# runtime, but a declared ceiling that is too low kills the call mid-batch.
GPU_DURATION_BASE = int(os.environ.get("GPU_DURATION_BASE", 20))
GPU_SECONDS_PER_IMAGE = float(os.environ.get("GPU_SECONDS_PER_IMAGE", 1.5))
GPU_SECONDS_PER_TEXT = float(os.environ.get("GPU_SECONDS_PER_TEXT", 0.2))

_STARTED_AT = time.time()


# ======================================================================
# Model - loaded once, at import, and reused by every request
# ======================================================================


def _load_model() -> tuple[SentenceTransformer, int, str]:
    """Load Qwen3-VL onto cuda at startup, per the ZeroGPU model-loading rule."""
    dtype = torch.float16
    model = SentenceTransformer(MODEL_ID, device="cuda", model_kwargs={"torch_dtype": dtype})
    model.eval()

    raw_dim = model.get_sentence_embedding_dimension()
    if not raw_dim:
        raise RuntimeError(f"{MODEL_ID} reported no embedding dimension")
    dim = int(raw_dim)
    if dim != EXPECTED_DIMENSION:
        # Loud, at startup: a width change must never reach the index quietly.
        raise RuntimeError(
            f"{MODEL_ID} produced {dim}-d embeddings, expected {EXPECTED_DIMENSION}. "
            "The FAISS manifest pins this width; refusing to serve."
        )
    return model, dim, str(dtype).replace("torch.", "")


_MODEL: SentenceTransformer | None = None
_DIMENSION = 0
_DTYPE = ""
_LOAD_ERROR = ""

try:
    _MODEL, _DIMENSION, _DTYPE = _load_model()
    log.info("loaded %s (%d-d, %s)", MODEL_ID, _DIMENSION, _DTYPE)
except Exception as exc:  # noqa: BLE001 - /health must survive to report why
    _LOAD_ERROR = f"{type(exc).__name__}: {exc}"
    log.error("model failed to load: %s", _LOAD_ERROR)


# ======================================================================
# Schemas - identical to colab_server/server.py
# ======================================================================


class EmbedImagesRequest(BaseModel):
    """Base64-encoded images. `data:` prefixes are accepted and stripped."""

    images: list[str] = Field(min_length=1)
    batch_size: int = Field(default=8, ge=1, le=32)
    dimension: int | None = Field(default=None, ge=8, description="Optional Matryoshka truncation.")

    @field_validator("images")
    @classmethod
    def _limit_batch(cls, v: list[str]) -> list[str]:
        if len(v) > MAX_BATCH_ITEMS:
            raise ValueError(f"at most {MAX_BATCH_ITEMS} images per request")
        return v


class EmbedTextRequest(BaseModel):
    texts: list[str] = Field(min_length=1)
    is_query: bool = Field(
        default=True,
        description="Queries and documents are encoded asymmetrically by this model.",
    )
    batch_size: int = Field(default=16, ge=1, le=64)
    dimension: int | None = Field(default=None, ge=8)

    @field_validator("texts")
    @classmethod
    def _limit_batch(cls, v: list[str]) -> list[str]:
        if len(v) > MAX_BATCH_ITEMS:
            raise ValueError(f"at most {MAX_BATCH_ITEMS} texts per request")
        return v


class EmbedResponse(BaseModel):
    embeddings: list[list[float]]
    dimension: int
    count: int
    model: str
    normalized: bool = True
    elapsed_ms: float


class HealthResponse(BaseModel):
    status: Literal["ok", "loading", "error"]
    model: str
    model_id: str
    backend: str
    device: str
    dtype: str
    dimension: int
    uptime_seconds: float
    detail: str | None = None


# ======================================================================
# GPU work - the only code that holds a real GPU
# ======================================================================


def _image_duration(images: list[Image.Image], batch_size: int) -> int:
    return int(GPU_DURATION_BASE + GPU_SECONDS_PER_IMAGE * len(images))


def _text_duration(texts: list[str], is_query: bool, batch_size: int) -> int:
    return int(GPU_DURATION_BASE + GPU_SECONDS_PER_TEXT * len(texts))


@spaces.GPU(duration=_image_duration)
@torch.inference_mode()
def embed_images_gpu(images: list[Image.Image], batch_size: int) -> list[list[float]]:
    """Embed page/tile images directly. No captioning, no text proxy."""
    assert _MODEL is not None
    encode = getattr(_MODEL, "encode_document", _MODEL.encode)
    vectors = encode(
        images, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True
    )
    return [v.astype("float32").tolist() for v in vectors]


@spaces.GPU(duration=_text_duration)
@torch.inference_mode()
def embed_texts_gpu(texts: list[str], is_query: bool, batch_size: int) -> list[list[float]]:
    """Embed text into the *same* Qwen space the images live in."""
    assert _MODEL is not None
    encode = (
        _MODEL.encode_query
        if is_query and hasattr(_MODEL, "encode_query")
        else getattr(_MODEL, "encode_document", _MODEL.encode)
    )
    vectors = encode(
        texts, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True
    )
    return [v.astype("float32").tolist() for v in vectors]


# ======================================================================
# Helpers
# ======================================================================


def _decode_image(payload: str, position: int) -> Image.Image:
    raw = payload.split(",", 1)[1] if payload.startswith("data:") else payload
    try:
        data = base64.b64decode(raw, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise HTTPException(400, f"image[{position}] is not valid base64: {exc}") from exc
    try:
        return Image.open(io.BytesIO(data)).convert("RGB")
    except (UnidentifiedImageError, OSError) as exc:
        raise HTTPException(400, f"image[{position}] is not a readable image: {exc}") from exc


def _truncate(vectors: list[list[float]], dimension: int | None) -> list[list[float]]:
    """Matryoshka truncation, renormalised so cosine scores stay comparable."""
    if dimension is None or dimension >= _DIMENSION:
        return vectors
    out: list[list[float]] = []
    for vector in vectors:
        head = torch.tensor(vector[:dimension], dtype=torch.float32)
        out.append(torch.nn.functional.normalize(head, dim=0).tolist())
    return out


def _require_model() -> None:
    if _MODEL is None:
        raise HTTPException(503, f"model unavailable: {_LOAD_ERROR or 'not loaded'}")


async def require_token(authorization: Annotated[str | None, Header()] = None) -> None:
    """Optional bearer auth. Enabled by setting EMBEDDING_SERVER_TOKEN."""
    if not AUTH_TOKEN:
        return
    if authorization != f"Bearer {AUTH_TOKEN}":
        raise HTTPException(status_code=401, detail="invalid or missing bearer token")


# ======================================================================
# App - gradio.Server is a FastAPI subclass; ZeroGPU requires the Gradio SDK,
# and custom routes take priority over Gradio's own.
# ======================================================================

app = Server()


@app.exception_handler(ValueError)
async def _bad_request(request: Request, exc: ValueError) -> JSONResponse:
    return JSONResponse(status_code=400, content={"error": "bad_request", "detail": str(exc)})


@app.get("/health", response_model=HealthResponse)
async def health() -> HealthResponse:
    ready = _MODEL is not None
    return HealthResponse(
        status="ok" if ready else "error",
        model=MODEL_ID,
        model_id=MODEL_ID,
        backend="sentence-transformers",
        device="cuda-zerogpu",
        dtype=_DTYPE or "unknown",
        dimension=_DIMENSION,
        uptime_seconds=round(time.time() - _STARTED_AT, 2),
        detail=_LOAD_ERROR or None,
    )


@app.post("/embed", response_model=EmbedResponse)
async def embed_images(
    request: EmbedImagesRequest, _: Annotated[None, Header()] = None
) -> EmbedResponse:
    await require_token(_)
    _require_model()
    started = time.time()

    images = [_decode_image(payload, i) for i, payload in enumerate(request.images)]
    vectors = embed_images_gpu(images, request.batch_size)
    vectors = _truncate(vectors, request.dimension)

    return EmbedResponse(
        embeddings=vectors,
        dimension=len(vectors[0]) if vectors else 0,
        count=len(vectors),
        model=MODEL_ID,
        elapsed_ms=round((time.time() - started) * 1000, 2),
    )


@app.post("/embed_text", response_model=EmbedResponse)
async def embed_text(
    request: EmbedTextRequest, _: Annotated[None, Header()] = None
) -> EmbedResponse:
    await require_token(_)
    _require_model()
    started = time.time()

    vectors = embed_texts_gpu(request.texts, request.is_query, request.batch_size)
    vectors = _truncate(vectors, request.dimension)

    return EmbedResponse(
        embeddings=vectors,
        dimension=len(vectors[0]) if vectors else 0,
        count=len(vectors),
        model=MODEL_ID,
        elapsed_ms=round((time.time() - started) * 1000, 2),
    )


@app.get("/")
async def root() -> dict[str, object]:
    return {
        "service": "VisualOps Embedding Server (ZeroGPU)",
        "model": MODEL_ID,
        "dimension": _DIMENSION,
        "endpoints": ["/health", "/embed", "/embed_text"],
    }


if __name__ == "__main__":
    app.launch(show_error=True)