"""Cached open-stream AV pieces for Qwen2.5-Omni turn-taking LoRA.

The modality towers are frozen for the first fine-tuning experiment.  This
module serializes their exact output embeddings and Qwen's local TMRoPE
coordinates, then reconstructs the same still-open media turn without using a
KV cache.  The resulting sequence can be passed through the Thinker with
``use_cache=False`` so gradients reach LoRA adapters in every retained block.
"""

from __future__ import annotations

import json
import math
from collections import OrderedDict
from pathlib import Path
from typing import Any, Iterable, Sequence

import torch


SCHEMA_VERSION = 1
PIECE_TENSOR_NAMES = (
    "prefix_embeddings",
    "prefix_position_ids",
    "payload_embeddings",
    "payload_position_ids",
    "closure_ids",
    "closure_position_ids",
    "closure_embeddings",
)


def load_json(path: str | Path) -> dict[str, Any]:
    return json.loads(Path(path).read_text())


def save_json(path: str | Path, payload: dict[str, Any]) -> None:
    Path(path).write_text(json.dumps(payload, indent=2) + "\n")


def torch_load(path: str | Path) -> Any:
    """Load tensor-only cache files across old and new PyTorch defaults."""

    try:
        return torch.load(Path(path), map_location="cpu", weights_only=True)
    except TypeError:
        return torch.load(Path(path), map_location="cpu")


def cache_block_filename(block_index: int) -> str:
    return f"block_{block_index:06d}.pt"


def pieces_to_cache(
    pieces: Any,
    *,
    closure_embeddings: torch.Tensor,
    floating_dtype: torch.dtype = torch.bfloat16,
) -> dict[str, torch.Tensor]:
    """Detach one ``OpenMediaPieces`` object into a portable CPU payload."""

    output: dict[str, torch.Tensor] = {}
    for name in PIECE_TENSOR_NAMES:
        value = closure_embeddings if name == "closure_embeddings" else getattr(pieces, name)
        value = value.detach().cpu().contiguous()
        if value.is_floating_point():
            value = value.to(floating_dtype)
        output[name] = value
    validate_piece(output)
    return output


def validate_piece(piece: dict[str, torch.Tensor]) -> None:
    missing = [name for name in PIECE_TENSOR_NAMES if name not in piece]
    if missing:
        raise ValueError(f"Cached AV piece is missing {missing}")
    prefix = piece["prefix_embeddings"]
    payload = piece["payload_embeddings"]
    closure = piece["closure_embeddings"]
    if prefix.ndim != 3 or prefix.shape[0] != 1:
        raise ValueError(f"Invalid prefix embeddings: {tuple(prefix.shape)}")
    if payload.ndim != 3 or payload.shape[0] != 1:
        raise ValueError(f"Invalid payload embeddings: {tuple(payload.shape)}")
    if closure.ndim != 3 or closure.shape[0] != 1:
        raise ValueError(f"Invalid closure embeddings: {tuple(closure.shape)}")
    hidden_sizes = {int(prefix.shape[-1]), int(payload.shape[-1]), int(closure.shape[-1])}
    if len(hidden_sizes) != 1:
        raise ValueError(f"Cached piece hidden sizes disagree: {sorted(hidden_sizes)}")
    expected_positions = (
        ("prefix_position_ids", prefix.shape[1]),
        ("payload_position_ids", payload.shape[1]),
        ("closure_position_ids", closure.shape[1]),
    )
    for name, length in expected_positions:
        value = piece[name]
        if value.shape != (3, 1, length):
            raise ValueError(
                f"{name} must have shape {(3, 1, length)}, got {tuple(value.shape)}"
            )
    if piece["closure_ids"].shape != (1, closure.shape[1]):
        raise ValueError(
            "closure_ids/closure_embeddings mismatch: "
            f"{tuple(piece['closure_ids'].shape)} vs {tuple(closure.shape)}"
        )


def make_anchor(
    *,
    step_index: int,
    steps_per_block: int,
    memory_blocks: int,
    step_seconds: float,
    num_target_steps: int,
    num_horizons: int,
) -> dict[str, Any]:
    """Describe the exact completed-history/current-tail selection at time t."""

    if step_index < 0:
        raise ValueError("step_index must be non-negative")
    block_index = step_index // steps_per_block
    tail_steps = step_index % steps_per_block + 1
    first_history = max(0, block_index - (memory_blocks - 1))
    history_blocks = list(range(first_history, block_index))
    return {
        "step_index": int(step_index),
        "time_seconds": float((step_index + 1) * step_seconds),
        "block_index": int(block_index),
        "tail_steps": int(tail_steps),
        "history_blocks": history_blocks,
        "target_index": int(step_index),
        "has_future_target": bool(step_index + num_horizons < num_target_steps),
    }


