"""Shared AV preparation, independent strict-window Thinker replay banks.

Only encoder-native block embeddings are shared. High-layer KV is never
cropped from a longer context to stand in for a shorter-context extraction.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Sequence

import torch
from tqdm.auto import tqdm

import streaming_turn_model as core

DEFAULT_MEMORY_SECONDS = (2.0, 4.0, 6.0, 8.0, 10.0, 12.0)


def validate_memory_seconds(values: Sequence[float], native_seconds: float) -> tuple[float, ...]:
    """Reject ambiguous/non-native budgets before loading media or a model."""
    if not math.isfinite(native_seconds) or native_seconds <= 0:
        raise ValueError("native block seconds must be finite and positive")
    caps = tuple(float(value) for value in values)
    if not caps or len(set(caps)) != len(caps):
        raise ValueError("memory seconds list must be nonempty with no duplicates")
    for cap in caps:
        if not math.isfinite(cap) or cap < native_seconds:
            raise ValueError("every memory cap must be finite and at least one native block")
        blocks = core._require_integer_ratio(cap, native_seconds, "memory / native block")
        if cap != blocks * native_seconds:
            raise ValueError("memory caps must be exact native-block multiples, not rounded budgets")
    return tuple(sorted(caps))


def context_directory(cap: float) -> str:
    number = str(int(cap)) if float(cap).is_integer() else repr(float(cap))
    return f"context_{number}s"


def output_paths(output: Path, caps: Sequence[float]) -> dict[float, Path]:
    return {cap: output.parent / context_directory(cap) / output.name for cap in caps}


@dataclass
class _Bank:
    helper: core.OpenStreamReplayExtractor
    history_start: float
    committed_cache: Any = None
    payload_base: int | None = None
    states: list[torch.Tensor] = field(default_factory=list)
    checks: list[dict[str, Any]] = field(default_factory=list)
    history_starts: list[float] = field(default_factory=list)
    rebuilds: int = 0
    replayed_blocks: int = 0


class MultiWindowOpenStreamReplayExtractor(core.OpenStreamReplayExtractor):
    """Same tail-aware semantics as separate single-window extractions.

