"""Mmap-backed teacher bank for DINOv3 + V-JEPA2 + VideoMAE-v2 features."""

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
from typing import Any, Mapping, Sequence

import torch
from safetensors import safe_open
from torch.utils.data import Dataset


BANK_KIND = "triple_dinov3_vjepa2_videomae_v2_matched_teacher_bank"
BANK_SCHEMA_VERSION = 1
BRANCH_SHAPES = {
    "dinov3": (1024, 1024),
    "vjepa2": (1024, 1024),
    "videomae_v2": (1280, 1024),
}


def _canonical_bytes(value: Any) -> bytes:
    return json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
        allow_nan=False,
    ).encode("utf-8")


def canonical_sha256(value: Any) -> str:
    return hashlib.sha256(_canonical_bytes(value)).hexdigest()


def file_sha256(path: str | os.PathLike) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        while True:
            block = handle.read(8 * 1024 * 1024)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def _strict_sha256(value: object, label: str) -> str:
    normalized = str(value)
    return normalized


class TripleVideoFeatureBank(Dataset):
    """Metadata dataset plus random-access frozen FP16 feature rows."""

    def __init__(
        self,
        manifest_path: str | os.PathLike,
        *,
        expected_manifest_sha256: str,
        expected_content_hash: str,
        expected_source_manifest_sha256: str,
        expected_source_csv_sha256: str | None = None,
        expected_sample_identity_sha256: str | None = None,
        expected_sample_count: int = 5000,
        expected_videomae_revision: str | None = None,
        expected_videomae_download_manifest_content_hash: str | None = None,
        expected_vjepa_model_sha256: str | None = None,
        expected_dinov3_checkpoint_sha256: str | None = None,
        expected_dinov3_source_commit: str | None = None,
        expected_shard_size: int = 32,
        validate_all_shards: bool = True,
    ) -> None:
        self.manifest_path = Path(manifest_path).expanduser().resolve()
        actual_manifest_sha = file_sha256(self.manifest_path)
        with self.manifest_path.open("r", encoding="utf-8") as handle:
            manifest = json.load(handle)
        content_hash = manifest["content_hash"]

        samples = manifest["samples"]
        self.samples = []
        self._sample_index = {}
        for expected_index, item in enumerate(samples):
            self._sample_index[item["sample_id"]] = expected_index
            self.samples.append(dict(item))

        self.root = self.manifest_path.parent
        self.shards = []
        self._row_to_shard: list[tuple[int, int]] = [None] * len(samples)
        for shard_index, item in enumerate(manifest["shards"]):
            start = int(item["sample_start"])
            end = int(item["sample_end"])
            relative = Path(str(item["path"]))
            if relative.is_absolute() or ".." in relative.parts:
                raise ValueError("triple bank shard path must stay under bank root")
            path = (self.root / relative).resolve()
            if self.root not in path.parents or not path.is_file():
                raise FileNotFoundError(path)
            for row in range(start, end):
                self._row_to_shard[row] = (shard_index, row - start)
            self.shards.append({**item, "resolved_path": str(path)})
        self.manifest = manifest
        self.content_hash = content_hash
        self.manifest_sha256 = actual_manifest_sha

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, index: int) -> dict[str, object]:
        item = self.samples[int(index)]
        return {
            "prompts": item["prompt"],
            "sample_ids": item["sample_id"],
            "teacher_bank_rows": int(item["index"]),
        }

    @staticmethod
    def collate_fn(items: Sequence[Mapping[str, object]]) -> dict[str, object]:
        return {
            "prompts": [str(item["prompts"]) for item in items],
            "sample_ids": [str(item["sample_ids"]) for item in items],
            "teacher_bank_rows": torch.tensor([
                int(item["teacher_bank_rows"]) for item in items
            ], dtype=torch.long),
        }

    def load_rows(
        self,
        sample_ids: Sequence[str],
    ) -> dict[str, torch.Tensor]:
        """Read only requested rows, preserving order and duplicate IDs."""
        ids = tuple(str(value) for value in sample_ids)
        output = {name: [None] * len(ids) for name in BRANCH_SHAPES}
        by_shard: dict[int, list[tuple[int, int]]] = {}
        for output_index, sample_id in enumerate(ids):
            row = self._sample_index[sample_id]
            shard_index, offset = self._row_to_shard[row]
            by_shard.setdefault(shard_index, []).append((output_index, offset))
        for shard_index, requests in by_shard.items():
            path = self.shards[shard_index]["resolved_path"]
            with safe_open(path, framework="pt", device="cpu") as handle:
                slices = {name: handle.get_slice(name) for name in BRANCH_SHAPES}
                for output_index, offset in requests:
                    for name in BRANCH_SHAPES:
                        output[name][output_index] = slices[name][
                            offset:offset + 1
                        ].squeeze(0)
        return {
            name: torch.stack(rows, dim=0)
            for name, rows in output.items()
        }


__all__ = [
    "BANK_KIND",
    "BANK_SCHEMA_VERSION",
    "BRANCH_SHAPES",
    "TripleVideoFeatureBank",
    "canonical_sha256",
    "file_sha256",
]