"""WeMM-JEV text and vision inference."""

from __future__ import annotations

import contextlib
import gc
import os
import sys
import threading
from pathlib import Path

from ._compat import Engine, Unsupported
from . import _images as vimg
from ._images import resize_policy

os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

_LOAD_LOCK = threading.Lock()


class D3Engine(Engine):
    name = "d3"
    latency = (
        "Device-synchronized in-process request wall time including prompt rendering and tokenization (and, "
        "for image and video requests, their decoding and preprocessing); one request per call, its questions in "
        "request order in batches of batch_size; excludes model loading."
    )

    def __init__(
        self,
        model: str | None = None,
        revision: str | None = None,
        device: str | None = None,
        batch_size: int = 8,
        verify: str = "fast",
        model_name: str | None = "WeMM-JEV",
        permutation_average: str | bool = False,
        max_pixels: int | None = vimg.MAX_PIXELS,
        image_root: str | None = None,
        replicas: int | None = None,
        full_warmup: str | bool = False,
        enable_videos: str | bool = False,
        max_batch_tokens: int = 8192,
        **options,
    ):
        if options:
            raise TypeError(f"unknown engine options {sorted(options)}")
        if not model:
            model = str(Path(__file__).absolute().parent.parent)
        flag = str(permutation_average).strip().lower()
        if flag not in ("0", "1", "false", "true"):
            raise ValueError("permutation_average must be 0 or 1")
        averaged = flag in ("1", "true")
        fw_flag = str(full_warmup).strip().lower()
        if fw_flag not in ("0", "1", "false", "true"):
            raise ValueError("full_warmup must be 0 or 1")
        self.full_warmup = fw_flag in ("1", "true")
        video_flag = str(enable_videos).strip().lower()
        if video_flag not in ("0", "1", "false", "true"):
            raise ValueError("enable_videos must be 0 or 1")
        self.enable_videos = video_flag in ("1", "true")

        import torch

        self.torch = torch
        self.model_id = str(model)
        self.revision = revision
        self.batch_size = int(batch_size)
        self.max_batch_tokens = int(max_batch_tokens)
        if self.batch_size < 1 or self.max_batch_tokens < 0:
            raise ValueError("batch_size must be positive and max_batch_tokens nonnegative")
        self.verify = verify
        self.model_name = model_name
        self.permutation_average = averaged
        self.max_pixels = max_pixels
        self.image_root = image_root

        gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
        if device in (None, "cuda"):
            primary_device = "cuda:0" if gpus > 0 else "cpu"
            self._multi_gpu = gpus > 1
            self._num_gpus = max(1, gpus)
        else:
            primary_device = str(device)
            self._multi_gpu = False
            self._num_gpus = 1

        super().__init__(
            model=self.model_id,
            revision=revision,
            device=primary_device,
            batch_size=self.batch_size,
            verify=verify,
            model_name=model_name,
            max_pixels=max_pixels,
            enable_videos=self.enable_videos,
            max_batch_tokens=self.max_batch_tokens,
            **({"permutation_average": 1} if averaged else {}),
        )

        hf_model, decision = self._load_replica(primary_device)
        self.model = hf_model
        self.decision = decision
        self._slots = {0: (hf_model, decision, threading.Lock())}
        self._local = threading.local()
        self._next_slot = 0
        self._slots_lock = threading.Lock()

        if replicas is not None:
            n_rep = int(replicas)
            if n_rep > 1:
                if gpus < n_rep:
                    raise RuntimeError(f"replicas={n_rep} needs {n_rep} CUDA devices, found {gpus}")
                self._multi_gpu = True
                self._num_gpus = n_rep
                for idx in range(1, n_rep):
                    m_i, d_i = self._load_replica(f"cuda:{idx}")
                    self._slots[idx] = (m_i, d_i, threading.Lock())

        self.provenance = {
            **self.decision.provenance(),
            "max_pixels": self.max_pixels,
            "max_batch_tokens": self.max_batch_tokens,
            "image_resize": resize_policy(self.max_pixels) + "; then the checkpoint's own processor",
        }

    def _load_replica(self, device_str: str):
        with _LOAD_LOCK:
            local = Path(os.fspath(self.model_id)).expanduser()
            if local.is_dir():
                pkg_dir = os.path.abspath(os.fspath(local))
                if pkg_dir not in sys.path:
                    sys.path.insert(0, pkg_dir)
            from transformers import AutoModel

            if device_str.startswith("cuda"):
                self.torch.cuda.set_device(device_str)
            hf_model = AutoModel.from_pretrained(
                self.model_id,
                trust_remote_code=True,
                revision=self.revision,
                device=device_str,
                batch_size=self.batch_size,
                verify=self.verify,
            )
            decision = getattr(hf_model, "runtime", None)
            if decision is None:
                raise ValueError(
                    f"{self.model_id!r} did not load as a d3 decision package (missing decision_config.json)"
                )
            if self.model_name:
                decision.model_name = self.model_name
            if self.permutation_average:
                decision.permutation_average = True
            # Text/image evaluation uses PIL and needs no video decoder. This also
            # prevents provenance().video_contract() from importing optional cv2.
            if not self.enable_videos and getattr(decision, "video_unavailable", None) is None:
                decision.video_unavailable = "disabled for text/image evaluation (enable_videos=0)"
            return hf_model, decision

    def _acquire_slot(self):
        slot = getattr(self._local, "slot", None)
        if slot is None:
            with self._slots_lock:
                slot = self._next_slot if self._multi_gpu else 0
                self._next_slot = (self._next_slot + 1) % max(1, self._num_gpus)
            self._local.slot = slot
        entry = self._slots.get(slot)
        if entry is None:
            with self._slots_lock:
                entry = self._slots.get(slot)
                if entry is None:
                    hf_model, decision = self._load_replica(f"cuda:{slot}")
                    entry = (hf_model, decision, threading.Lock())
                    self._slots[slot] = entry
        return entry

    @contextlib.contextmanager
    def _device_ctx(self, decision):
        dev = getattr(decision, "device", None)
        if dev is not None and getattr(dev, "type", None) == "cuda":
            with self.torch.cuda.device(dev):
                yield
        else:
            yield

    def warmup(self):
        with self._device_ctx(self.decision):
            super().warmup()
            if self.full_warmup or getattr(self.decision, "fast", None) is not None:
                self.warmup_seconds = self.decision.warmup()

    def runtime(self):
        return self.decision.runtime_info()

    def synchronize(self):
        for _, decision, _ in list(self._slots.values()):
            decision.synchronize()

    def _prepare_images(self, images, decision):
        if not images:
            return None, None
        from PIL import Image

        pil_images = []
        image_pixels = []
        for item in images:
            if isinstance(item, Image.Image):
                im = item.convert("RGB")
                if self.max_pixels and im.width * im.height > self.max_pixels:
                    im = im.resize(vimg.fit(im.width, im.height, self.max_pixels), Image.LANCZOS)
            else:
                im = vimg.pil(item, self.image_root, self.max_pixels).convert("RGB")
            pil_images.append(im)
            image_pixels.append(im.width * im.height)
        image_tokens = None
        proc = getattr(decision, "processor", None)
        if proc is not None and getattr(proc, "image_processor", None) is not None:
            ip = proc.image_processor
            merge = getattr(ip, "merge_size", 2) ** 2
            try:
                image_tokens = sum(
                    ip.get_number_of_image_patches(im.height, im.width, {}) // merge
                    for im in pil_images
                )
            except Exception:
                image_tokens = None
        raw = {"image_pixels": image_pixels}
        if image_tokens is not None:
            raw["image_tokens"] = image_tokens
        return pil_images, raw

    def __call__(self, state, questions, images=None, videos=None, prompt=None):
        if videos and not self.enable_videos:
            raise Unsupported("video evaluation requires --option enable_videos=1 and a working video decoder")
        _, decision, lock = self._acquire_slot()
        pil_images, raw = self._prepare_images(images, decision)
        with lock, self._device_ctx(decision):
            prepared = decision.prepare(state, questions, pil_images, videos)
            over = [
                e["message"]
                for e in prepared.errors.values()
                if e["error"] == "max_length_exceeded"
            ]
            if over:
                raise Unsupported(over[0])
            if prepared.errors:
                raise ValueError(
                    "invalid questions: "
                    + "; ".join(f"{k}: {e['message']}" for k, e in prepared.errors.items())
                )
            probabilities, tokens = self._run_prepared(decision, prepared)
            response = decision.respond(prepared, probabilities, tokens)
        failed = {
            k: a["message"] for k, a in response["answers"].items() if "error" in a
        }
        if failed:
            raise ValueError(f"invalid model output: {failed}")
        return response, raw

    def _run_prepared(self, decision, prepared):
        """Bound padded tokens per forward without truncating questions or context."""
        original = decision.batch_size
        sizes = list(getattr(prepared, "lengths", {}).values())
        sizes.extend(len(s) for s in getattr(prepared, "sequences", {}).values())
        sizes.extend(len(s) for s in getattr(prepared, "reversed_sequences", {}).values())
        sizes.extend(getattr(prepared, "reversed_lengths", {}).values())
        batch = original
        if self.max_batch_tokens and sizes:
            width = max(sizes) * (2 if self.permutation_average else 1)
            batch = min(batch, max(1, self.max_batch_tokens // max(1, width)))
        try:
            while True:
                decision.batch_size = batch
                try:
                    return decision.run(prepared)
                except self.torch.OutOfMemoryError:
                    if batch == 1:
                        raise
                # Leave the except block first to release the failed forward's traceback/tensors.
                batch = max(1, batch // 2)
                gc.collect()
                self.torch.cuda.empty_cache()
        finally:
            decision.batch_size = original

    def close(self):
        with self._slots_lock:
            self._slots.clear()
        self.model = None
        self.decision = None
        if self.torch.cuda.is_available():
            self.torch.cuda.empty_cache()