"""
Binary content classifier: 'photography' vs 'other'.

Uses OpenCLIP ViT-B/32 image features (ONNX) and two pre-computed text
prompt centroids. Text centroids are computed once with the native
open_clip text encoder and cached to {models_dir}/classifier/vectors.npz
so steady-state worker startup doesn't pay the PyTorch cost.
"""
from __future__ import annotations

import logging
from pathlib import Path

import numpy as np

from app.config import VisionSettings
from app.services.vision.base import ClassificationResult, ContentClassifier
from app.services.vision.embed import CLIPVisualEncoder

logger = logging.getLogger(__name__)


PROMPTS = {
    "photography": [
        "a photograph taken with a camera",
        "a real photo of a real scene or person",
        "a candid photograph",
        "a portrait photograph",
        "a landscape photograph",
    ],
    "other": [
        "a screenshot of a computer screen",
        "a screenshot of a phone screen",
        "a screen capture of a user interface",
        "a scanned document",
        "a photo of a document with printed text",
        "a photo of a receipt",
        "a photo of a bill or invoice",
        "an internet meme with text overlay",
        "a funny image with caption text",
        "a digital illustration or graphic design",
    ],
}


def _compute_text_centroids() -> dict[str, np.ndarray]:
    """Compute the 'photography' and 'other' centroid vectors using the
    open_clip text encoder. Only called on the cache-miss path."""
    import open_clip
    import torch

    logger.info("Computing CLIP text centroids for binary classifier")
    model, _, _ = open_clip.create_model_and_transforms(
        "ViT-B-32", pretrained="laion2b_s34b_b79k"
    )
    model.eval()
    tokenizer = open_clip.get_tokenizer("ViT-B-32")

    centroids: dict[str, np.ndarray] = {}
    for label, prompts in PROMPTS.items():
        tokens = tokenizer(prompts)
        with torch.no_grad():
            feats = model.encode_text(tokens)
            feats = feats / feats.norm(dim=-1, keepdim=True)
            avg = feats.mean(dim=0)
            avg = avg / avg.norm()
        centroids[label] = avg.numpy().astype(np.float32)
    return centroids


class CLIPContentClassifier(ContentClassifier):
    def __init__(self, settings: VisionSettings):
        self._min_confidence = settings.classifier.min_confidence
        self._encoder = CLIPVisualEncoder(settings)

        cache_dir = Path(settings.models_dir) / "classifier"
        cache_dir.mkdir(parents=True, exist_ok=True)
        cache_path = cache_dir / "vectors.npz"

        if cache_path.exists():
            logger.info("Loading cached text centroids from %s", cache_path)
            data = np.load(cache_path)
            self._photo = data["photography"].astype(np.float32)
            self._other = data["other"].astype(np.float32)
        else:
            centroids = _compute_text_centroids()
            self._photo = centroids["photography"]
            self._other = centroids["other"]
            np.savez(cache_path, photography=self._photo, other=self._other)
            logger.info("Cached text centroids to %s", cache_path)

    def classify(self, image: np.ndarray) -> ClassificationResult:
        vec = self._encoder.encode(image)
        s_photo = float(np.dot(vec, self._photo))
        s_other = float(np.dot(vec, self._other))

        if s_photo >= s_other:
            label = "photography"
            margin = s_photo - s_other
        else:
            label = "other"
            margin = s_other - s_photo

        # 0.01 margin → ~0.3 conf, 0.03+ → ~1.0
        confidence = min(1.0, margin * 30)
        return ClassificationResult(label=label, confidence=confidence)