"""Small pieces shared by the streaming session and the server.

    StageSampling / resolve_sampling            which timesteps a body frame is denoised at
    wrapper_ctor_kwargs_from_ckpt_config        checkpoint config -> model constructor arguments
    cond_dict, frame_tokens                     conditioning dict, tokens per frame
    warm_z0_clean                               the clean write of the seed frame into memory
"""
from __future__ import annotations

from dataclasses import dataclass

import torch

from gta_streaming_kv_cache import SummaryKVCache
from noise_schedule import SampleGrid, sample_grid


# --------------------------------------------------------------------------------------
# Sampling plan
# --------------------------------------------------------------------------------------
@dataclass(frozen=True)
class StageSampling:
    """Where a body frame is denoised: a grid of timesteps and the order of model calls on it.

    Between two calls the current estimate of the clean frame is re-noised with FRESH noise to the next
    timestep (``summary_streaming_inference.renoise_body_x0``).
    """

    grid: SampleGrid
    # Grid indices the model is called at, in order, high -> low t, repeats allowed (a repeat re-noises the estimate
    # back to the same t and denoises again). None = every grid index once. Index 0 (t = 0) is never a call.
    step_plan: tuple[int, ...] | None = None

    def __post_init__(self):
        n = len(self.grid)
        if n < 2:
            raise ValueError(f"sampling grid too short: {self.grid.timesteps}")
        if self.step_plan is not None:
            plan = tuple(int(i) for i in self.step_plan)
            if not plan or any(i < 1 or i >= n for i in plan):
                raise ValueError(f"step_plan must be non-empty grid indices in 1..{n - 1}; got {plan}")
            if any(plan[k] < plan[k + 1] for k in range(len(plan) - 1)):
                raise ValueError(f"step_plan must be non-increasing (high -> low t); got {plan}")

    @property
    def indices(self) -> list[int]:
        """Grid indices of the model calls, in order."""
        if self.step_plan is not None:
            return [int(i) for i in self.step_plan]
        return list(range(len(self.grid) - 1, 0, -1))

    @property
    def timesteps(self) -> list[int]:
        """The timesteps (0..999) of those calls, e.g. ``[999, 878, 669]``."""
        return [int(self.grid.timesteps[i]) for i in self.indices]

    @property
    def label(self) -> str:
        ts = self.timesteps
        return f"{len(ts)} [{','.join(str(t) for t in ts)}]" if len(ts) <= 6 else f"{len(ts)}"


def resolve_sampling(config: dict, *, steps: int | None = None, tag: str = "sampling") -> StageSampling:
    """The sampling plan a checkpoint config asks for.

    ``steps`` (default: the config's own ``len(timesteps)``) is the number of model calls per frame. The grid is
    rebuilt from the config's schedule (``shifted_cosine_ratio``) and ``timestep_shift``; the caller checks it against
    the config's ``timesteps`` list.
    """
    if not isinstance(config, dict):
        raise TypeError(f"ckpt config must be a dict; got {type(config)}")
    if str(config["noise_schedule"]) != "shifted_cosine":
        raise ValueError(f"noise_schedule must be 'shifted_cosine'; got {config['noise_schedule']!r}")
    n = len(config["timesteps"]) if steps is None else int(steps)
    grid = sample_grid(steps=n, shift=float(config["timestep_shift"]), ratio=float(config["shifted_cosine_ratio"]))
    samp = StageSampling(grid=grid)
    print(f"[{tag}] steps={n} shift={float(config['timestep_shift']):g} t=[{','.join(str(t) for t in samp.timesteps)}]")
    return samp


# --------------------------------------------------------------------------------------
# Checkpoint config -> model
# --------------------------------------------------------------------------------------
_ARCH_KEYS = ("dim", "ffn_dim", "num_heads", "num_layers", "summary_tokens", "raw_window_frames", "summary_window_frames")


def num_layers_from_ckpt_config(cfg: dict) -> int:
    """Number of trunk blocks."""
    n = int(cfg["num_layers"])
    if n < 1:
        raise ValueError(f"num_layers must be >= 1; got {n}")
    return n


def wrapper_ctor_kwargs_from_ckpt_config(cfg: dict) -> dict:
    """``SummaryformerWrapper`` kwargs from a checkpoint config.

    Memory geometry is X / Y / W: ``summary_tokens`` / ``raw_window_frames`` / ``summary_window_frames``.
    """
    if not isinstance(cfg, dict):
        raise TypeError(f"ckpt config must be a dict; got {type(cfg)}")
    missing = [k for k in _ARCH_KEYS if k not in cfg]
    if missing:
        raise KeyError(f"ckpt config missing {missing}.")
    return dict(
        dim=int(cfg["dim"]),
        ffn_dim=int(cfg["ffn_dim"]),
        num_heads=int(cfg["num_heads"]),
        num_layers=num_layers_from_ckpt_config(cfg),
        summary_tokens=int(cfg["summary_tokens"]),
        raw_window_frames=int(cfg["raw_window_frames"]),
        summary_window_frames=int(cfg["summary_window_frames"]),
    )


# --------------------------------------------------------------------------------------
# Geometry / conditioning
# --------------------------------------------------------------------------------------
def frame_tokens(model, h: int, w: int) -> int:
    """Tokens of one frame on an ``h x w`` latent grid: one token per latent cell."""
    return int(h) * int(w)


def cond_dict(*, fps=None, action_context=None, z0_focus_region=None) -> dict:
    """``z0_focus_region`` is a seed-frame knob (Bool [B,H,W])."""
    out = {
        "fps": fps,
        "action_context": action_context,
    }
    if z0_focus_region is not None:
        out["z0_focus_region"] = z0_focus_region
    return out


@torch.no_grad()
def warm_z0_clean(wrapper, kv: SummaryKVCache, cond_pos: dict, xhat0: torch.Tensor) -> None:
    """Write the clean seed frame's record into the Z0 slot (one pass at t = 0), then ``kv.note_z0_written()``.

    ``xhat0 [B, 1, C, H, W]``. No ``begin_frame`` and no cursor commit: the seed frame is the sink, not a RAW slot.
    """
    if not (isinstance(xhat0, torch.Tensor) and xhat0.dim() == 5 and int(xhat0.size(1)) == 1):
        raise TypeError(f"xhat0 must be Tensor[B,1,C,H,W]; got {getattr(xhat0, 'shape', type(xhat0))}")
    t0 = torch.zeros((int(xhat0.size(0)),), device=xhat0.device, dtype=torch.float32)
    wrapper.forward_z0(x_z0=xhat0, conditional_dict=cond_pos, t_z0=t0, kv_cache=kv)
    kv.note_z0_written()