class CachedBlockStore:
    """Small CPU LRU for block shard files shared by adjacent samples."""

    def __init__(self, cache_dir: str | Path, max_blocks: int = 16) -> None:
        if max_blocks <= 0:
            raise ValueError("max_blocks must be positive")
        self.cache_dir = Path(cache_dir)
        self.max_blocks = int(max_blocks)
        self._blocks: OrderedDict[int, dict[str, Any]] = OrderedDict()

    def get(self, block_index: int) -> dict[str, Any]:
        cached = self._blocks.get(block_index)
        if cached is not None:
            self._blocks.move_to_end(block_index)
            return cached
        path = self.cache_dir / cache_block_filename(block_index)
        if not path.is_file():
            raise FileNotFoundError(path)
        payload = torch_load(path)
        if int(payload.get("block_index", -1)) != block_index:
            raise ValueError(f"Block index mismatch in {path}")
        variants = payload.get("variants")
        if not isinstance(variants, dict) or not variants:
            raise ValueError(f"No variants in {path}")
        self._blocks[block_index] = payload
        self._blocks.move_to_end(block_index)
        while len(self._blocks) > self.max_blocks:
            self._blocks.popitem(last=False)
        return payload

    def variant(self, block_index: int, tail_steps: int) -> dict[str, torch.Tensor]:
        block = self.get(block_index)
        variants = block["variants"]
        piece = variants.get(str(tail_steps), variants.get(tail_steps))
        if piece is None:
            raise KeyError(
                f"Block {block_index} has no tail_steps={tail_steps}; "
                f"available={sorted(str(key) for key in variants)}"
            )
        validate_piece(piece)
        return piece


def _strip_embedding_batch(value: torch.Tensor) -> torch.Tensor:
    if value.ndim != 3 or value.shape[0] != 1:
        raise ValueError(f"Expected [1,L,D], got {tuple(value.shape)}")
    return value[0]


def _strip_position_batch(value: torch.Tensor) -> torch.LongTensor:
    if value.ndim != 3 or value.shape[:2] != (3, 1):
        raise ValueError(f"Expected [3,1,L], got {tuple(value.shape)}")
    return value[:, 0, :].long()


