"""Decision Index engine for a d3 package loaded through d3q (text and the public vision suite).

    PYTHONPATH=kit_engine python -m decision_index run --engine d3q_engine:D3QEngine \
        --option model=~/models/d3-NVFP4 --out runs/d3-NVFP4 --compact

The request path is the package's own d3_engine.D3Engine (prepare, run, respond; batch_size questions per forward
pass, request order), so results compare one to one with the official BF16 run. Only loading differs:
d3q.load.load_d3 keeps INT4/NVFP4 weights packed and runs them on the 4-bit kernels, and loads a BF16 package with
the stock runtime. The package's own d3q/ is used when it ships one, otherwise this project's.

The vision runner passes ``image_root`` (row images are relative to it) and a per-row ``prompt``: the prompt is the
letter-readout cue for generative engines and is ignored here, since d3 reads every option's probability directly.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

from decision_index.engines import Engine, Unsupported

PROJECT = Path(__file__).resolve().parent.parent
# Keep Triton autotune results on disk, so later processes reuse them (read when the kernels are imported).
os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")


class D3QEngine(Engine):
    name = "d3"
    latency = (
        "Device-synchronized in-process request wall time including prompt rendering and tokenization (and, "
        "for image 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=None, device="cuda:0", batch_size=8, image_root=None, **options):
        if options:
            raise TypeError(f"unknown engine options {sorted(options)}")
        if not model:
            raise ValueError("pass --option model=<package dir>")
        package = Path(model).expanduser().resolve()
        # D3Q_MEM_FRACTION=<0..1>: cap this process's share of GPU memory when several runs share one GPU, so a long
        # request hits d3q's split-on-OOM inside its own share instead of starving the other processes.
        fraction = os.environ.get("D3Q_MEM_FRACTION")
        if fraction:
            import torch

            torch.cuda.set_per_process_memory_fraction(float(fraction), torch.device(device))
        sys.path.insert(0, str(package if (package / "d3q").is_dir() else PROJECT))
        from d3q.load import load_d3

        super().__init__(model=str(package), device=device, batch_size=batch_size)
        self.decision = load_d3(package, device=device, batch_size=int(batch_size))
        self.provenance = self.decision.provenance()
        self.image_root = Path(image_root) if image_root else None
        import torch

        self.after_load_bytes = torch.cuda.memory_allocated(device)
        torch.cuda.reset_peak_memory_stats(device)
        self.device = device

    def warmup(self):
        super().warmup()
        self.warmup_seconds = self.decision.warmup()

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

    def synchronize(self):
        self.decision.synchronize()

    def close(self):
        # D3Q_MEMORY_LOG=<path>: record GPU memory after loading and at peak over the run, for the model card.
        path = os.environ.get("D3Q_MEMORY_LOG")
        if path:
            import json

            import torch

            Path(path).write_text(json.dumps({
                "after_load_gib": round(self.after_load_bytes / 2**30, 2),
                "peak_gib": round(torch.cuda.max_memory_allocated(self.device) / 2**30, 2),
                "gpu": torch.cuda.get_device_name(self.device),
            }))

    def __call__(self, state, questions, images=None, videos=None, prompt=None):
        if images and self.image_root:
            images = [
                str(self.image_root / i) if isinstance(i, str) and not os.path.isabs(i) and "://" not in i else i
                for i in images
            ]
        prepared = self.decision.prepare(state, questions, 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.decision.run(prepared)
        response = self.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, None