"""embedFrames and classifyContent with SigLIP (Apache-2.0) on ONNX Runtime.

classifyContent is zero-shot: each category is a set of text prompts and the
score is the highest sigmoid probability among them. It is reported as
experimental; the gate thresholds decide what the scores mean.
"""

from __future__ import annotations

import threading
from typing import Any

import cv2
import numpy as np

from .. import media_io
from ..schemas import ClassifyContentRequest, TimesRequest
from .common import Ctx, clamp01, open_input, unique_sorted
from .devices import onnx_providers

SIGLIP_SIZE = 224
SIGLIP_TEXT_LENGTH = 64
SIGLIP_PAD_ID = 1
BATCH = 16

CATEGORY_PROMPTS: dict[str, tuple[str, ...]] = {
    "explicit_nudity": ("a photo of a naked person", "a photo containing explicit nudity"),
    "suggestive": ("a photo of a person in lingerie", "a sexually suggestive photo"),
    "violence": ("a photo of a violent fight", "a photo showing blood and injuries"),
    "minors_risk": ("a photo of a child", "a photo of a young teenager"),
    "self_harm": ("a photo of self-harm", "a photo of cuts on a wrist"),
    "drugs": ("a photo of illegal drugs", "a photo of drug paraphernalia"),
    "weapons": ("a photo of a gun", "a photo of a person holding a knife as a weapon"),
}


def preprocess(frames: list[np.ndarray]) -> np.ndarray:
    """BGR frames -> SigLIP pixel_values (N, 3, 224, 224), normalized to [-1, 1]."""
    batch = []
    for frame in frames:
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(rgb, (SIGLIP_SIZE, SIGLIP_SIZE), interpolation=cv2.INTER_CUBIC)
        batch.append(resized.astype(np.float32) / 127.5 - 1.0)
    return np.ascontiguousarray(np.stack(batch).transpose(0, 3, 1, 2))


def _session(path: str, settings: Any) -> Any:
    import onnxruntime

    return onnxruntime.InferenceSession(path, providers=onnx_providers(settings))


class SiglipEmbedder:
    model = "siglip-base-patch16-224"

    def __init__(self, path: str, settings: Any) -> None:
        self.session = _session(path, settings)
        names = [o.name for o in self.session.get_outputs()]
        self.output = "pooler_output" if "pooler_output" in names else names[-1]
        self.dim = 768

    def embed(self, frames: list[np.ndarray]) -> np.ndarray:
        (vectors,) = self.session.run([self.output], {"pixel_values": preprocess(frames)})
        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        return vectors / np.maximum(norms, 1e-12)


def load_frame_embedder(files: dict[str, str], settings: Any) -> SiglipEmbedder:
    return SiglipEmbedder(files["vision_model_quantized.onnx"], settings)


class SiglipZeroShot:
    model = "siglip-base-patch16-224-zero-shot"

    def __init__(self, model_path: str, tokenizer_path: str, settings: Any) -> None:
        from tokenizers import Tokenizer

        tokenizer = Tokenizer.from_file(tokenizer_path)
        tokenizer.enable_truncation(SIGLIP_TEXT_LENGTH)
        tokenizer.enable_padding(length=SIGLIP_TEXT_LENGTH, pad_id=SIGLIP_PAD_ID)
        self.labels: list[str] = []
        prompts: list[str] = []
        for category, texts in CATEGORY_PROMPTS.items():
            for text in texts:
                self.labels.append(category)
                prompts.append(text.lower())
        self.input_ids = np.array([e.ids for e in tokenizer.encode_batch(prompts)], dtype=np.int64)
        self.session = _session(model_path, settings)

    def classify(self, frames: list[np.ndarray]) -> list[dict[str, float]]:
        (logits,) = self.session.run(
            ["logits_per_image"], {"input_ids": self.input_ids, "pixel_values": preprocess(frames)}
        )
        probs = 1.0 / (1.0 + np.exp(-logits))
        result = []
        for row in probs:
            scores: dict[str, float] = {}
            for label, value in zip(self.labels, row, strict=True):
                scores[label] = max(scores.get(label, 0.0), clamp01(float(value)))
            result.append(scores)
        return result


def load_content_classifier(files: dict[str, str], settings: Any) -> SiglipZeroShot:
    return SiglipZeroShot(files["model_quantized.onnx"], files["tokenizer.json"], settings)


def _batched(frames: Any, size: int) -> Any:
    times: list[int] = []
    batch: list[np.ndarray] = []
    for t_ms, frame in frames:
        times.append(t_ms)
        batch.append(frame)
        if len(batch) == size:
            yield times, batch
            times, batch = [], []
    if batch:
        yield times, batch


def embed_run(ctx: Ctx, body: TimesRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    embedder = ctx.registry.runtime("embedFrames")
    vectors: list[list[float]] = []
    # Vectors follow the request order, duplicates included.
    for _, batch in _batched(media_io.frames_at(real, info, list(body.timesMs), cancel, 512), BATCH):
        vectors.extend([float(v) for v in row] for row in embedder.embed(batch))
    return {"model": embedder.model, "dim": int(embedder.dim), "vectors": vectors}


def classify_run(ctx: Ctx, body: ClassifyContentRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    classifier = ctx.registry.runtime("classifyContent")
    results: dict[int, dict[str, float]] = {}

    def score(frames: Any) -> None:
        for times, batch in _batched(frames, BATCH):
            for t_ms, scores in zip(times, classifier.classify(batch), strict=True):
                results[t_ms] = scores

    score(media_io.stream_frames(real, info, cancel, body.sampleFps, 512))
    extra = [t for t in unique_sorted(body.extraTimesMs or []) if t not in results]
    score(media_io.frames_at(real, info, extra, cancel, 512))
    return {
        "model": classifier.model,
        "frames": [{"tMs": t, "scores": results[t]} for t in sorted(results)],
    }