#!/usr/bin/env python3
"""Validate every shard of a BookTTS teacher dataset."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path

import numpy as np


def audit(args: argparse.Namespace) -> None:
    metadata_path = args.dataset / "dataset.json"
    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    token_map = np.load(args.dataset / "token_id_map.npy")
    samples = 0
    tokens = 0
    bytes_total = metadata_path.stat().st_size + (args.dataset / "token_id_map.npy").stat().st_size
    digest = hashlib.sha256()
    for record in metadata.get("shards", []):
        path = args.dataset / record["file"]
        digest.update(path.read_bytes())
        bytes_total += path.stat().st_size
        with np.load(path) as shard:
            ids = shard["original_input_ids"]
            mask = shard["attention_mask"]
            teacher = shard["teacher"]
            if ids.shape != mask.shape or teacher.shape[:2] != ids.shape or teacher.shape[2:] != (768,):
                raise ValueError(f"Incompatible shapes in {path.name}: {ids.shape}, {mask.shape}, {teacher.shape}")
            if ids.shape[0] != record["samples"]:
                raise ValueError(f"Sample count mismatch in {path.name}")
            if ids.size and (ids.min() < 0 or ids.max() >= token_map.shape[0]):
                raise ValueError(f"Token id outside token map in {path.name}")
            if not np.isfinite(teacher).all():
                raise ValueError(f"Non-finite teacher value in {path.name}")
            if not np.all((mask == 0) | (mask == 1)):
                raise ValueError(f"Non-binary attention mask in {path.name}")
            samples += ids.shape[0]
            tokens += int(mask.sum())
    if samples != metadata["samples"]:
        raise ValueError(f"Dataset sample count is {samples}, metadata says {metadata['samples']}")
    report = {
        "dataset": str(args.dataset),
        "samples": samples,
        "shards": len(metadata.get("shards", [])),
        "tokens": tokens,
        "vocab_size": metadata["vocab_size"],
        "token_map_entries": int(token_map.shape[0]),
        "bytes": bytes_total,
        "shards_sha256": digest.hexdigest().upper(),
        "teacher_dtype": metadata.get("teacher_dtype"),
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=Path, required=True)
    return parser.parse_args()


if __name__ == "__main__":
    sys.exit(audit(parse_args()) or 0)