from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

REPO_ID = "hexgrad/Kokoro-82M"
SAMPLE_RATE = 24000

ROOT = Path(__file__).resolve().parent
ASSET_DIR = ROOT
CONFIG_FILE = ASSET_DIR / "config.json"
MODEL_FILE = ASSET_DIR / "kokoro-v1_0.pth"
VOICE_DIR = ASSET_DIR / "voices"

LANG_NAMES = {
    "a": "American English",
    "b": "British English",
    "e": "Spanish",
    "f": "French",
    "h": "Hindi",
    "i": "Italian",
    "j": "Japanese",
    "p": "Brazilian Portuguese",
    "z": "Mandarin Chinese",
}

LANG_ALIASES = {
    "auto": "auto",
    "a": "a",
    "en-us": "a",
    "us": "a",
    "american": "a",
    "b": "b",
    "en-gb": "b",
    "gb": "b",
    "british": "b",
    "e": "e",
    "es": "e",
    "spanish": "e",
    "f": "f",
    "fr": "f",
    "fr-fr": "f",
    "french": "f",
    "h": "h",
    "hi": "h",
    "hindi": "h",
    "i": "i",
    "it": "i",
    "italian": "i",
    "j": "j",
    "ja": "j",
    "japanese": "j",
    "p": "p",
    "pt": "p",
    "pt-br": "p",
    "portuguese": "p",
    "brazilian-portuguese": "p",
    "z": "z",
    "zh": "z",
    "zh-cn": "z",
    "mandarin": "z",
    "chinese": "z",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate WAV audio with local Kokoro-82M model and voices."
    )
    parser.add_argument(
        "text_args",
        nargs="*",
        help="Text to synthesize. If omitted, uses --text, --text-file, stdin, or a short sample.",
    )
    parser.add_argument("-t", "--text", help="Text to synthesize.")
    parser.add_argument("--text-file", type=Path, help="UTF-8 text file to synthesize.")
    parser.add_argument(
        "-v",
        "--voice",
        default="af_heart",
        help="Voice name, .pt path, or comma-separated blend. Example: af_heart or af_heart,af_bella",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("kokoro_sample.wav"),
        help="Output WAV path.",
    )
    parser.add_argument("--speed", type=float, default=1.0, help="Speech speed multiplier.")
    parser.add_argument(
        "-l",
        "--lang",
        default="auto",
        help="Language code or alias. Default auto-infers from the first voice letter.",
    )
    parser.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda", "mps"],
        default="auto",
        help="Torch device to use.",
    )
    parser.add_argument(
        "--split-pattern",
        default=r"\n+",
        help="Regex used to split text into sections before synthesis.",
    )
    parser.add_argument(
        "--gap",
        type=float,
        default=0.12,
        help="Seconds of silence to insert between generated chunks.",
    )
    parser.add_argument("--list-voices", action="store_true", help="Print local voices and exit.")
    parser.add_argument("--verbose", action="store_true", help="Print generated chunks and phonemes.")
    return parser.parse_args()


def require_assets() -> None:
    missing = [p for p in (CONFIG_FILE, MODEL_FILE, VOICE_DIR) if not p.exists()]
    if missing:
        formatted = "\n".join(f"  - {p}" for p in missing)
        raise SystemExit(f"Missing Kokoro asset(s):\n{formatted}")


def require_python_version() -> None:
    if not ((3, 10) <= sys.version_info[:2] < (3, 13)):
        version = ".".join(map(str, sys.version_info[:3]))
        raise SystemExit(
            "Kokoro's current PyPI runtime requires Python 3.10, 3.11, or 3.12.\n"
            f"You are running Python {version}. Use Python 3.12 for this script."
        )


def available_voices() -> list[str]:
    if not VOICE_DIR.exists():
        return []
    return sorted(path.stem for path in VOICE_DIR.glob("*.pt"))


def print_voices() -> None:
    voices = available_voices()
    print(f"{len(voices)} local voices in {VOICE_DIR}:")
    for code, language in LANG_NAMES.items():
        group = [voice for voice in voices if voice.startswith(code)]
        if group:
            print(f"\n{code} - {language}")
            print("  " + ", ".join(group))


def resolve_voice_files(voice_arg: str) -> list[Path]:
    voices = [part.strip() for part in voice_arg.split(",") if part.strip()]
    if not voices:
        raise SystemExit("No voice was provided.")

    resolved: list[Path] = []
    for voice in voices:
        candidate = Path(voice).expanduser()
        if candidate.suffix == ".pt":
            if candidate.exists():
                resolved.append(candidate.resolve())
                continue
            local_by_name = VOICE_DIR / candidate.name
            if local_by_name.exists():
                resolved.append(local_by_name)
                continue

        local = VOICE_DIR / f"{voice}.pt"
        if local.exists():
            resolved.append(local)
            continue

        raise SystemExit(f"Unknown voice '{voice}'. Run `python kokoro/sample.py --list-voices`.")

    return resolved


def normalize_lang(value: str) -> str:
    key = (value or "auto").lower()
    lang = LANG_ALIASES.get(key, key)
    if lang != "auto" and lang not in LANG_NAMES:
        valid = ", ".join(sorted(LANG_NAMES))
        raise SystemExit(f"Unknown language '{value}'. Valid codes: {valid}, or auto.")
    return lang


def infer_lang(voice_files: list[Path], requested: str) -> str:
    lang = normalize_lang(requested)
    if lang != "auto":
        return lang

    inferred = {path.stem[0].lower() for path in voice_files if path.stem}
    inferred &= set(LANG_NAMES)
    if len(inferred) == 1:
        return next(iter(inferred))
    if len(inferred) > 1:
        langs = ", ".join(sorted(inferred))
        raise SystemExit(f"Voice blend spans languages ({langs}); pass --lang explicitly.")
    raise SystemExit("Could not infer language from voice name; pass --lang explicitly.")


