#!/usr/bin/env python3
"""
Download diverse voice samples for Qwen3-TTS voice cloning.
Uses voice_registry.json for idempotent, permanent speaker-to-name assignments.
Fetches gender from upstream SPEAKERS.TXT (OpenSLR) to ensure correct assignment.

Requires: pip install datasets soundfile

Usage:
    python scripts/download_diverse_voices.py
    python scripts/download_diverse_voices.py --registry voice_registry.json
    python scripts/download_diverse_voices.py --corpora librispeech-test-clean
"""

import os
import json
import argparse
import urllib.request
from pathlib import Path

try:
    from datasets import load_dataset
    import soundfile as sf
    HAS_DATASETS = True
except ImportError:
    HAS_DATASETS = False
    print("Install required libraries: pip install datasets soundfile")

# GitHub mirror of LibriSpeech SPEAKERS.TXT (plain text, easier to parse)
SPEAKERS_TXT_GITHUB = "https://raw.githubusercontent.com/oscarknagg/voicemap/master/data/LibriSpeech/SPEAKERS.TXT"


def fetch_speaker_genders():
    """Fetch gender info from upstream LibriSpeech SPEAKERS.TXT."""
    print("Fetching speaker genders from upstream SPEAKERS.TXT...")
    try:
        req = urllib.request.Request(SPEAKERS_TXT_GITHUB, headers={"User-Agent": "uncloseai-speech"})
        with urllib.request.urlopen(req, timeout=15) as resp:
            text = resp.read().decode("utf-8")
    except Exception as e:
        print(f"  WARNING: Failed to fetch SPEAKERS.TXT: {e}")
        return {}

    genders = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith(";"):
            continue
        # Format: ID  | SEX | SUBSET | MINUTES | NAME
        parts = [p.strip() for p in line.split("|")]
        if len(parts) >= 2:
            try:
                speaker_id = parts[0].strip()
                sex = parts[1].strip().upper()
                if sex in ("F", "M"):
                    genders[speaker_id] = "female" if sex == "F" else "male"
            except (ValueError, IndexError):
                continue

    print(f"  Loaded genders for {len(genders)} speakers")
    return genders


def load_registry(registry_path):
    """Load voice registry from JSON file."""
    if not registry_path.exists():
        return None
    with open(registry_path) as f:
        return json.load(f)


def save_registry(registry, registry_path):
    """Save voice registry to JSON file (append-only, never remove voices)."""
    with open(registry_path, "w") as f:
        json.dump(registry, f, indent=2)
    print(f"  Registry saved: {registry_path}")


def get_used_names(registry):
    """Get set of names already assigned in the registry."""
    return set(registry.get("voices", {}).keys())


def get_next_name(registry, gender):
    """Get the next available name from the pool for the given gender."""
    used = get_used_names(registry)
    pool = registry.get("name_pools", {}).get(gender, [])
    for name in pool:
        if name not in used:
            return name
    return None


def assign_speakers_to_names(registry, corpus_id, speakers_by_gender, speaker_genders):
    """Assign names to new speakers deterministically.

    New speakers are sorted by ID (ascending) and assigned names in pool order.
    Existing assignments are never changed.
    """
    # Build reverse lookup: (corpus, speaker_id) -> name
    assigned = {}
    for name, info in registry.get("voices", {}).items():
        key = (info["corpus"], info["speaker_id"])
        assigned[key] = name

    new_assignments = []

    for gender in ("female", "male"):
        # Sort unassigned speakers by ID for determinism
        speakers = sorted(speakers_by_gender.get(gender, []), key=lambda s: int(s))
        for sid in speakers:
            key = (corpus_id, sid)
            if key in assigned:
                continue  # Already has a name

            name = get_next_name(registry, gender)
            if name is None:
                print(f"  WARNING: No more {gender} names available, skipping speaker {sid}")
                continue

            registry["voices"][name] = {
                "corpus": corpus_id,
                "speaker_id": sid,
                "gender": gender,
                "locked": True,
            }
            new_assignments.append((name, sid, gender))
            print(f"  NEW: {name:12s} <- speaker {sid} ({gender})")

    return new_assignments


def pick_best_sample(samples, min_dur=4.0, max_dur=12.0, target=7.0):
    """Pick the best sample: prefer 5-10 seconds, clean, complete sentence."""
    best = None
    best_score = float('-inf')

    for s in samples:
        audio = s["audio"]
        dur = len(audio["array"]) / audio["sampling_rate"]
        text = s.get("text", "")

        # Skip too short
        if dur < 3.0:
            continue

        # Score: prefer target duration, penalize extremes
        score = -abs(dur - target)

        # Bonus for ending with period (complete sentence)
        if text.strip().endswith('.'):
            score += 2.0

        # Bonus for being in ideal range
        if min_dur <= dur <= max_dur:
            score += 5.0

        # Penalty for very long text (harder for model)
        if len(text) > 300:
            score -= 3.0

        if score > best_score:
            best = s
            best_score = score

    return best


