"""
OpenCLIP ViT-B/32 embedder using ONNX Runtime.

Expects two ONNX files under {models_dir}/embed/:
  - visual.onnx   (image encoder)
  - textual.onnx  (text encoder)

These are exported from open_clip via bootstrap_models.py.
"""
import logging
from pathlib import Path

import numpy as np
import onnxruntime as ort

from app.config import VisionSettings
from app.services.vision.base import Embedder

logger = logging.getLogger(__name__)

# OpenCLIP ViT-B/32 preprocessing constants (ImageNet norm)
_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
_INPUT_SIZE = 224


def _preprocess_image(image: np.ndarray) -> np.ndarray:
    """Resize, center-crop, normalize an RGB uint8 image to NCHW float32."""
    from PIL import Image

    img = Image.fromarray(image).convert("RGB")
    # Resize shortest edge to _INPUT_SIZE, then center crop
    w, h = img.size
    scale = _INPUT_SIZE / min(w, h)
    img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
    w, h = img.size
    left = (w - _INPUT_SIZE) // 2
    top = (h - _INPUT_SIZE) // 2
    img = img.crop((left, top, left + _INPUT_SIZE, top + _INPUT_SIZE))

    arr = np.array(img, dtype=np.float32) / 255.0
    arr = (arr - _MEAN) / _STD
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    return arr[np.newaxis]  # NCHW


class OpenCLIPEmbedder(Embedder):
    def __init__(self, settings: VisionSettings):
        model_dir = Path(settings.models_dir) / "embed"
        visual_path = model_dir / "visual.onnx"
        textual_path = model_dir / "textual.onnx"

        opts = ort.SessionOptions()
        opts.inter_op_num_threads = 2
        opts.intra_op_num_threads = 2

        logger.info("Loading visual encoder from %s", visual_path)
        self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])

        logger.info("Loading textual encoder from %s", textual_path)
        self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])

    def embed_image(self, image: np.ndarray) -> np.ndarray:
        inp = _preprocess_image(image)
        input_name = self._visual.get_inputs()[0].name
        out = self._visual.run(None, {input_name: inp})[0][0]
        out = out / np.linalg.norm(out)
        return out.astype(np.float32)

    def embed_text(self, text: str) -> np.ndarray:
        import open_clip
        tokenizer = open_clip.get_tokenizer("ViT-B-32")
        tokens = tokenizer([text]).numpy().astype(np.int64)
        input_name = self._textual.get_inputs()[0].name
        out = self._textual.run(None, {input_name: tokens})[0][0]
        out = out / np.linalg.norm(out)
        return out.astype(np.float32)

    @property
    def dim(self) -> int:
        return 512