"""Scroll wheel event sequence generator."""
from __future__ import annotations

import json
import logging
import math
from pathlib import Path

import numpy as np
import torch

from tools.scroll.models import ScrollCVAE

logger = logging.getLogger(__name__)

_BUNDLED_SCROLL_MODELS = Path(__file__).resolve().parent.parent.parent / "data" / "scroll_models"


def _build_condition(
    distance: float,
    direction: int,
    mode: str,
    viewport_height: float = 900.0,
) -> np.ndarray:
    """Build 7-dim condition vector matching the trainer layout.

    Dims: [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3]
    """
    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

    viewport_norm = viewport_height / 1000.0

    return np.array([
        distance / 5000.0,
        math.log(max(distance, 1.0) / 500.0),
        float(direction),
        viewport_norm,
        *mode_onehot,
    ], dtype=np.float32)


def generate_scroll(
    start_scrollY: int,
    target_scrollY: int,
    mode: str = "target",
    model_dir: str | None = None,
) -> list[dict]:
    """Generate a realistic scroll event sequence.

    Args:
        start_scrollY: Current scroll position (px from top).
        target_scrollY: Target scroll position.
        mode: "target" | "fast" | "precise"
        model_dir: Path to scroll model files. None = bundled.

    Returns:
        List of {"deltaY": int, "deltaMode": 0, "t": int}.
        Positive deltaY = scroll down, negative = scroll up.
    """
    model_dir_path = Path(model_dir) if model_dir else _BUNDLED_SCROLL_MODELS
    model_pt = model_dir_path / "scroll_model.pt"
    config_json = model_dir_path / "scroll_config.json"

    if not model_pt.exists():
        raise FileNotFoundError(f"Scroll model not found at {model_pt}")

    seq_len = 32
    if config_json.exists():
        cfg = json.loads(config_json.read_text())
        seq_len = cfg.get("seq_len", 32)

    model = ScrollCVAE(seq_len=seq_len)
    model.load_state_dict(torch.load(model_pt, map_location="cpu", weights_only=True))
    model.eval()

    distance = abs(target_scrollY - start_scrollY)
    direction = 1 if target_scrollY > start_scrollY else -1
    distance = max(distance, 10)

    cond = _build_condition(float(distance), direction, mode)
    cond_t = torch.from_numpy(cond).unsqueeze(0)

    with torch.no_grad():
        z = torch.randn(1, model.latent_dim)
        decoded = model.decode(z, cond_t).squeeze(0).numpy()

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

    # De-normalise delta: use softmax-like normalisation so they sum to ~1
    delta_weights = np.exp(delta_norm)
    delta_weights = delta_weights / delta_weights.sum()
    delta_px = delta_weights * distance * direction

    # Quantise to realistic wheel increments
    quantum = 40 if mode == "precise" else 120

    delta_quantised = np.round(delta_px / quantum) * quantum
    for i in range(len(delta_quantised)):
        if delta_quantised[i] == 0:
            delta_quantised[i] = quantum * direction

    # Adjust last event so total matches target distance
    current_total = delta_quantised.sum()
    diff = (distance * direction) - current_total
    delta_quantised[-1] += diff

    # Timestamps from log_dt
    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.exp(log_dt).clip(5, 80)

    # Scale to realistic total duration
    if mode == "fast":
        expected_duration = distance * 0.2 + 100
    elif mode == "precise":
        expected_duration = distance * 1.5 + 300
    else:
        expected_duration = distance * 0.4 + 200

    dt_ms = dt_ms * (expected_duration / max(dt_ms.sum(), 1.0))
    dt_ms = dt_ms.clip(5, 80)

    t_abs = np.cumsum(dt_ms).astype(int)
    t_abs = np.concatenate([[0], t_abs[:-1]])

    # Ensure monotonic
    for i in range(1, len(t_abs)):
        if t_abs[i] <= t_abs[i - 1]:
            t_abs[i] = t_abs[i - 1] + 5

    # Build events, removing zero-delta (keep at least 5)
    events = []
    for i in range(seq_len):
        dy = int(delta_quantised[i])
        if dy != 0 or len(events) < 5:
            events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})

    return events