def reconstruct_open_stream_sequence(
    *,
    store: CachedBlockStore,
    static: dict[str, Any],
    anchor: dict[str, Any],
    current_piece: dict[str, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.LongTensor]:
    """Materialize one differentiable Thinker input from frozen AV embeddings.

    Complete history blocks use their 2.0-second variant.  Only the current
    block uses its 0.2--2.0-second phase.  The first retained block keeps the
    chat/media prefix; later blocks append payload only.  The final block's
    media closure and the artificial assistant readout query are appended once.
    """

    steps_per_block = int(static["steps_per_block"])
    positions_per_chunk = int(static["positions_per_chunk"])
    history_blocks = [int(value) for value in anchor["history_blocks"]]
    current_block = int(anchor["block_index"])
    selected: list[tuple[int, int]] = [
        (block_index, steps_per_block) for block_index in history_blocks
    ]
    selected.append((current_block, int(anchor["tail_steps"])))
    if not selected:
        raise ValueError("An anchor must select at least its current block")

    embedding_parts: list[torch.Tensor] = []
    position_parts: list[torch.LongTensor] = []
    first_payload_base: int | None = None
    last_piece: dict[str, torch.Tensor] | None = None
    last_shift = 0

    for sequence_index, (block_index, tail_steps) in enumerate(selected):
        is_current = sequence_index == len(selected) - 1
        piece = (
            current_piece
            if is_current and current_piece is not None
            else store.variant(block_index, tail_steps)
        )
        validate_piece(piece)
        local_payload_positions = _strip_position_batch(piece["payload_position_ids"])
        local_payload_base = int(local_payload_positions.min().item())
        if first_payload_base is None:
            first_payload_base = local_payload_base
        desired_payload_base = first_payload_base + sequence_index * positions_per_chunk
        position_shift = desired_payload_base - local_payload_base

        if sequence_index == 0:
            embedding_parts.append(_strip_embedding_batch(piece["prefix_embeddings"]))
            position_parts.append(_strip_position_batch(piece["prefix_position_ids"]))
        embedding_parts.append(_strip_embedding_batch(piece["payload_embeddings"]))
        position_parts.append(local_payload_positions + position_shift)
        last_piece = piece
        last_shift = position_shift

    if last_piece is None:
        raise RuntimeError("No cached AV piece was selected")
    closure_embeddings = _strip_embedding_batch(last_piece["closure_embeddings"])
    closure_positions = (
        _strip_position_batch(last_piece["closure_position_ids"]) + last_shift
    )
    embedding_parts.append(closure_embeddings)
    position_parts.append(closure_positions)

    query_embeddings = static["query_embeddings"]
    if query_embeddings.ndim != 3 or query_embeddings.shape[0] != 1:
        raise ValueError(
            f"static query_embeddings must be [1,Q,D], got {tuple(query_embeddings.shape)}"
        )
    query_embeddings = query_embeddings[0]
    query_start = int(closure_positions.max().item()) + 1
    query_scalar = torch.arange(
        query_start,
        query_start + query_embeddings.shape[0],
        dtype=torch.long,
    )
    query_positions = query_scalar.view(1, -1).expand(3, -1)
    embedding_parts.append(query_embeddings)
    position_parts.append(query_positions)

    embeddings = torch.cat(embedding_parts, dim=0).contiguous()
    position_ids = torch.cat(position_parts, dim=1).contiguous()
    if position_ids.shape != (3, embeddings.shape[0]):
        raise RuntimeError(
            "Reconstructed position/embedding mismatch: "
            f"{tuple(position_ids.shape)} vs {tuple(embeddings.shape)}"
        )
    return embeddings, position_ids


def collate_cached_sequences(
    batch: Sequence[dict[str, Any]],
) -> dict[str, Any]:
    """Right-pad variable token sequences while preserving each readout index."""

    if not batch:
        raise ValueError("Cannot collate an empty batch")
    hidden_size = int(batch[0]["inputs_embeds"].shape[-1])
    max_length = max(int(item["inputs_embeds"].shape[0]) for item in batch)
    dtype = batch[0]["inputs_embeds"].dtype
    inputs_embeds = torch.zeros(
        len(batch), max_length, hidden_size, dtype=dtype
    )
    position_ids = torch.zeros(3, len(batch), max_length, dtype=torch.long)
    attention_mask = torch.zeros(len(batch), max_length, dtype=torch.long)
    readout_indices = torch.empty(len(batch), dtype=torch.long)
    targets = torch.stack([item["targets"] for item in batch])

    for batch_index, item in enumerate(batch):
        embeddings = item["inputs_embeds"]
        positions = item["position_ids"]
        length = int(embeddings.shape[0])
        if embeddings.shape[-1] != hidden_size:
            raise ValueError("Mixed hidden sizes in one batch")
        if positions.shape != (3, length):
            raise ValueError(
                f"Invalid position_ids {tuple(positions.shape)} for length {length}"
            )
        inputs_embeds[batch_index, :length] = embeddings
        position_ids[:, batch_index, :length] = positions
        attention_mask[batch_index, :length] = 1
        readout_indices[batch_index] = length - 1

    return {
        "inputs_embeds": inputs_embeds,
        "position_ids": position_ids,
        "attention_mask": attention_mask,
        "readout_indices": readout_indices,
        "targets": targets,
        "step_indices": torch.tensor(
            [int(item["step_index"]) for item in batch], dtype=torch.long
        ),
    }


def anchor_intervals(
    anchor: dict[str, Any],
    *,
    block_seconds: float,
    step_seconds: float,
) -> list[tuple[float, float]]:
    """Return source intervals used by the reference KV replay for an anchor."""

    intervals = [
        (index * block_seconds, (index + 1) * block_seconds)
        for index in anchor["history_blocks"]
    ]
    block_index = int(anchor["block_index"])
    current_start = block_index * block_seconds
    current_end = current_start + int(anchor["tail_steps"]) * step_seconds
    intervals.append((current_start, current_end))
    return intervals


def forward_last_state(
    language_model: torch.nn.Module,
    *,
    inputs_embeds: torch.Tensor,
    position_ids: torch.LongTensor,
    attention_mask: torch.LongTensor,
    readout_indices: torch.LongTensor,
) -> torch.Tensor:
    """Run full-sequence causal Thinker attention and gather each query state."""

    outputs = language_model(
        inputs_embeds=inputs_embeds,
        position_ids=position_ids,
        attention_mask=attention_mask,
        use_cache=False,
        output_hidden_states=False,
        return_dict=True,
        cache_position=torch.arange(
            inputs_embeds.shape[1], device=inputs_embeds.device
        ),
    )
    batch_indices = torch.arange(
        inputs_embeds.shape[0], device=inputs_embeds.device
    )
    return outputs.last_hidden_state[batch_indices, readout_indices]


def cache_size_bytes(paths: Iterable[Path]) -> int:
    return sum(path.stat().st_size for path in paths if path.is_file())