"""ScrollModel — ONNX Runtime-backed scroll event generation."""
from __future__ import annotations

import json
import math
from collections.abc import Sequence
from pathlib import Path
from typing import Literal, Optional

import numpy as np
import onnxruntime as ort

from ai_mouse._assets import resolve
from ai_mouse.errors import ModelLoadError

_DURATION_TABLE = {
    "fast":    lambda d: d * 0.2 + 100.0,
    "precise": lambda d: d * 1.5 + 300.0,
    "target":  lambda d: d * 0.4 + 200.0,
}

_QUANTUM = {"precise": 40, "fast": 120, "target": 120}


class ScrollModel:
    """Persistent ONNX Runtime session for scroll event generation."""

    def __init__(
        self,
        model_path: str | Path | None = None,
        providers: Sequence[str] | None = None,
        seed: int | None = None,
    ) -> None:
        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None

        onnx_path = resolve(path_obj, "scroll_decoder.onnx")
        cfg_path = resolve(path_obj, "scroll_config.json")
        cfg = json.loads(cfg_path.read_text())

        self._seq_len = int(cfg["seq_len"])
        self._latent_dim = int(cfg["latent_dim"])
        self._cond_dim = int(cfg["cond_dim"])

        try:
            self._session = ort.InferenceSession(
                str(onnx_path),
                providers=list(providers) if providers else ["CPUExecutionProvider"],
            )
        except Exception as exc:
            raise ModelLoadError(f"Failed to load scroll ONNX session: {exc}") from exc

        self._rng = np.random.default_rng(seed)

    def generate(
        self,
        start_scroll_y: int,
        target_scroll_y: int,
        mode: Literal["target", "fast", "precise"] = "target",
        viewport_height: int = 800,
        seed: int | None = None,
    ) -> list[dict]:
        rng = np.random.default_rng(seed) if seed is not None else self._rng

        distance = abs(target_scroll_y - start_scroll_y)
        direction = 1 if target_scroll_y > start_scroll_y else -1
        distance = max(distance, 10)

        cond = self._build_condition(float(distance), direction, mode, viewport_height)
        z = rng.standard_normal((1, self._latent_dim)).astype(np.float32)
        decoded = self._session.run(["seq"], {"z": z, "cond": cond[None]})[0][0]

        delta_norm = decoded[:, 0]
        log_dt = decoded[:, 1]

        delta_weights = np.exp(delta_norm)
        delta_weights /= delta_weights.sum()
        delta_px = delta_weights * distance * direction

        quantum = _QUANTUM[mode]
        delta_q = np.round(delta_px / quantum) * quantum
        for i in range(len(delta_q)):
            if delta_q[i] == 0:
                delta_q[i] = quantum * direction
        delta_q[-1] += (distance * direction) - delta_q.sum()

        if len(log_dt) > 3:
            median_log = float(np.median(log_dt))
            log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
            log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
        dt_ms = np.clip(np.exp(log_dt), 5, 80)
        expected = _DURATION_TABLE[mode](distance)
        dt_ms = np.clip(dt_ms * (expected / max(dt_ms.sum(), 1.0)), 5, 80)

        t_abs = np.cumsum(dt_ms).astype(int)
        t_abs = np.concatenate([[0], t_abs[:-1]])
        for i in range(1, len(t_abs)):
            if t_abs[i] <= t_abs[i - 1]:
                t_abs[i] = t_abs[i - 1] + 5

        events: list[dict] = []
        for i in range(self._seq_len):
            dy = int(delta_q[i])
            if dy != 0 or len(events) < 5:
                events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
        return events

    def _build_condition(
        self,
        distance: float,
        direction: int,
        mode: str,
        viewport_height: int,
    ) -> np.ndarray:
        mode_onehot = [0.0, 0.0, 0.0]
        if mode == "target":
            mode_onehot[0] = 1.0
        elif mode == "fast":
            mode_onehot[1] = 1.0
        elif mode == "precise":
            mode_onehot[2] = 1.0
        return np.array(
            [
                distance / 5000.0,
                math.log(max(distance, 1.0) / 500.0),
                float(direction),
                viewport_height / 1000.0,
                *mode_onehot,
            ],
            dtype=np.float32,
        )

    def close(self) -> None:
        self._session = None  # type: ignore[assignment]

    def __enter__(self) -> "ScrollModel":
        return self

    def __exit__(self, *exc) -> None:
        self.close()
