#!/usr/bin/env python3
"""
Download diverse voice samples for Qwen3-TTS voice cloning.
Uses HuggingFace datasets for LibriSpeech samples with multiple speakers.
"""

import os
import json
import argparse
from pathlib import Path

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

# LibriSpeech test-clean speaker metadata
# Format: speaker_id: (gender, description)
LIBRISPEECH_SPEAKERS = {
    # Female speakers
    "1089": ("female", "clear, professional"),
    "1188": ("female", "warm, measured"),
    "1221": ("female", "expressive, storyteller"),
    "1320": ("female", "soft, gentle"),
    "3570": ("female", "bright, energetic"),
    "3575": ("female", "calm, neutral"),
    "4446": ("female", "mature, authoritative"),
    "4507": ("female", "young, clear"),
    "5142": ("female", "warm, friendly"),
    "6829": ("female", "crisp, articulate"),
    "6930": ("female", "smooth, radio-like"),
    "7729": ("female", "light, airy"),
    "8230": ("female", "rich, full"),
    "8455": ("female", "neutral, newsreader"),

    # Male speakers
    "1284": ("male", "deep, resonant"),
    "1580": ("male", "clear, narrator"),
    "2094": ("male", "warm, baritone"),
    "2830": ("male", "young, energetic"),
    "4077": ("male", "mature, professor"),
    "4970": ("male", "smooth, announcer"),
    "5105": ("male", "neutral, clear"),
    "5639": ("male", "deep, dramatic"),
    "7021": ("male", "light, conversational"),
    "7127": ("male", "authoritative, news"),
    "7176": ("male", "warm, storyteller"),
    "8224": ("male", "crisp, professional"),
    "8463": ("male", "rich, bass"),
}

# Voice name mappings for OpenAI-compatible names
VOICE_MAPPINGS = {
    # OpenAI standard voices
    "alloy": {"speaker": "1089", "desc": "neutral, balanced female"},
    "echo": {"speaker": "1284", "desc": "deep, clear male"},
    "fable": {"speaker": "1221", "desc": "expressive, storyteller female"},
    "onyx": {"speaker": "5639", "desc": "deep, dramatic male"},
    "nova": {"speaker": "1188", "desc": "warm, friendly female"},
    "shimmer": {"speaker": "1320", "desc": "soft, gentle female"},

    # Extended voices - female
    "aurora": {"speaker": "3570", "desc": "bright, energetic female"},
    "bella": {"speaker": "4446", "desc": "mature, authoritative female"},
    "clara": {"speaker": "5142", "desc": "warm, conversational female"},
    "dawn": {"speaker": "6829", "desc": "crisp, articulate female"},
    "ember": {"speaker": "6930", "desc": "smooth, radio-like female"},
    "fiona": {"speaker": "7729", "desc": "light, airy female"},
    "grace": {"speaker": "8230", "desc": "rich, full female"},
    "hazel": {"speaker": "8455", "desc": "neutral, newsreader female"},
    "iris": {"speaker": "3575", "desc": "calm, neutral female"},
    "jade": {"speaker": "4507", "desc": "young, clear female"},

    # Extended voices - male
    "atlas": {"speaker": "1580", "desc": "clear, narrator male"},
    "blaze": {"speaker": "2830", "desc": "young, energetic male"},
    "cedar": {"speaker": "4077", "desc": "mature, professor male"},
    "drake": {"speaker": "4970", "desc": "smooth, announcer male"},
    "eric": {"speaker": "5105", "desc": "neutral, clear male"},
    "felix": {"speaker": "7021", "desc": "light, conversational male"},
    "grant": {"speaker": "7127", "desc": "authoritative, news male"},
    "hugo": {"speaker": "7176", "desc": "warm, storyteller male"},
    "ivan": {"speaker": "8224", "desc": "crisp, professional male"},
    "jack": {"speaker": "2094", "desc": "warm, baritone male"},
    "knox": {"speaker": "8463", "desc": "rich, bass male"},
}


def download_librispeech_samples(output_dir: Path, max_per_speaker: int = 1):
    """Download LibriSpeech test-clean samples for each speaker."""

    if not HAS_DATASETS:
        print("Error: datasets library not installed")
        return {}

    print("Loading LibriSpeech test-clean dataset...")
    dataset = load_dataset(
        "openslr/librispeech_asr",
        "clean",
        split="test",
        trust_remote_code=True
    )

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

    # Download samples for our target speakers
    downloaded = {}
    output_dir.mkdir(parents=True, exist_ok=True)

    for voice_name, voice_info in VOICE_MAPPINGS.items():
        speaker_id = voice_info["speaker"]

        if speaker_id not in speaker_samples:
            print(f"  Warning: Speaker {speaker_id} not found for voice '{voice_name}'")
            continue

        sample = speaker_samples[speaker_id][0]
        audio = sample["audio"]
        transcript = sample["text"]

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

        downloaded[voice_name] = {
            "file": str(output_path),
            "transcript": transcript,
            "speaker_id": speaker_id,
            "description": voice_info["desc"],
            "sample_rate": audio["sampling_rate"]
        }

        print(f"  Downloaded: {voice_name} (speaker {speaker_id})")

    return downloaded


def generate_voice_config(voices: dict, output_file: Path):
    """Generate voice_to_speaker.yaml config."""

    lines = [
        "# uncloseai-speech Voice Configuration",
        "# Auto-generated with diverse LibriSpeech speakers",
        "",
        "tts-1-qwen:",
        "  # OpenAI-compatible voice names with diverse speakers",
        "  # Each voice has a unique speaker from LibriSpeech test-clean",
        "",
    ]

    for voice_name, info in sorted(voices.items()):
        # Use relative path from app root
        rel_path = info["file"].replace("/app/", "").replace(str(Path.cwd()) + "/", "")
        if not rel_path.startswith("voices/"):
            rel_path = f"voices/samples/{voice_name}.wav"

        lines.extend([
            f"  # {info['description']}",
            f"  {voice_name}:",
            f"    ref_audio: {rel_path}",
            f'    ref_text: "{info["transcript"]}"',
            f"    language: English",
            "",
        ])

    with open(output_file, "w") as f:
        f.write("\n".join(lines))

    print(f"\nGenerated config: {output_file}")


def main():
    parser = argparse.ArgumentParser(description="Download diverse voice samples")
    parser.add_argument("--output-dir", "-o", default="voices/samples",
                        help="Output directory for voice samples")
    parser.add_argument("--config-output", "-c", default="config/voice_to_speaker.yaml",
                        help="Output path for voice config")
    parser.add_argument("--max-per-speaker", "-m", type=int, default=1,
                        help="Max samples per speaker")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)

    print(f"Downloading voice samples to: {output_dir}")
    voices = download_librispeech_samples(output_dir, args.max_per_speaker)

    if voices:
        print(f"\nDownloaded {len(voices)} voice samples")

        # Generate config
        config_path = Path(args.config_output)
        config_path.parent.mkdir(parents=True, exist_ok=True)
        generate_voice_config(voices, config_path)

        # Also save metadata
        metadata_path = output_dir / "voices_metadata.json"
        with open(metadata_path, "w") as f:
            json.dump(voices, f, indent=2)
        print(f"Saved metadata: {metadata_path}")
    else:
        print("No voices downloaded")


if __name__ == "__main__":
    main()
