"""Turn one PLD language into SpeechT5 training inputs, cached on the Hub.

Why this exists: `halohalo/finetune_tts.py` reads PLD from a raw corpus on a
local disk (`PLD_RAW`), which a Colab VM does not have. The Hub copy
(`sapinsapin/pld`) is 201 train shards of every language interleaved, so
pulling one language means touching all ~24 GB.

That is a one-time cost, and this script pays it once: stream the shards one at
a time (never more than one on disk), keep the rows for one language, run the
same text/mel/x-vector preprocessing the original training used, and write the
result to parquet. Push that to a Hub dataset repo and every later run --
including every resume after a free-tier VM disappears -- pulls ~1 GB instead
of 24.

Preprocessing is deliberately identical to `finetune_tts.py`:
  - text: curly apostrophes normalized, digit-bearing lines dropped
  - audio: SpeechT5Processor log-mel targets at 16 kHz
  - speaker: one speechbrain x-vector per *clip*, F.normalize'd, never averaged
  - caps: <=220 input ids, <=960 mel frames

Usage (on the VM):
    python scripts/tts_data.py --language ceb --push-to Splintir/pld-ceb-tts-proc
"""

from __future__ import annotations

import argparse
import io
import os
import re
import shutil
import sys
import time
from pathlib import Path

import numpy as np

SR = 16000
PLD_REPO = "sapinsapin/pld"
N_TRAIN_SHARDS = 201
DIGIT_RE = re.compile(r"\d")

# Columns read from each shard. `audio` dominates the transfer; the rest are
# cheap and decide whether a row is kept at all.
COLUMNS = [
    "audio",
    "sentence",
    "language",
    "speech_type",
    "num_words",
    "text_is_prompt",
    "speaker_id",
    "duration",
]

MAX_INPUT_IDS = 220
MAX_MEL_FRAMES = 960


def clean_text(text: str) -> str | None:
    """SpeechT5's tokenizer is character-level Latin: normalize apostrophes and
    reject digits rather than teach the model to skip them."""
    text = text.replace("’", "'").replace("‘", "'").strip()
    if not text or DIGIT_RE.search(text):
        return None
    return text


def keep_row(row: dict, language: str) -> bool:
    """The TTS filter from halolib.finetune._PLD_FILTERS, plus the language."""
    return (
        row["language"] == language
        and row["speech_type"] == "read"
        and not row["text_is_prompt"]
        and row["num_words"] >= 3
    )


def build_embedder():
    import torch

    # speechbrain 1.1 registers optional integrations (k2, wordemb, ...) as lazy
    # modules that import on any attribute access. Loading the Xvector lobe goes
    # through pydoc, which probes dunders on every module in sys.modules and so
    # force-imports integrations whose dependencies are absent. Dunders never
    # come from a lazy import, so refusing them is safe. Same patch as
    # sapin-hil/tts.py.
    from speechbrain.utils import importutils as _importutils

    _lazy_getattr = _importutils.LazyModule.__getattr__
    _importutils.LazyModule.__getattr__ = (
        lambda self, attr: (_ for _ in ()).throw(AttributeError(attr))
        if attr.startswith("__") else _lazy_getattr(self, attr))

    from speechbrain.inference.speaker import EncoderClassifier
    from speechbrain.utils.fetching import LocalStrategy

    savedir = Path(os.environ.get("HF_HOME", "~/.cache")).expanduser() / "speechbrain-xvect"
    return EncoderClassifier.from_hparams(
        source="speechbrain/spkrec-xvect-voxceleb",
        savedir=str(savedir),
        # COPY, not the default SYMLINK: symlinking needs Developer Mode on
        # Windows and fails with WinError 1314 otherwise. Harmless on Linux.
        local_strategy=LocalStrategy.COPY,
        run_opts={"device": "cuda" if torch.cuda.is_available() else "cpu"},
    )


def iter_shard_rows(language: str, token: str | None, shards: range):
    """Yield matching rows shard by shard, holding one shard on disk at a time."""
    import pyarrow.parquet as pq
    from huggingface_hub import hf_hub_download

    scratch = Path("/content/_pld_shard") if Path("/content").exists() else Path("./_pld_shard")

    for i in shards:
        name = f"data/train-{i:05d}-of-{N_TRAIN_SHARDS:05d}.parquet"
        shutil.rmtree(scratch, ignore_errors=True)
        scratch.mkdir(parents=True, exist_ok=True)
        path = hf_hub_download(
            PLD_REPO, name, repo_type="dataset", token=token, local_dir=str(scratch)
        )
        table = pq.read_table(path, columns=COLUMNS)
        # to_pylist on the whole shard materializes 1500 audio blobs (~120 MB);
        # row-group at a time keeps the peak an order of magnitude lower.
        for batch in table.to_batches(max_chunksize=100):
            for row in batch.to_pylist():
                if keep_row(row, language):
                    yield i, row
        del table
        shutil.rmtree(scratch, ignore_errors=True)


