#!/usr/bin/env python3
"""Transcribe cloned-voices/*.wav with whisper-large-v3 and write ref_texts.

Outputs:
  - cloned-voices/whisper_refs.json         (raw transcripts keyed by voice name)
  - voice_to_speaker.default.yaml           (ref_text rewritten in place)
  - cloned-voices/voices_metadata.json      (ref_text rewritten in place)

F5-TTS cloning quality depends on ref_text matching the prosody of ref_audio
(commas, periods, casing). LibriSpeech ground-truth labels are ALL CAPS with
no punctuation, which is the wrong signal for a flow-matching TTS conditioned
on text. Whisper hears what F5 will hear, so its transcript is the better ref.

Run on a GPU host (4090 / 3090). Idempotent: rerun any time cloned-voices/
changes. Use `make whisper-refs` rather than calling this directly.

Pass --from-cache to skip ASR and only re-apply normalization from an existing
whisper_refs.json — useful after tweaking normalize_text(). No GPU needed.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent.parent
VOICES_DIR = REPO_ROOT / "cloned-voices"
YAML_PATH = REPO_ROOT / "voice_to_speaker.default.yaml"
METADATA_PATH = VOICES_DIR / "voices_metadata.json"
WHISPER_JSON = VOICES_DIR / "whisper_refs.json"
MODEL_ID = os.environ.get("WHISPER_MODEL", "openai/whisper-large-v3")


def normalize_text(text: str) -> str:
    """Clean Whisper output for use as F5 ref_text.

    Whisper occasionally hallucinates `"'` cluster characters when it interprets
    a fragment as quoted dialogue — at sentence start, mid-clause, anywhere. It
    also returns lowercase output and drops terminal punctuation on some clips.
    F5 tokenizes the raw string, so messy ref_text → noisy conditioning signal.
    Designed to be idempotent: running it twice on the same input is a no-op.
    """
    text = re.sub(r"\s+", " ", text).strip()
    # `"'` never appears as legitimate English punctuation — always Whisper noise
    text = text.replace('"\'', "")
    # strip leading whitespace and opening-quote junk
    text = re.sub(r'^[\s"\'`]+', "", text)
    # strip trailing whitespace and any quote/apostrophe chars (legit closing quote
    # after a terminal `.` is also dropped — F5 only cares about the prosody marker)
    text = re.sub(r'[\s"\'`]+$', "", text)
    # if a terminal punct is followed by stray apostrophe-then-period (`.'.`),
    # collapse to the terminal — happens when a closing-quoted line gets a `.` appended
    text = re.sub(r"([.!?])['\"`]+\.?$", r"\1", text)
    text = re.sub(r"\s+", " ", text).strip()
    if text and text[0].isalpha():
        text = text[0].upper() + text[1:]
    if text and text[-1] not in ".!?":
        text = text + "."
    return text


def load_pipeline():
    import torch
    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device.startswith("cuda") else torch.float32
    print(f"[whisper_refs] loading {MODEL_ID} on {device} ({dtype})", flush=True)
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        MODEL_ID, torch_dtype=dtype, low_cpu_mem_usage=True, use_safetensors=True
    ).to(device)
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    return pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=dtype,
        device=device,
        return_timestamps=False,
    )


def transcribe_all(asr) -> dict[str, str]:
    wavs = sorted(VOICES_DIR.glob("*.wav"))
    print(f"[whisper_refs] {len(wavs)} wavs to transcribe", flush=True)
    out: dict[str, str] = {}
    for i, wav in enumerate(wavs, 1):
        result = asr(
            str(wav),
            generate_kwargs={"language": "en", "task": "transcribe"},
        )
        text = normalize_text(result["text"])
        voice = wav.stem
        out[voice] = text
        print(f"[whisper_refs] [{i:2d}/{len(wavs)}] {voice:10s} -> {text}", flush=True)
    return out


def rewrite_yaml(refs: dict[str, str]) -> int:
    """Surgical replace of ref_text values keyed by ref_audio filename.

    Avoids a full YAML round-trip so comments/order/whitespace stay untouched.
    Each voice block contains a `ref_audio: cloned-voices/<voice>.wav` line
    followed (next non-blank line) by `ref_text: "..."`. Match on the audio
    path and rewrite the next ref_text line.
    """
    lines = YAML_PATH.read_text().splitlines(keepends=True)
    audio_re = re.compile(r"^(\s*)ref_audio:\s*cloned-voices/([^\s.]+)\.wav\s*$")
    text_re = re.compile(r"^(\s*)ref_text:\s*.*$")
    changed = 0
    i = 0
    while i < len(lines):
        m = audio_re.match(lines[i].rstrip("\n"))
        if not m:
            i += 1
            continue
        voice = m.group(2)
        if voice not in refs:
            i += 1
            continue
        # find the next ref_text line within the same block (no blank line break)
        j = i + 1
        while j < len(lines) and lines[j].strip() != "":
            tm = text_re.match(lines[j].rstrip("\n"))
            if tm:
                indent = tm.group(1)
                escaped = refs[voice].replace("\\", "\\\\").replace('"', '\\"')
                lines[j] = f'{indent}ref_text: "{escaped}"\n'
                changed += 1
                break
            j += 1
        i = j + 1
    YAML_PATH.write_text("".join(lines))
    return changed


def rewrite_metadata(refs: dict[str, str]) -> int:
    data = json.loads(METADATA_PATH.read_text())
    changed = 0
    for voice, entry in data.items():
        if voice in refs and entry.get("ref_text") != refs[voice]:
            entry["ref_text"] = refs[voice]
            changed += 1
    METADATA_PATH.write_text(json.dumps(data, indent=2) + "\n")
    return changed


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument(
        "--from-cache",
        action="store_true",
        help="skip ASR; re-apply normalize_text() to cached whisper_refs.json",
    )
    args = parser.parse_args()

    if not VOICES_DIR.is_dir():
        print(f"[whisper_refs] no cloned-voices dir at {VOICES_DIR}", file=sys.stderr)
        return 2

    if args.from_cache:
        if not WHISPER_JSON.is_file():
            print(f"[whisper_refs] --from-cache but {WHISPER_JSON} missing", file=sys.stderr)
            return 2
        raw = json.loads(WHISPER_JSON.read_text())
        refs = {voice: normalize_text(text) for voice, text in raw.items()}
        print(f"[whisper_refs] re-normalized {len(refs)} entries from cache", flush=True)
    else:
        asr = load_pipeline()
        refs = transcribe_all(asr)

    WHISPER_JSON.write_text(json.dumps(refs, indent=2, ensure_ascii=False) + "\n")
    print(f"[whisper_refs] wrote {WHISPER_JSON} ({len(refs)} entries)", flush=True)
    yaml_changed = rewrite_yaml(refs)
    meta_changed = rewrite_metadata(refs)
    print(
        f"[whisper_refs] yaml: {yaml_changed} ref_text replaced | "
        f"metadata: {meta_changed} ref_text replaced",
        flush=True,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())