"""MouseModel — ONNX Runtime-backed mouse trajectory generation."""
from __future__ import annotations

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

import numpy as np
import onnxruntime as ort

from ai_mouse._assets import resolve
from ai_mouse._coord import decode_trajectory
from ai_mouse._postprocess import (
    build_timestamps,
    damp_start,
    gaussian_smooth,
    resample_arc,
    sample_duration,
    soften_forward,
    truncnorm_sample,
    warp_endpoints,
)
from ai_mouse.errors import GenerationError, ModelLoadError

_N_EULER_STEPS = 10


class MouseModel:
    """Persistent ONNX Runtime session for mouse trajectory 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, "flow_model.onnx")
        cfg_path = resolve(path_obj, "train_config.json")
        click_path = resolve(path_obj, "click_dist.json")
        dur_path = resolve(path_obj, "duration_dist.json")

        cfg = json.loads(cfg_path.read_text())
        self._seq_len = int(cfg["seq_len"])
        self._cond_dim = int(cfg.get("cond_dim", 3))
        self._click_params = json.loads(click_path.read_text())
        self._duration_dist = json.loads(dur_path.read_text())

        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 ONNX session: {exc}") from exc

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

    def generate(
        self,
        start: tuple[int, int],
        end: tuple[int, int],
        n_points: int = 64,
        speed: float | None = None,
        click: bool = True,
        seed: int | None = None,
    ) -> list[tuple[int, int, int]]:
        rng = np.random.default_rng(seed) if seed is not None else self._rng

        sx, sy = float(start[0]), float(start[1])
        ex, ey = float(end[0]), float(end[1])
        dist = max(math.hypot(ex - sx, ey - sy), 1.0)

        total_duration = sample_duration(self._duration_dist, dist, rng)
        if speed is not None and speed > 0:
            total_duration /= speed
        total_duration = max(total_duration, 10.0)

        cond = np.array(
            [
                dist / 2000.0,
                math.log(dist / 100.0),
                math.log(total_duration / 500.0),
            ],
            dtype=np.float32,
        )[None]

        x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
        dt = 1.0 / _N_EULER_STEPS
        for step in range(_N_EULER_STEPS):
            t = np.full((1,), step * dt, dtype=np.float32)
            v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0]
            x = x + v * dt

        if not np.all(np.isfinite(x)):
            raise GenerationError("Trajectory contains NaN/Inf after Euler integration")

        forward = x[0, :, 0].copy()
        lateral = x[0, :, 1].copy()
        log_dt = x[0, :, 2].copy()

        forward = soften_forward(forward)
        lateral = damp_start(lateral)
        forward = gaussian_smooth(forward, sigma=1.0)
        lateral = gaussian_smooth(lateral, sigma=1.0)
        forward, lateral = warp_endpoints(forward, lateral)

        log_dt = np.clip(log_dt, 0.0, 5.0)
        log_dt[0] = 0.0

        normalised = np.stack([forward, lateral], axis=1)
        pixels = decode_trajectory(normalised, start, end)

        if n_points != self._seq_len:
            pixels = resample_arc(pixels, n_points)
            log_dt = np.interp(
                np.linspace(0, 1, n_points),
                np.linspace(0, 1, self._seq_len),
                log_dt,
            )

        ts = build_timestamps(log_dt, total_duration)

        moves: list[tuple[int, int, int]] = [
            (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i])))
            for i in range(n_points)
        ]
        if not click:
            return moves

        click_dur = int(
            truncnorm_sample(
                float(self._click_params["mu"]),
                float(self._click_params["sigma"]),
                float(self._click_params["low"]),
                float(self._click_params["high"]),
                rng,
            )
        )
        click_dur = max(click_dur, int(float(self._click_params["low"])))
        last_t = moves[-1][2]
        cx, cy = moves[-1][0], moves[-1][1]
        return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)]

    def sample_click_duration_ms(self, seed: int | None = None) -> int:
        rng = np.random.default_rng(seed) if seed is not None else self._rng
        v = truncnorm_sample(
            float(self._click_params["mu"]),
            float(self._click_params["sigma"]),
            float(self._click_params["low"]),
            float(self._click_params["high"]),
            rng,
        )
        return max(int(v), int(float(self._click_params["low"])))

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

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

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