"""Strict no-box full-frame Qwen AV sidecar contract and loader."""

from __future__ import annotations

import math
from pathlib import Path
from typing import Any

import torch

from local_turn_features import FramePersonFeatures, reference_alignment


FULLFRAME_AV_FORMAT = "qwen-fullframe-av-tokens-v1"


def load_fullframe_av_features(
    path: Path,
    reference: dict[str, Any],
    *,
    sample_id: str,
    dataset: str,
    expected_contract: dict[str, Any] | None = None,
    history_max_gap_seconds: float | None = None,
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
    """Load aligned Qwen audio and ragged full-frame spatial token sets."""

    times, indices = reference_alignment(reference)
    payload = torch.load(path, map_location="cpu", weights_only=False)
    if not isinstance(payload, dict) or payload.get("format") != FULLFRAME_AV_FORMAT:
        raise ValueError(f"Invalid full-frame feature format: {path}")
    if payload.get("sample_id") != sample_id or payload.get("dataset") != dataset:
        raise ValueError(f"Full-frame sample/dataset mismatch: {path}")
    if dataset != "easycom":
        raise ValueError("Full-frame F2/F3 currently support EasyCom only")
    if payload.get("partial", False):
        raise ValueError(f"Partial smoke-test features cannot be used for training: {path}")
    local_indices = torch.as_tensor(payload.get("label_indices", []))
    local_times = torch.as_tensor(payload.get("times", []), dtype=torch.float64)
    if local_indices.dtype != torch.long or not torch.equal(indices, local_indices):
        raise ValueError(f"Full-frame label_indices mismatch: {path}")
    if local_times.shape != times.shape or not torch.allclose(
        local_times, times, atol=1e-7, rtol=0
    ):
        raise ValueError(f"Full-frame times mismatch: {path}")
    config = payload.get("config", {})
    if (
        config.get("causal_policy") != "strict-half-open-fullframe-grid"
        or config.get("uses_speech_labels") is not False
        or config.get("uses_boxes") is not False
        or config.get("visual_feature_kind")
        != "qwen-vision-post-merger-spatial-token"
        or config.get("visual_source")
        != "easycom-full-frame-after-session-banner-crop"
        or config.get("audio_source") != "reused-qwen-local-mixed-audio"
    ):
        raise ValueError(f"Full-frame AV features lack the causal/no-box contract: {path}")
    contract = {"format": FULLFRAME_AV_FORMAT, "modalities": "av", "config": config}
    if expected_contract is not None and contract != expected_contract:
        raise ValueError(f"Full-frame contract differs across sessions/checkpoint: {path}")

    audio = payload.get("audio")
    audio_valid = payload.get("audio_valid")
    audio_dim = config.get("audio_dim")
    if not isinstance(audio_dim, int) or audio_dim <= 0:
        raise ValueError(f"Missing audio_dim: {path}")
    if (
        not isinstance(audio, torch.Tensor)
        or audio.shape != (len(times), audio_dim)
        or not audio.is_floating_point()
        or not torch.isfinite(audio).all()
    ):
        raise ValueError(f"Invalid full-frame sidecar audio: {path}")
    if (
        not isinstance(audio_valid, torch.Tensor)
        or audio_valid.dtype != torch.bool
        or audio_valid.shape != (len(times),)
    ):
        raise ValueError(f"Invalid full-frame audio_valid mask: {path}")
    if torch.count_nonzero(audio[~audio_valid]):
        raise ValueError(f"Invalid audio rows must be stored as zeros: {path}")

    visual_dim = config.get("visual_dim")
    visual_fps = config.get("visual_fps")
    tokens = payload.get("visual_tokens")
    if not isinstance(visual_dim, int) or visual_dim <= 0:
        raise ValueError(f"Missing visual_dim: {path}")
    if not isinstance(visual_fps, (int, float)) or not math.isfinite(visual_fps) or visual_fps <= 0:
        raise ValueError(f"Invalid visual_fps: {path}")
    if not isinstance(tokens, dict):
        raise ValueError(f"Missing full-frame visual tokens: {path}")
    features = tokens.get("features")
    frame_times = tokens.get("frame_times")
    frame_offsets = tokens.get("frame_offsets")
    grid_shapes = tokens.get("grid_shapes")
    source_clip_indices = tokens.get("frame_source_clip_indices")
    if (
        not isinstance(features, torch.Tensor)
        or features.ndim != 2
        or features.shape[1] != visual_dim
        or not features.is_floating_point()
        or not torch.isfinite(features).all()
    ):
        raise ValueError(f"Invalid full-frame token features: {path}")
    if (
        not isinstance(frame_times, torch.Tensor)
        or frame_times.ndim != 1
        or not frame_times.is_floating_point()
    ):
        raise ValueError(f"Invalid full-frame frame_times: {path}")
    frame_times = frame_times.to(dtype=torch.float64)
    if (
        not len(frame_times)
        or not torch.isfinite(frame_times).all()
        or bool((frame_times < 0).any())
        or bool((frame_times[1:] <= frame_times[:-1]).any())
    ):
        raise ValueError(f"Full-frame times must be finite/increasing: {path}")
    expected_frame_count = max(
        0,
        math.ceil(float(times[-1]) * float(visual_fps) - 1e-8),
    )
    expected_frame_times = (
        torch.arange(expected_frame_count, dtype=torch.float64)
        / float(visual_fps)
    )
    if frame_times.shape != expected_frame_times.shape or not torch.allclose(
        frame_times, expected_frame_times, atol=1e-8, rtol=0
    ):
        raise ValueError(f"Full-frame times are not the fixed causal grid: {path}")
    if (
        not isinstance(frame_offsets, torch.Tensor)
        or frame_offsets.dtype != torch.long
        or frame_offsets.shape != (len(frame_times) + 1,)
        or int(frame_offsets[0]) != 0
        or int(frame_offsets[-1]) != len(features)
        or bool((frame_offsets[1:] <= frame_offsets[:-1]).any())
    ):
        raise ValueError(f"Invalid full-frame frame_offsets: {path}")
    if (
        not isinstance(grid_shapes, torch.Tensor)
        or grid_shapes.dtype != torch.long
        or grid_shapes.shape != (len(frame_times), 2)
        or bool((grid_shapes <= 0).any())
    ):
        raise ValueError(f"Invalid full-frame grid_shapes: {path}")
    if not torch.equal(frame_offsets[1:] - frame_offsets[:-1], grid_shapes.prod(-1)):
        raise ValueError(f"Token counts do not match full-frame grids: {path}")
    if (
        not isinstance(source_clip_indices, torch.Tensor)
        or source_clip_indices.dtype != torch.long
        or source_clip_indices.shape != (len(frame_times),)
        or bool((source_clip_indices < 0).any())
        or bool((source_clip_indices[1:] < source_clip_indices[:-1]).any())
    ):
        raise ValueError(f"Invalid full-frame source-clip indices: {path}")
    if float(frame_times[-1]) >= float(times[-1]) - 1e-8:
        raise ValueError(f"Full-frame timeline extends past the reference cutoff: {path}")

    row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1
    chosen = torch.full(times.shape, float("nan"), dtype=torch.float64)
    valid_rows = row_to_frame >= 0
    chosen[valid_rows] = frame_times[row_to_frame[valid_rows]]
    if bool((chosen[valid_rows] >= times[valid_rows]).any()):
        raise ValueError(f"Full-frame mapping selected boundary/future evidence: {path}")

    previous = None
    if history_max_gap_seconds is not None:
        if not math.isfinite(history_max_gap_seconds) or history_max_gap_seconds <= 0:
            raise ValueError("history_max_gap_seconds must be finite and positive")
        previous = torch.full((len(features),), -1, dtype=torch.long)
        for frame in range(1, len(frame_times)):
            gap = float(frame_times[frame] - frame_times[frame - 1])
            same_clip = bool(source_clip_indices[frame] == source_clip_indices[frame - 1])
            same_grid = torch.equal(grid_shapes[frame], grid_shapes[frame - 1])
            if not same_clip or not same_grid or gap > history_max_gap_seconds + 1e-8:
                continue
            current_start, current_end = int(frame_offsets[frame]), int(frame_offsets[frame + 1])
            old_start, old_end = int(frame_offsets[frame - 1]), int(frame_offsets[frame])
            if current_end - current_start != old_end - old_start:
                raise AssertionError("Equal grids have unequal token counts")
            previous[current_start:current_end] = torch.arange(old_start, old_end)

    frames = FramePersonFeatures(
        features,
        frame_offsets,
        row_to_frame.to(torch.long),
        chosen,
        previous,
    )
    packed_audio = torch.cat((audio.float(), audio_valid[:, None].float()), dim=-1)
    return packed_audio, frames, contract