def read_text(args: argparse.Namespace) -> str:
    parts: list[str] = []
    if args.text:
        parts.append(args.text)
    if args.text_file:
        parts.append(args.text_file.read_text(encoding="utf-8"))
    if args.text_args:
        parts.append(" ".join(args.text_args))
    if parts:
        return "\n".join(parts)
    if not sys.stdin.isatty():
        return sys.stdin.read()
    return "Kokoro is ready to generate speech with any local voice."


def same_path(entry: str, target: Path) -> bool:
    try:
        return Path(entry or ".").resolve() == target
    except OSError:
        return False


def import_runtime():
    # The local asset directory is named "kokoro", so temporarily remove it
    # and its parent from sys.path to avoid shadowing the installed package.
    removed: list[tuple[int, str]] = []
    for index, entry in reversed(list(enumerate(sys.path))):
        if same_path(entry, ROOT) or same_path(entry, ROOT.parent):
            removed.append((index, entry))
            sys.path.pop(index)

    try:
        import numpy as np
        import soundfile as sf
        import torch
        from kokoro import KModel, KPipeline
    except Exception as exc:
        install = (
            "Missing Kokoro runtime dependency.\n"
            "Install Python packages with:\n"
            "  python -m pip install -r kokoro/requirements.txt\n\n"
            "Use Python 3.10, 3.11, or 3.12. Python 3.13+ is too new for Kokoro's current PyPI package.\n"
            "For Japanese/Chinese voices, requirements.txt includes the misaki extras."
        )
        raise RuntimeError(install) from exc
    finally:
        for index, entry in sorted(removed):
            sys.path.insert(index, entry)

    return np, sf, torch, KModel, KPipeline


def select_device(torch, requested: str) -> str:
    if requested == "cuda" and not torch.cuda.is_available():
        raise SystemExit("CUDA was requested, but torch.cuda.is_available() is false.")
    if requested == "mps":
        has_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
        if not has_mps:
            raise SystemExit("MPS was requested, but torch.backends.mps.is_available() is false.")
        if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") != "1":
            raise SystemExit("Set PYTORCH_ENABLE_MPS_FALLBACK=1 before using --device mps.")
    if requested != "auto":
        return requested
    if torch.cuda.is_available():
        return "cuda"
    has_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
    if has_mps and os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") == "1":
        return "mps"
    return "cpu"


def load_voice_tensor(torch, voice_files: list[Path]):
    packs = [torch.load(str(path), weights_only=True) for path in voice_files]
    if len(packs) == 1:
        return packs[0]
    return torch.mean(torch.stack(packs), dim=0)


def language_dependency_hint(lang: str, exc: Exception) -> RuntimeError:
    suffix = ""
    if lang in {"e", "f", "h", "i", "p"}:
        suffix = "\nInstall espeak-ng and make sure it is available on PATH."
    elif lang == "j":
        suffix = "\nInstall Japanese support with: python -m pip install \"misaki[ja]>=0.9.4\""
    elif lang == "z":
        suffix = "\nInstall Chinese support with: python -m pip install \"misaki[zh]>=0.9.4\""
    return RuntimeError(f"Failed to initialize {LANG_NAMES.get(lang, lang)} pipeline.{suffix}")


def synthesize(args: argparse.Namespace) -> None:
    require_assets()
    require_python_version()
    voice_files = resolve_voice_files(args.voice)
    lang = infer_lang(voice_files, args.lang)
    text = read_text(args).strip()
    if not text:
        raise SystemExit("No text to synthesize.")

    np, sf, torch, KModel, KPipeline = import_runtime()
    device = select_device(torch, args.device)
    voice_tensor = load_voice_tensor(torch, voice_files)

    model = KModel(repo_id=REPO_ID, config=str(CONFIG_FILE), model=str(MODEL_FILE))
    model = model.to(device).eval()

    try:
        pipeline = KPipeline(lang_code=lang, repo_id=REPO_ID, model=model)
    except Exception as exc:
        raise language_dependency_hint(lang, exc) from exc

    audios = []
    generator = pipeline(
        text,
        voice=voice_tensor,
        speed=args.speed,
        split_pattern=args.split_pattern,
    )

    for index, result in enumerate(generator):
        if args.verbose:
            print(f"[{index}] {result.graphemes}")
            print(f"    phonemes: {result.phonemes}")
        audio = result.audio
        if audio is None:
            continue
        if hasattr(audio, "detach"):
            audio = audio.detach().cpu().numpy()
        audios.append(np.asarray(audio, dtype=np.float32))

    if not audios:
        raise SystemExit("Kokoro did not generate any audio. Check the text and language.")

    if len(audios) == 1 or args.gap <= 0:
        merged = np.concatenate(audios)
    else:
        silence = np.zeros(max(0, int(SAMPLE_RATE * args.gap)), dtype=np.float32)
        chunks = []
        for index, audio in enumerate(audios):
            if index:
                chunks.append(silence)
            chunks.append(audio)
        merged = np.concatenate(chunks)

    output = args.output.expanduser()
    output.parent.mkdir(parents=True, exist_ok=True)
    sf.write(str(output), merged, SAMPLE_RATE)

    voice_names = ",".join(path.stem for path in voice_files)
    print(f"Wrote {output} using voice={voice_names}, lang={lang}, device={device}")


def main() -> None:
    args = parse_args()
    if args.list_voices:
        require_assets()
        print_voices()
        return
    try:
        synthesize(args)
    except RuntimeError as exc:
        raise SystemExit(str(exc)) from exc


if __name__ == "__main__":
    main()