"""
Thorsten-Voice/Kokoro - German TTS inference example.

Downloads model, config and voicepack from the Hugging Face Hub
(no local files required) and synthesizes German speech in Thorsten's voice.

Usage:
    python inference.py "Hallo, hier spricht Thorsten." output.wav
    python inference.py "Hallo, hier spricht Thorsten." output.wav ep10
    python inference.py "Hallo, hier spricht Thorsten." output.wav ep3
"""
import sys

import numpy as np
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download

REPO_ID = "Thorsten-Voice/Kokoro"
BASE_REPO_ID = "hexgrad/Kokoro-82M"  # architecture reference only, no download from here
SAMPLE_RATE = 24000

# All 10 Stage 2 checkpoints are available (ep1 .. ep10). Epoch 5 is the
# default: in informal listening comparisons it was judged more natural, with
# a slightly slower, less "clipped" speaking pace. Epoch 10 had marginally
# better pitch (F0) loss and a faster, tighter delivery. See the model card
# for the full per-epoch loss table.
CHECKPOINTS = {f"ep{n}": {"model": f"model_ep{n}.pth", "voice": f"voices/thorsten_ep{n}.pt"} for n in range(1, 11)}
CHECKPOINTS["ep5"] = {"model": "model.pth", "voice": "voices/thorsten.pt"}  # default files reuse epoch 5
CHECKPOINTS["default"] = CHECKPOINTS["ep5"]


def load_pipeline(variant: str = "default", device: str | None = None):
    from kokoro import KModel, KPipeline

    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    if variant not in CHECKPOINTS:
        raise ValueError(f"Unknown checkpoint '{variant}'. Available: {sorted(CHECKPOINTS)}")
    files = CHECKPOINTS[variant]

    config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
    model_path = hf_hub_download(repo_id=REPO_ID, filename=files["model"])
    voice_path = hf_hub_download(repo_id=REPO_ID, filename=files["voice"])

    kmodel = KModel(repo_id=BASE_REPO_ID, config=config_path, model=model_path)
    kmodel = kmodel.to(device).eval()

    pipeline = KPipeline(lang_code="d", repo_id=BASE_REPO_ID, model=kmodel)

    # Workaround: misaki's German G2P frontend can emit 'ʏ' (short ü, e.g. in
    # "Brücke") which is not part of Kokoro's 178-symbol vocabulary (only the
    # long-ü symbol 'y' is). Left unpatched, this silently breaks short-ü
    # words at inference time. We wrap the G2P call to apply the same
    # substitution that was used when preparing the training data.
    _original_g2p = pipeline.g2p

    def _patched_g2p(text):
        phonemes, tokens = _original_g2p(text)
        return phonemes.replace("ʏ", "y"), tokens

    pipeline.g2p = _patched_g2p

    voice = torch.load(voice_path, map_location="cpu", weights_only=True)
    return pipeline, voice, device


def synthesize(text: str, output_path: str = "output.wav", speed: float = 1.0, variant: str = "default") -> None:
    pipeline, voice, device = load_pipeline(variant)
    print(f"Using device: {device} | checkpoint: {variant}")

    audio_chunks = []
    for _, phonemes, audio in pipeline(text, voice=voice, speed=speed):
        print(f"Phonemes: {phonemes}")
        audio_chunks.append(audio)

    if not audio_chunks:
        print("WARNING: no audio was generated")
        return

    combined = np.concatenate(audio_chunks)
    sf.write(output_path, combined, SAMPLE_RATE)
    print(f"Saved: {output_path} ({len(combined) / SAMPLE_RATE:.1f}s)")


if __name__ == "__main__":
    input_text = sys.argv[1] if len(sys.argv) > 1 else "Hallo, hier spricht Thorsten."
    output_file = sys.argv[2] if len(sys.argv) > 2 else "output.wav"
    checkpoint_variant = sys.argv[3] if len(sys.argv) > 3 else "default"
    synthesize(input_text, output_file, variant=checkpoint_variant)