Each bank has an independent KV history. All banks consume the same prepared
current tail and a bounded shared ring of completed AV block embeddings.
Temporary tails, chat closure and readout queries are not committed until a
native boundary; only the completed media payload is committed there.
"""

    def __init__(self, model, processor, config, *, memory_seconds_list, fixed_grid_av=None):
        self.memory_seconds_list = validate_memory_seconds(
            memory_seconds_list, config.native_block_seconds
        )
        super().__init__(
            model, processor, replace(config, memory_seconds=max(self.memory_seconds_list)),
            fixed_grid_av=fixed_grid_av,
        )

    @torch.inference_mode()
    def extract(
        self, *, video_path: str | Path | None, audio_path: str | Path | None,
        start: float, end: float, use_audio_in_video: bool,
        verify_every: int = 0, verify_continuous_every: int = 0,
        show_progress: bool = True,
    ) -> dict[float, core.ExtractionResult]:
        if verify_every < 0 or verify_continuous_every < 0:
            raise ValueError("verification intervals must be non-negative")
        if not math.isclose(start, 0.0, abs_tol=1e-9) or not math.isclose(
            self.config.step_seconds, 0.2, abs_tol=1e-9
        ):
            raise ValueError("multiwindow dense banks require start=0 and step_seconds=0.2")
        if not math.isfinite(end) or end <= start:
            raise ValueError("end must be finite and greater than start")
        core._require_integer_ratio(end - start, self.config.step_seconds, "duration / step")
        times = core._time_grid(start, end, self.config.step_seconds)
        native = self.config.native_block_seconds
        positions_per_chunk = core._tmrope_positions_per_chunk(self.model)
        banks = {
            cap: _Bank(
                helper=core.OpenStreamReplayExtractor(
                    self.model, self.processor, replace(self.config, memory_seconds=cap),
                    fixed_grid_av=self.fixed_grid_av,
                ),
                history_start=start,
            )
            for cap in self.memory_seconds_list
        }
        completed = core.CompletedBlockPieceCache(
            origin=start, block_seconds=native,
            max_blocks=round(max(self.memory_seconds_list) / native) - 1,
        )
        block_start = start
        prepared_seconds = 0.0
        with tqdm(total=len(times), desc="multiwindow-replay", unit="cutoff",
                  disable=not show_progress, dynamic_ncols=True) as progress:
            for step, cutoff in enumerate(times, start=1):
                for cap, bank in banks.items():
                    if bank.committed_cache is not None and block_start - bank.history_start >= cap - 1e-6:
                        bank.history_start = block_start - (cap - native)
                        intervals = core._block_intervals(bank.history_start, block_start, native)
                        bank.rebuilds += 1
                        bank.replayed_blocks += len(intervals)
                        if intervals:
                            bank.committed_cache, _, bank.payload_base, _, _ = bank.helper._cache_from_intervals(
                                video_path=video_path, audio_path=audio_path, intervals=intervals,
                                use_audio_in_video=use_audio_in_video,
                                prepared_blocks=completed.get_many(intervals),
                            )
                        else:
                            bank.committed_cache = None
                            bank.payload_base = None

                # Exactly one production processor/tower invocation per cutoff.
                pieces = self._prepare_pieces(
                    video_path=video_path, audio_path=audio_path,
                    interval=(block_start, cutoff), use_audio_in_video=use_audio_in_video,
                )
                prepared_seconds += cutoff - block_start
                boundary = math.isclose(cutoff, block_start + native, abs_tol=1e-6)
                for cap, bank in banks.items():
                    if bank.payload_base is None:
                        bank.payload_base = int(pieces.payload_position_ids.min().item())
                    block_index = round((block_start - bank.history_start) / native)
                    media_cache, _, closure_positions = core._append_open_pieces(
                        self.model, pieces, past_key_values=core.clone_cache(bank.committed_cache),
                        rope_offset=bank.payload_base + block_index * positions_per_chunk,
                        include_prefix=bank.committed_cache is None,
                    )
                    state = core._read_open_state(
                        self.model, media_cache, closure_ids=pieces.closure_ids,
                        closure_position_ids=closure_positions, query_ids=bank.helper.query_ids,
                        last_n_layers=self.config.readout_last_n_layers,
                        layer_fusion=self.config.readout_layer_fusion,
                    )
                    bank.states.append(state.cpu())
                    bank.history_starts.append(bank.history_start)
                    intervals = core._block_intervals(bank.history_start, cutoff, native)
                    if verify_every and step % verify_every == 0:
                        # Deliberately independent processor/tower calls, not the shared ring.
                        reference = bank.helper.state_from_scratch(
                            video_path=video_path, audio_path=audio_path, intervals=intervals,
                            use_audio_in_video=use_audio_in_video,
                        )
                        bank.checks.append(core._compare_states(cutoff, state, reference, kind="cache_replay"))
                    if verify_continuous_every and step % verify_continuous_every == 0:
                        reference = bank.helper.continuous_reference.state_for_intervals(
                            video_path=video_path, audio_path=audio_path,
                            intervals=[(bank.history_start, cutoff)], use_audio_in_video=use_audio_in_video,
                        )
                        bank.checks.append(core._compare_states(
                            cutoff, state, reference, kind="continuous_window_semantics"
                        ))
                    if boundary:
                        bank.committed_cache = media_cache
                if boundary:
                    completed.put(core.CachedOpenMediaBlock(block_start, cutoff, pieces))
                    block_start = cutoff
                progress.set_postfix_str(f"media={cutoff:.1f}s banks={len(banks)}", refresh=False)
                progress.update(1)

        time_tensor = torch.tensor(times, dtype=torch.float64)
        boundary_mask = core._periodic_boundary_mask(times, origin=start, period=native)
        results = {}
        for cap, bank in banks.items():
            history = torch.tensor(bank.history_starts, dtype=torch.float64)
            span = time_tensor - history
            if (span <= 0).any() or (span > cap + 1e-6).any():
                raise RuntimeError(f"invalid context span in {cap:g}s bank")
            results[cap] = core.ExtractionResult(
                times=time_tensor.clone(), states=torch.cat(bank.states, dim=0),
                verification=bank.checks,
                diagnostics={
                    "history_start_times": history,
                    "actual_context_seconds": span,
                    "state_source_times": time_tensor.clone(),
                    "state_age_seconds": torch.zeros_like(time_tensor),
                    "initialization_mask": torch.zeros(len(times), dtype=torch.bool),
                    "state_update_mask": torch.ones(len(times), dtype=torch.bool),
                    "payload_commit_mask": boundary_mask.clone(),
                    "native_boundary_mask": boundary_mask.clone(),
                    "production_processor_calls_shared": len(times),
                    "production_prepared_media_seconds_shared": prepared_seconds,
                    "production_thinker_tail_calls": len(times),
                    "rolling_rebuilds": bank.rebuilds,
                    "replayed_complete_blocks": bank.replayed_blocks,
                    "shared_completed_block_capacity": round(max(self.memory_seconds_list) / native) - 1,
                    "memory_seconds_list": list(self.memory_seconds_list),
                    "sharing": "AV_embeddings_only_independent_Thinker_KV",
                    "verification_excluded_from_production_counters": True,
                },
            )
        return results