"""Turn one PLD speaker into a VITS training set, cached on the Hub.

    python scripts/vits_data.py --language ceb --survey
    python scripts/vits_data.py --language ceb --speaker-id top --push-to Splintir/pld-ceb-vits

Why this exists separately from `tts_data.py`: that script computes log-mel
spectrograms and x-vectors, because SpeechT5 consumes both. VITS consumes raw
waveform and text and nothing else -- it learns its own alignment and carries
one baked-in voice, so there is no speaker embedding to compute.

**Single speaker, deliberately.** MMS/VITS checkpoints hold exactly one voice.
Finetuning a one-voice model on PLD's many speakers averages them into mush,
which is the most likely reason the single-speaker SpeechT5 `-solo` run beat the
full-corpus `-v2` run on every statistic. Applying that lesson before the run
this time rather than after it.

`--survey` prints the speaker distribution and exits, so the choice of speaker
is made against clip counts and total duration rather than assumed. A VITS
finetune wants tens of minutes at minimum; if the top speaker is thin, the
survey says so before any GPU time is spent.

The scan reuses `tts_data.py`'s shard iterator: PLD's train split is 201 shards
of every language interleaved (~24 GB), so one language means touching all of
them, one shard on disk at a time.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from tts_data import PLD_REPO, iter_shard_rows  # noqa: E402

DIGIT_RE = re.compile(r"\d")

# MMS tokenizers are character-level and lowercase, with a per-language vocab
# that excludes digits. Rejecting a row is honest; keeping it would teach the
# model that "1990" is silence.
KEEP_RE = re.compile(r"[^a-zñáéíóú' -]")


def clean_text(text: str) -> str | None:
    text = text.replace("’", "'").replace("‘", "'")
    text = text.lower().strip()
    if not text or DIGIT_RE.search(text):
        return None
    text = KEEP_RE.sub("", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text or None


def survey(language: str, token: str | None, shards: int) -> None:
    """Print who speaks this language and for how long, then stop."""
    clips: Counter[str] = Counter()
    secs: defaultdict[str, float] = defaultdict(float)

    # gender/age are not in tts_data.COLUMNS and the iterator only reads those,
    # so the survey reports clips and duration -- the two numbers that decide
    # whether a speaker can carry a finetune.
    for n, (_, row) in enumerate(iter_shard_rows(language, token, range(shards)), 1):
        sid = row["speaker_id"]
        clips[sid] += 1
        secs[sid] += float(row.get("duration") or 0.0)
        if n % 2000 == 0:
            print(f"  ... {n} rows, {len(clips)} speakers", flush=True)

    total = sum(clips.values())
    print(f"\n{language}: {total} usable clips, {len(clips)} speakers, "
          f"{sum(secs.values()) / 3600:.1f} h total\n")
    print(f"{'speaker_id':<26}{'clips':>8}{'minutes':>10}")
    for sid, n in clips.most_common(20):
        print(f"{sid:<26}{n:>8}{secs[sid] / 60:>10.1f}")
    print("\nA VITS finetune wants >= ~30 min from one speaker. Pick from the "
          "top rows and rerun with --speaker-id.")


def cache_dir(language: str, speaker: str) -> Path:
    return Path("vits_cache") / language / speaker


def load_cached(language: str, speaker: str):
    """Reuse a completed scan. Scanning 201 shards to find ~15 minutes of audio
    costs ~24 GB of transfer, so it must never be repeated because a later step
    failed."""
    manifest = cache_dir(language, speaker) / "manifest.jsonl"
    if not manifest.exists():
        return None
    rows = [json.loads(line) for line in
            manifest.read_text(encoding="utf-8").splitlines() if line.strip()]
    rows = [r for r in rows if (cache_dir(language, speaker) / r["file"]).exists()]
    if not rows:
        return None
    print(f"reusing {len(rows)} cached clips from "
          f"{cache_dir(language, speaker)} (delete it to force a rescan)",
          flush=True)
    return rows


def collect(language: str, speaker: str, token: str | None, shards: int,
            max_seconds: float):
    """Gather one speaker's clips. Returns (records, resolved_speaker_id)."""
    import io

    import soundfile as sf

    # `top` cannot be resolved until the corpus has been scanned once, so the
    # first pass counts and the second keeps. Two passes over 24 GB is slow;
    # buffering every language's audio in RAM instead is worse.
    if speaker == "top":
        counts: Counter[str] = Counter()
        for _, row in iter_shard_rows(language, token, range(shards)):
            counts[row["speaker_id"]] += 1
        if not counts:
            raise SystemExit(f"no usable {language} rows -- check the filters")
        speaker, n = counts.most_common(1)[0]
        print(f"resolved `top` -> {speaker} ({n} clips)", flush=True)

    out = cache_dir(language, speaker)
    out.mkdir(parents=True, exist_ok=True)
    manifest = (out / "manifest.jsonl").open("w", encoding="utf-8")

    records, total = [], 0.0
    seen_shard = -1
    for shard, row in iter_shard_rows(language, token, range(shards)):
        # One speaker is a handful of clips scattered over 201 shards, so a
        # clip-count progress line can stay silent for hours. Report the scan
        # itself instead -- otherwise a live run is indistinguishable from a
        # hung one.
        if shard != seen_shard:
            seen_shard = shard
            print(f"  shard {shard + 1}/{shards}  kept {len(records)} clips, "
                  f"{total / 60:.1f} min", flush=True)
        if row["speaker_id"] != speaker:
            continue
        text = clean_text(row["sentence"])
        if not text:
            continue
        raw = row["audio"]["bytes"]
        # Decode once here rather than trusting the shard's declared duration:
        # the trainer segments on real sample counts, and a mismatch shows up as
        # a silent crash deep in the collator.
        try:
            wav, rate = sf.read(io.BytesIO(raw), dtype="float32", always_2d=False)
        except Exception as exc:                      # noqa: BLE001
            print(f"  skipped unreadable clip: {exc}", flush=True)
            continue
        if wav.ndim > 1:
            wav = wav.mean(axis=1)
        secs = len(wav) / rate
        if not 1.0 <= secs <= 15.0:
            continue

        # Write PLD's own encoded bytes straight through rather than re-encoding
        # the decoded array: no quality loss, and `datasets` can build an Audio
        # column from encoded bytes without torchcodec, which it needs for raw
        # arrays and file paths alike.
        ext = Path(row["audio"].get("path") or "clip.wav").suffix or ".wav"
        name = f"{len(records):04d}{ext}"
        (out / name).write_bytes(raw)
        manifest.write(json.dumps({"file": name, "text": text,
                                   "seconds": round(secs, 3)}) + "\n")
        manifest.flush()

        records.append({"file": name, "text": text, "seconds": round(secs, 3)})
        total += secs
        if max_seconds and total >= max_seconds:
            break

    manifest.close()
    print(f"\n{speaker}: {len(records)} clips, {total / 60:.1f} min "
          f"-> cached in {out}", flush=True)
    if total < 900:
        print("WARNING: under 15 minutes. Expect a weak finetune -- consider "
              "pooling a second speaker of the same gender and dialect.",
              flush=True)
    return records, speaker


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--language", default="ceb", help="PLD ISO 639-3 code")
    ap.add_argument("--speaker-id", default="top",
                    help="`top` picks the speaker with the most clips")
    ap.add_argument("--survey", action="store_true",
                    help="print the speaker distribution and exit")
    ap.add_argument("--shards", type=int, default=201)
    ap.add_argument("--max-seconds", type=float, default=0,
                    help="stop after this much audio (0 = no cap)")
    ap.add_argument("--push-to", default="",
                    help="Hub dataset repo, e.g. Splintir/pld-ceb-vits")
    args = ap.parse_args()

    token = os.environ.get("HF_TOKEN")
    if not token:
        env = Path(__file__).resolve().parent.parent / ".env"
        if env.exists():
            for line in env.read_text(encoding="utf-8").splitlines():
                key, _, value = line.strip().partition("=")
                if key == "HF_TOKEN" and value:
                    token = value.strip()

    print(f"scanning {args.shards} {PLD_REPO} train shards for {args.language} ...",
          flush=True)

    if args.survey:
        survey(args.language, token, args.shards)
        return

    speaker = args.speaker_id
    records = None if speaker == "top" else load_cached(args.language, speaker)
    if records is None:
        records, speaker = collect(args.language, args.speaker_id, token,
                                   args.shards, args.max_seconds)
    if not records:
        raise SystemExit(f"no clips for speaker {speaker}")

    from datasets import Audio, Dataset

    src = cache_dir(args.language, speaker)
    rows = [{"audio": {"bytes": (src / r["file"]).read_bytes(),
                       "path": r["file"]},
             "text": r["text"]}
            for r in records]
    ds = Dataset.from_list(rows).cast_column("audio", Audio(sampling_rate=16000))
    # A held-out slice the trainer can score against, kept small: VITS eval is
    # slow (it renders audio) and the number that decides anything is the
    # 50-line bench, not this. Proportional with a floor and a ceiling -- a
    # flat floor alone puts more clips in eval than train on a small set.
    n_eval = max(4, min(16, len(ds) // 10))
    ds = ds.train_test_split(test_size=n_eval, seed=0)
    print(ds)

    if args.push_to:
        ds.push_to_hub(args.push_to, token=token, private=True)
        print(f"pushed -> {args.push_to}  (speaker {speaker})")
    else:
        out = Path("vits_data") / args.language
        ds.save_to_disk(str(out))
        print(f"saved -> {out}  (pass --push-to to upload)")


if __name__ == "__main__":
    main()