def main():
    parser = argparse.ArgumentParser(description="Download diverse voice samples for Qwen3-TTS")
    parser.add_argument("-o", "--output-dir", default="cloned-voices",
                        help="Output directory for voice WAV files")
    parser.add_argument("-c", "--config-output", default="voice_to_speaker.default.yaml",
                        help="Output path for voice config YAML")
    parser.add_argument("--config-runtime", default="config/voice_to_speaker.yaml",
                        help="Runtime config path (also written if dir exists)")
    parser.add_argument("--registry", default="voice_registry.json",
                        help="Path to voice registry JSON (default: voice_registry.json)")
    parser.add_argument("--corpora", nargs="+", default=["librispeech-test-clean"],
                        help="Corpora to download (default: librispeech-test-clean)")
    parser.add_argument("--max-samples", type=int, default=20,
                        help="Max samples to collect per speaker for selection")
    args = parser.parse_args()

    if not HAS_DATASETS:
        print("ERROR: Install required libraries first:")
        print("  pip install datasets soundfile")
        return 1

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    registry_path = Path(args.registry)

    # Load or create registry
    registry = load_registry(registry_path)
    if registry is None:
        print(f"ERROR: Registry not found at {registry_path}")
        print("  Create it or copy from voice_registry.json")
        return 1

    print(f"Loaded registry: {len(registry.get('voices', {}))} existing voices")

    # Fetch genders from upstream
    speaker_genders = fetch_speaker_genders()
    if not speaker_genders:
        print("ERROR: Could not fetch speaker genders. Cannot assign gendered names.")
        return 1

    print(f"\n=== Downloading Diverse Voice Samples for Qwen3-TTS ===\n")
    print(f"Output directory: {output_dir}")
    print(f"Config output: {args.config_output}")
    print(f"Registry: {registry_path}")
    print(f"Corpora: {', '.join(args.corpora)}")

    # Process each corpus
    all_voice_names = set()

    for corpus_id in args.corpora:
        corpus_config = registry.get("corpora", {}).get(corpus_id)
        if corpus_config is None:
            print(f"\nERROR: Unknown corpus '{corpus_id}'. Available: {list(registry.get('corpora', {}).keys())}")
            continue

        print(f"\n--- Corpus: {corpus_id} ---")
        print(f"  {corpus_config.get('description', '')}")

        # Load dataset
        print(f"  Loading {corpus_config['dataset']} ({corpus_config['config']}/{corpus_config['split']})...")
        print("  (First run downloads ~1.5 GB, cached after that)\n")
        dataset = load_dataset(
            corpus_config["dataset"],
            corpus_config["config"],
            split=corpus_config["split"],
            trust_remote_code=True
        )

        # Group samples by speaker
        print("  Grouping samples by speaker...")
        speaker_samples = {}
        for sample in dataset:
            sid = str(sample["speaker_id"])
            if sid not in speaker_samples:
                speaker_samples[sid] = []
            if len(speaker_samples[sid]) < args.max_samples:
                speaker_samples[sid].append(sample)

        print(f"  Found {len(speaker_samples)} speakers\n")

        # Split speakers by gender
        speakers_by_gender = {"female": [], "male": []}
        for sid in speaker_samples:
            gender = speaker_genders.get(sid)
            if gender in ("female", "male"):
                speakers_by_gender[gender].append(sid)
            else:
                print(f"  WARNING: No gender for speaker {sid}, skipping")

        print(f"  Female speakers: {len(speakers_by_gender['female'])}")
        print(f"  Male speakers: {len(speakers_by_gender['male'])}")

        # Assign names to any new speakers
        new = assign_speakers_to_names(registry, corpus_id, speakers_by_gender, speaker_genders)
        if new:
            print(f"\n  Assigned {len(new)} new voices")
            save_registry(registry, registry_path)
        else:
            print(f"\n  No new speakers to assign")

        # Download samples for all voices in this corpus
        corpus_voices = {
            name: info for name, info in registry["voices"].items()
            if info["corpus"] == corpus_id
        }

        print(f"\n  Downloading {len(corpus_voices)} voice samples...\n")

        for voice_name, voice_info in sorted(corpus_voices.items()):
            sid = voice_info["speaker_id"]
            if sid not in speaker_samples:
                print(f"  WARNING: Speaker {sid} ({voice_name}) not in dataset")
                continue

            samples = speaker_samples[sid]
            best = pick_best_sample(samples)
            if best is None:
                print(f"  WARNING: No suitable sample for '{voice_name}' (speaker {sid})")
                continue

            audio = best["audio"]
            transcript = best["text"].strip()
            duration = len(audio["array"]) / audio["sampling_rate"]

            # Save WAV
            out_path = output_dir / f"{voice_name}.wav"
            sf.write(str(out_path), audio["array"], audio["sampling_rate"])

            voice_info["ref_audio"] = f"cloned-voices/{voice_name}.wav"
            voice_info["ref_text"] = transcript
            voice_info["duration"] = round(duration, 1)

            all_voice_names.add(voice_name)
            print(f"  {voice_name:12s} | {voice_info['gender']:6s} | speaker {sid:5s} | {duration:.1f}s")

    print(f"\nDownloaded {len(all_voice_names)} voices total\n")

    # Build voices dict for config generation
    voices = {}
    for name in sorted(all_voice_names):
        info = registry["voices"][name]
        if "ref_audio" in info:
            voices[name] = info

    # Generate YAML config
    print("Generating voice config...")

    female_names = sorted(n for n, v in voices.items() if v["gender"] == "female")
    male_names = sorted(n for n, v in voices.items() if v["gender"] == "male")

    # Order by name pool position for consistent output
    female_pool = registry["name_pools"]["female"]
    male_pool = registry["name_pools"]["male"]
    female_names.sort(key=lambda n: female_pool.index(n) if n in female_pool else 999)
    male_names.sort(key=lambda n: male_pool.index(n) if n in male_pool else 999)

    female_list = ", ".join(female_names)
    male_list = ", ".join(male_names)

    lines = [
        "# uncloseai-speech Voice Configuration",
        "# Diverse voice samples from LibriSpeech test-clean (public domain)",
        "# Each voice is a DISTINCT SPEAKER for Qwen3-TTS voice cloning",
        "# Gender verified from upstream LibriSpeech SPEAKERS.TXT",
        "# Assignments locked in voice_registry.json (idempotent, append-only)",
        "#",
        f"# Female voices: {female_list}",
        f"# Male voices: {male_list}",
        "#",
        "# Source: LibriSpeech test-clean (public domain, LibriVox recordings)",
        f"# {len(voices)} distinct speakers ({len(female_names)} female, {len(male_names)} male)",
        "",
        "tts-1-qwen:",
        "",
    ]

    # Write female voices first, then male (in pool order)
    for voice_name in female_names + male_names:
        if voice_name not in voices:
            continue
        v = voices[voice_name]
        lines.append(f"  # {v['gender']} - speaker {v['speaker_id']}")
        lines.append(f"  {voice_name}:")
        lines.append(f"    ref_audio: {v['ref_audio']}")
        safe_text = v['ref_text'].replace('"', '\\"')
        lines.append(f'    ref_text: "{safe_text}"')
        lines.append(f"    language: English")
        lines.append("")

    # Disabled engines
    lines.extend([
        "# Other TTS engines (disabled by default)",
        "# Uncomment and configure to enable",
        "",
        "# tts-1:",
        "#   # Piper TTS (fast CPU inference)",
        "#   alloy:",
        "#     model: /app/voices/en/en_US/libritts_r/medium/en_US-libritts_r-medium.onnx",
        "#     speaker: 79",
        "",
        "# tts-1-hd:",
        "#   # XTTS v2 (voice cloning)",
        "#   alloy:",
        "#     model: xtts",
        "#     speaker: voices/alloy.wav",
        "",
    ])

    config_text = "\n".join(lines)

    # Write default config
    with open(args.config_output, "w") as f:
        f.write(config_text)
    print(f"  Written: {args.config_output}")

    # Write runtime config if dir exists
    runtime_dir = Path(args.config_runtime).parent
    if runtime_dir.exists():
        with open(args.config_runtime, "w") as f:
            f.write(config_text)
        print(f"  Written: {args.config_runtime}")

    # Save metadata JSON
    meta = {}
    for name in female_names + male_names:
        if name in voices:
            v = voices[name]
            meta[name] = {
                "ref_audio": v["ref_audio"],
                "ref_text": v["ref_text"],
                "speaker_id": v["speaker_id"],
                "gender": v["gender"],
                "duration": v["duration"],
            }

    meta_path = output_dir / "voices_metadata.json"
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)
    print(f"  Written: {meta_path}")

    print(f"\n=== Done! {len(voices)} diverse voices configured ===")
    print(f"\nVoice mapping:")
    for name in female_names + male_names:
        if name in voices:
            v = voices[name]
            print(f"  {name:12s} -> speaker {v['speaker_id']} ({v['gender']})")

    return 0


if __name__ == "__main__":
    exit(main() or 0)