def process(rows, processor, embedder, log_every: int = 250):
    """(audio, text) -> (input_ids, mel labels, x-vector), dropping what cannot train."""
    import soundfile as sf
    import torch

    kept = seen = 0
    t0 = time.time()
    for shard_i, row in rows:
        seen += 1
        text = clean_text(row["sentence"])
        if text is None:
            continue

        wav, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32")
        if sr != SR:
            raise SystemExit(f"expected {SR} Hz, shard {shard_i} gave {sr}")
        if wav.ndim > 1:
            wav = wav.mean(axis=1)

        example = processor(
            text=text, audio_target=wav, sampling_rate=SR, return_attention_mask=False
        )
        input_ids = example["input_ids"]
        labels = np.asarray(example["labels"][0], dtype=np.float32)
        if len(input_ids) > MAX_INPUT_IDS or len(labels) > MAX_MEL_FRAMES:
            continue

        with torch.no_grad():
            emb = embedder.encode_batch(torch.tensor(wav).unsqueeze(0))
            emb = torch.nn.functional.normalize(emb, dim=2).squeeze().cpu().numpy()

        kept += 1
        if kept % log_every == 0:
            rate = seen / max(time.time() - t0, 1e-9)
            print(
                f"  shard {shard_i:3d}  scanned {seen:6d}  kept {kept:6d}"
                f"  ({rate:.1f} rows/s)",
                flush=True,
            )

        yield {
            "input_ids": list(map(int, input_ids)),
            "labels": labels.reshape(-1).tolist(),
            "n_mel_frames": int(labels.shape[0]),
            "speaker_embeddings": emb.astype(np.float32).tolist(),
            "text": text,
            "speaker_id": row["speaker_id"],
            "duration": float(row["duration"]),
        }


def write_parquet(records, out_dir: Path, rows_per_file: int = 1000) -> int:
    """Stream records to parquet shards so peak memory stays near one shard."""
    import pyarrow as pa
    import pyarrow.parquet as pq

    out_dir.mkdir(parents=True, exist_ok=True)
    buf, n, part = [], 0, 0

    def flush() -> None:
        nonlocal buf, part
        if not buf:
            return
        pq.write_table(pa.Table.from_pylist(buf), out_dir / f"part-{part:04d}.parquet")
        part += 1
        buf = []

    for rec in records:
        buf.append(rec)
        n += 1
        if len(buf) >= rows_per_file:
            flush()
    flush()
    return n


def main() -> None:
    sys.stdout.reconfigure(encoding="utf-8")
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--language", required=True, help="PLD ISO 639-3 code, e.g. ceb")
    ap.add_argument("--out", default="/content/pld_proc", help="local parquet dir")
    ap.add_argument("--push-to", default="", help="Hub dataset repo to upload to")
    ap.add_argument("--private", action="store_true")
    ap.add_argument("--shards", type=int, default=N_TRAIN_SHARDS,
                    help="how many train shards to scan (fewer = smaller sample)")
    ap.add_argument("--max-samples", type=int, default=0, help="0 = every usable clip")
    args = ap.parse_args()

    os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
    token = os.environ.get("HF_TOKEN") or None

    from transformers import SpeechT5Processor

    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
    embedder = build_embedder()

    out_dir = Path(args.out) / args.language
    shutil.rmtree(out_dir, ignore_errors=True)

    print(f"scanning {args.shards} PLD train shards for {args.language} ...", flush=True)
    rows = iter_shard_rows(args.language, token, range(args.shards))
    records = process(rows, processor, embedder)
    if args.max_samples:
        import itertools

        records = itertools.islice(records, args.max_samples)

    t0 = time.time()
    n = write_parquet(records, out_dir)
    print(f"wrote {n} rows to {out_dir} in {(time.time() - t0) / 60:.1f} min")
    if n == 0:
        raise SystemExit(f"no usable {args.language} rows -- check the filters")

    size = sum(f.stat().st_size for f in out_dir.glob("*.parquet")) / 1024**3
    print(f"parquet size: {size:.2f} GB")

    if args.push_to:
        from huggingface_hub import HfApi

        api = HfApi(token=token)
        api.create_repo(args.push_to, repo_type="dataset", exist_ok=True,
                        private=args.private)
        api.upload_folder(
            folder_path=str(out_dir),
            repo_id=args.push_to,
            repo_type="dataset",
            path_in_repo=f"data/{args.language}",
            commit_message=f"Preprocessed {args.language} TTS inputs ({n} clips)",
        )
        print(f"pushed: https://huggingface.co/datasets/{args.push_to}")


if __name__ == "__main__":
    main()