"""
Dataset inspection script for Phase 1.

This script is written and ready to run, but has NOT been executed against
real data yet because no dataset is present under data/raw/ as of this
commit. It performs no assumptions about label schema beyond what it can
detect from the actual files it finds — if it can't determine something
(e.g. speaker ID), it reports that explicitly rather than guessing.

Usage (once data/raw/ is populated):
    python scripts/inspect_dataset.py --data-dir data/raw --out docs/dataset_report.json

What it checks:
  - file inventory (audio files, manifest/label files found)
  - per-file audio properties: sample rate, channels, duration, corruption
  - label schema detection (from any CSV/JSON manifest found)
  - class balance (if labels found)
  - duplicate detection (exact file hash + near-duplicate via audio hash)
  - speaker/session ID presence (from manifest columns, if any)
  - transcript / language metadata presence
"""

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

AUDIO_EXTENSIONS = {".wav", ".flac", ".mp3", ".ogg", ".m4a"}
MANIFEST_EXTENSIONS = {".csv", ".json", ".jsonl", ".tsv"}


def find_files(data_dir: Path):
    audio_files, manifest_files, other_files = [], [], []
    for p in data_dir.rglob("*"):
        if not p.is_file():
            continue
        suffix = p.suffix.lower()
        if suffix in AUDIO_EXTENSIONS:
            audio_files.append(p)
        elif suffix in MANIFEST_EXTENSIONS:
            manifest_files.append(p)
        else:
            other_files.append(p)
    return audio_files, manifest_files, other_files


def file_hash(path: Path, chunk_size: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(chunk_size):
            h.update(chunk)
    return h.hexdigest()


def inspect_audio_file(path: Path):
    """Return dict of measured properties, or an error record. No guessing."""
    try:
        import soundfile as sf
    except ImportError:
        return {"path": str(path), "error": "soundfile not installed"}

    try:
        info = sf.info(str(path))
        return {
            "path": str(path),
            "sample_rate": info.samplerate,
            "channels": info.channels,
            "duration_sec": round(info.frames / info.samplerate, 4)
            if info.samplerate
            else None,
            "frames": info.frames,
            "format": info.format,
            "corrupted": False,
        }
    except Exception as e:
        return {"path": str(path), "corrupted": True, "error": str(e)}


def try_load_manifests(manifest_files):
    """Load any CSV/JSON manifests found and report their columns.
    Does not assume which column is the label — reports raw schema so a
    human (or a follow-up script, once we know the real schema) can decide.
    """
    reports = []
    for mf in manifest_files:
        entry = {"path": str(mf)}
        try:
            if mf.suffix.lower() == ".csv" or mf.suffix.lower() == ".tsv":
                import csv

                delim = "\t" if mf.suffix.lower() == ".tsv" else ","
                with open(mf, newline="", encoding="utf-8", errors="replace") as f:
                    reader = csv.reader(f, delimiter=delim)
                    rows = list(reader)
                entry["n_rows"] = max(0, len(rows) - 1)
                entry["columns"] = rows[0] if rows else []
                entry["sample_rows"] = rows[1:4]
            elif mf.suffix.lower() in (".json", ".jsonl"):
                with open(mf, encoding="utf-8", errors="replace") as f:
                    if mf.suffix.lower() == ".jsonl":
                        lines = [json.loads(l) for l in f.readlines()[:5] if l.strip()]
                        entry["sample_records"] = lines
                    else:
                        data = json.load(f)
                        entry["top_level_type"] = type(data).__name__
                        if isinstance(data, list) and data:
                            entry["columns_guess"] = (
                                list(data[0].keys())
                                if isinstance(data[0], dict)
                                else None
                            )
                            entry["n_records"] = len(data)
                            entry["sample_records"] = data[:3]
                        elif isinstance(data, dict):
                            entry["top_level_keys"] = list(data.keys())
        except Exception as e:
            entry["error"] = str(e)
        reports.append(entry)
    return reports


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", type=Path, default=Path("data/raw"))
    ap.add_argument("--out", type=Path, default=Path("docs/dataset_report.json"))
    ap.add_argument(
        "--max-audio-inspect",
        type=int,
        default=None,
        help="Cap number of audio files to probe for sample rate/duration "
        "(useful for a quick pass on very large datasets). Default: all.",
    )
    args = ap.parse_args()

    if not args.data_dir.exists():
        print(
            f"ERROR: {args.data_dir} does not exist. Nothing to inspect. "
            f"Place the raw dataset there first.",
            file=sys.stderr,
        )
        sys.exit(1)

    audio_files, manifest_files, other_files = find_files(args.data_dir)

    report = {
        "data_dir": str(args.data_dir),
        "n_audio_files_found": len(audio_files),
        "n_manifest_files_found": len(manifest_files),
        "n_other_files_found": len(other_files),
        "manifest_files": [str(p) for p in manifest_files],
        "other_files_sample": [str(p) for p in other_files[:20]],
    }

    if not audio_files and not manifest_files:
        print(
            f"WARNING: no audio or manifest files found under {args.data_dir}. "
            f"Found {len(other_files)} other files. Nothing measured.",
        )

    # Manifest schema (no label assumptions)
    report["manifests"] = try_load_manifests(manifest_files)

    # Audio properties + corruption + duplicate check
    to_inspect = audio_files[: args.max_audio_inspect] if args.max_audio_inspect else audio_files
    audio_reports = [inspect_audio_file(p) for p in to_inspect]
    report["audio_files"] = audio_reports

    corrupted = [r for r in audio_reports if r.get("corrupted")]
    report["n_corrupted"] = len(corrupted)

    hashes = {}
    for p in to_inspect:
        try:
            h = file_hash(p)
            hashes.setdefault(h, []).append(str(p))
        except Exception:
            pass
    duplicates = {h: paths for h, paths in hashes.items() if len(paths) > 1}
    report["n_exact_duplicate_groups"] = len(duplicates)
    report["exact_duplicates_sample"] = dict(list(duplicates.items())[:5])

    # Sample rate / channel / duration distribution summary (only over successfully read files)
    ok = [r for r in audio_reports if not r.get("corrupted")]
    if ok:
        srs = sorted(set(r["sample_rate"] for r in ok if r.get("sample_rate") is not None))
        chans = sorted(set(r["channels"] for r in ok if r.get("channels") is not None))
        durations = [r["duration_sec"] for r in ok if r.get("duration_sec") is not None]
        report["sample_rates_found"] = srs
        report["channel_counts_found"] = chans
        if durations:
            durations_sorted = sorted(durations)
            n = len(durations_sorted)
            report["duration_stats_sec"] = {
                "n": n,
                "min": durations_sorted[0],
                "max": durations_sorted[-1],
                "mean": round(sum(durations_sorted) / n, 4),
                "median": durations_sorted[n // 2],
            }

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(report, f, indent=2)

    print(f"Wrote dataset report to {args.out}")
    print(json.dumps({k: v for k, v in report.items() if k != "audio_files"}, indent=2))


if __name__ == "__main__":
    main()