#!/usr/bin/env python3
"""Download and convert every voice from the pinned official Kokoro releases.

All network requests are anonymous and every downloaded LFS object is checked
against the SHA-256 exposed by the source repository at the pinned revision.
"""

from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
import json
from pathlib import Path
import shutil
import tempfile
from typing import Any

from huggingface_hub import HfApi, hf_hub_download
import numpy as np
import requests
import torch


ROOT = Path(__file__).resolve().parents[1]

BASE_REPOSITORY = "hexgrad/Kokoro-82M"
BASE_REVISION = "f3ff3571791e39611d31c381e3a41a3af07b4987"
ZH_REPOSITORY = "hexgrad/Kokoro-82M-v1.1-zh"
ZH_REVISION = "01e7505bd6a7a2ac4975463114c3a7650a9f7218"
LANGUAGE_MODELS_REPOSITORY = "software-mansion/react-native-executorch-kokoro"
LANGUAGE_MODELS_REVISION = "3744b57964eab7df6e8c48f0b84badb29e14df07"
PHONEMIS_REPOSITORY = "IgorSwat/Phonemis"
PHONEMIS_REVISION = "71eb1ce33bd586d38cbac037843b8539d7829c3b"

EXPECTED_VOICE_COUNTS = {
    BASE_REPOSITORY: 54,
    ZH_REPOSITORY: 103,
}

MODEL_FILE_SHA256 = {
    (BASE_REPOSITORY, "config.json"): "5abb01e2403b072bf03d04fde160443e209d7a0dad49a423be15196b9b43c17f",
    (BASE_REPOSITORY, "kokoro-v1_0.pth"): "496dba118d1a58f5f3db2efc88dbdc216e0483fc89fe6e47ee1f2c53f18ad1e4",
    (ZH_REPOSITORY, "config.json"): "bc333efa5ce4ceff433c8c8e5d027a1eca0166001e4e4a62bea2d26ff7a46890",
    (ZH_REPOSITORY, "kokoro-v1_1-zh.pth"): "b1d8410fa44dfb5c15471fd6c4225ea6b4e9ac7fa03c98e8bea47a9928476e2b",
}

LANGUAGE_MODEL_FILES = {
    "finetunes/kokoro_polish_converted.pth": (
        "kokoro-pl-fp32",
        "e3202dc4d1f6e65dddff8a8e8d2e091ce9a3a66cf989c6c581a81b2f1969af49",
    ),
    "finetunes/kokoro_german_converted.pth": (
        "kokoro-de-fp32",
        "b8b2ab322963e7662c6036035c76c34a6a5582f814917407c94d632a1c930f71",
    ),
}

LANGUAGE_VOICE_FILES = {
    "voices/pm_mateusz.bin": (
        "voices/pl/pm_mateusz.bin",
        "dc8f2919ede945e6962310b5204b912e64689d5698c571822bf75ad704870cb4",
    ),
    "voices/df_anna.bin": (
        "voices/de/df_anna.bin",
        "d583ccff3cdca2f7fae535cb998ac07e9fcb90f09737b9a41fa2734ec44a8f0b",
    ),
}

LANGUAGE_BY_PREFIX = {
    "a": "en-us",
    "b": "en-gb",
    "e": "es",
    "f": "fr",
    "h": "hi",
    "i": "it",
    "j": "ja",
    "p": "pt-br",
    "z": "zh",
}

PHONEMIS_FILES = {
    "data/de/phonemizer_de.bin": (
        "phonemizers/de/phonemizer_de.bin",
        "4888dc7e54dc66098551555096562063364091fc246da0d057b629587bedaa0b",
        7_094_120,
    ),
    "data/en-gb/lexicon_full.json": (
        "phonemizers/en-gb/lexicon_full.json",
        "52167ca536a93d56a02e8b1db29f572438cb7103737ef2c3f8dd9dbfc0b35e0a",
        6_664_570,
    ),
    "data/en-gb/phonemizer_en_gb.bin": (
        "phonemizers/en-gb/phonemizer_en_gb.bin",
        "3d4fe5a541c02de30879a5b84b88f5229f5b31f8fb2dcf7e89a1ce3a6c330334",
        7_094_120,
    ),
    "data/en-gb/tagger.json": (
        "phonemizers/en-gb/tagger.json",
        "af2fe9831e8560fa78ebf7d96da715ce5ecb43a3363bd701c952a6db206f169c",
        1_634_732,
    ),
    "data/en-us/lexicon_full.json": (
        "phonemizers/en-us/lexicon_full.json",
        "ef0b19a0126455e4216fb08083c8b50f7e85f98f6055738129a89d2095e635d7",
        6_254_347,
    ),
    "data/en-us/phonemizer_en_us.bin": (
        "phonemizers/en-us/phonemizer_en_us.bin",
        "e059561fb8d51eadfd2000965be30e31f0152e7e8c8b4fcf7859dbcce8557576",
        7_094_120,
    ),
    "data/en-us/tagger.json": (
        "phonemizers/en-us/tagger.json",
        "af2fe9831e8560fa78ebf7d96da715ce5ecb43a3363bd701c952a6db206f169c",
        1_634_732,
    ),
    "data/es/phonemizer_es.bin": (
        "phonemizers/es/phonemizer_es.bin",
        "8dc68946e12c1a233ac9153fd369f6418b930a3d731587ff64dc789f44d75a52",
        7_094_120,
    ),
    "data/fr/phonemizer_fr.bin": (
        "phonemizers/fr/phonemizer_fr.bin",
        "fa0018b750a3670328026107b44e02eafe3223b6878585ca0927786b7812d96a",
        7_094_120,
    ),
    "data/hi/phonemizer_hi.bin": (
        "phonemizers/hi/phonemizer_hi.bin",
        "dcee3272f96d7f1b7cc40c5df23060b502a5b9066f7cb20092413dff82a487f5",
        9_194_344,
    ),
    "data/it/phonemizer_it.bin": (
        "phonemizers/it/phonemizer_it.bin",
        "dca8d068d76134a40856cd874ea6e5e05c988914369ab643f52ab77a512067e3",
        7_094_120,
    ),
    "data/pl/phonemizer_pl.bin": (
        "phonemizers/pl/phonemizer_pl.bin",
        "ec85f4dc2c4ac7a72ff88b98b0664a4ed887bd15c0d5add2eb1d6a6ee05b73f2",
        7_094_120,
    ),
    "data/pt/phonemizer_pt.bin": (
        "phonemizers/pt-br/phonemizer_pt.bin",
        "89049ea03c52ffa7233a343d44a35059aae3b1231d772c8c498fbc4427756ecf",
        7_094_120,
    ),
}


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def download_hf_file(
    repository: str,
    revision: str,
    filename: str,
    cache_dir: Path,
) -> Path:
    return Path(
        hf_hub_download(
            repository,
            filename=filename,
            revision=revision,
            cache_dir=cache_dir,
            token=False,
        )
    )


def convert_voice(source: Path, destination: Path) -> None:
    voice = torch.load(source, map_location="cpu", weights_only=True)
    if not isinstance(voice, torch.Tensor):
        raise TypeError(f"Expected a tensor in {source}, got {type(voice).__name__}")
    values = voice.detach().cpu().numpy().astype("<f4", copy=False)
    if values.size != 510 * 256:
        raise RuntimeError(f"Unexpected voice shape {values.shape} in {source}")
    values = values.reshape(510, 256)
    if not np.isfinite(values).all():
        raise RuntimeError(f"Non-finite values in {source}")
    destination.parent.mkdir(parents=True, exist_ok=True)
    values.tofile(destination)


def voice_display_name(voice_id: str) -> str:
    name = voice_id.split("_", 1)[1]
    if name.isdigit():
        gender = "Female" if voice_id[1] == "f" else "Male"
        return f"Chinese {gender} {name}"
    return name.replace("_", " ").title()


def list_voice_files(api: HfApi, repository: str, revision: str) -> list[Any]:
    entries = [
        entry
        for entry in api.list_repo_tree(
            repository,
            revision=revision,
            recursive=True,
            expand=True,
            token=False,
        )
        if entry.path.startswith("voices/") and entry.path.endswith(".pt")
    ]
    expected = EXPECTED_VOICE_COUNTS[repository]
    if len(entries) != expected:
        raise RuntimeError(
            f"Expected {expected} voices in {repository}@{revision}, found {len(entries)}"
        )
    if any(entry.lfs is None for entry in entries):
        raise RuntimeError(f"A voice in {repository}@{revision} has no LFS metadata")
    return sorted(entries, key=lambda entry: entry.path)


def sync_voices(
    cache_dir: Path, workers: int
) -> tuple[list[dict[str, Any]], dict[str, str], list[dict[str, Any]]]:
    api = HfApi(token=False)
    sources = [
        (BASE_REPOSITORY, BASE_REVISION, "kokoro-v1.0-fp32"),
        (ZH_REPOSITORY, ZH_REVISION, "kokoro-v1.1-zh-fp32"),
    ]
    jobs: list[tuple[str, str, str, Any]] = []
    for repository, revision, model_id in sources:
        for entry in list_voice_files(api, repository, revision):
            jobs.append((repository, revision, model_id, entry))

    downloaded: dict[tuple[str, str], Path] = {}
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(
                download_hf_file,
                repository,
                revision,
                entry.path,
                cache_dir,
            ): (repository, entry.path)
            for repository, revision, _model_id, entry in jobs
        }
        for index, future in enumerate(as_completed(futures), start=1):
            repository, path = futures[future]
            downloaded[(repository, path)] = future.result()
            if index % 20 == 0 or index == len(futures):
                print(f"Downloaded {index}/{len(futures)} voice tensors")

    inventory = []
    for repository, revision, model_id, entry in jobs:
        source = downloaded[(repository, entry.path)]
        if source.stat().st_size != entry.lfs.size or sha256(source) != entry.lfs.sha256:
            raise RuntimeError(f"Source LFS verification failed: {repository}/{entry.path}")
        voice_id = Path(entry.path).stem
        language_id = LANGUAGE_BY_PREFIX[voice_id[0]]
        destination = ROOT / "voices" / language_id / f"{voice_id}.bin"
        convert_voice(source, destination)
        inventory.append(
            {
                "id": voice_id,
                "languageId": language_id,
                "displayName": voice_display_name(voice_id),
                "gender": "female" if voice_id[1] == "f" else "male",
                "modelId": model_id,
                "sourceRepository": repository,
                "sourceRevision": revision,
                "sourcePath": entry.path,
                "sourceSizeBytes": entry.lfs.size,
                "sourceSha256": entry.lfs.sha256,
                "artifactPath": destination.relative_to(ROOT).as_posix(),
                "artifactSizeBytes": destination.stat().st_size,
                "artifactSha256": sha256(destination),
            }
        )

    checkpoints: dict[str, str] = {}
    model_files = [
        (
            BASE_REPOSITORY,
            BASE_REVISION,
            "kokoro-v1_0.pth",
            "runtime/kokoro-v1.0-config.json",
            "kokoro-v1.0-fp32",
        ),
        (
            ZH_REPOSITORY,
            ZH_REVISION,
            "kokoro-v1_1-zh.pth",
            "runtime/kokoro-v1.1-zh-config.json",
            "kokoro-v1.1-zh-fp32",
        ),
    ]
    for repository, revision, checkpoint_name, config_target, model_id in model_files:
        checkpoint = download_hf_file(
            repository, revision, checkpoint_name, cache_dir
        )
        config = download_hf_file(repository, revision, "config.json", cache_dir)
        if sha256(checkpoint) != MODEL_FILE_SHA256[(repository, checkpoint_name)]:
            raise RuntimeError(f"Checkpoint verification failed: {repository}/{checkpoint_name}")
        if sha256(config) != MODEL_FILE_SHA256[(repository, "config.json")]:
            raise RuntimeError(f"Config verification failed: {repository}/config.json")
        target = ROOT / config_target
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(config, target)
        checkpoints[model_id] = str(checkpoint.resolve())

    language_model_inventory = []
    for source_path, (model_id, expected_sha) in LANGUAGE_MODEL_FILES.items():
        checkpoint = download_hf_file(
            LANGUAGE_MODELS_REPOSITORY,
            LANGUAGE_MODELS_REVISION,
            source_path,
            cache_dir,
        )
        if sha256(checkpoint) != expected_sha:
            raise RuntimeError(f"Checkpoint verification failed: {source_path}")
        checkpoints[model_id] = str(checkpoint.resolve())
        language_model_inventory.append(
            {
                "role": "checkpoint",
                "modelId": model_id,
                "sourcePath": source_path,
                "sourceSizeBytes": checkpoint.stat().st_size,
                "sourceSha256": expected_sha,
            }
        )

    for source_path, (target_path, expected_sha) in LANGUAGE_VOICE_FILES.items():
        source = download_hf_file(
            LANGUAGE_MODELS_REPOSITORY,
            LANGUAGE_MODELS_REVISION,
            source_path,
            cache_dir,
        )
        if sha256(source) != expected_sha:
            raise RuntimeError(f"Voice verification failed: {source_path}")
        destination = ROOT / target_path
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(source, destination)
        language_model_inventory.append(
            {
                "role": "voice",
                "sourcePath": source_path,
                "sourceSizeBytes": source.stat().st_size,
                "sourceSha256": expected_sha,
                "artifactPath": target_path,
                "artifactSizeBytes": destination.stat().st_size,
                "artifactSha256": sha256(destination),
            }
        )

    expected_voice_paths = {item["artifactPath"] for item in inventory}
    expected_voice_paths.update(target for target, _sha in LANGUAGE_VOICE_FILES.values())
    actual_voice_paths = {
        path.relative_to(ROOT).as_posix()
        for path in (ROOT / "voices").glob("*/*.bin")
    }
    if actual_voice_paths != expected_voice_paths:
        raise RuntimeError(
            "Voice directory does not exactly match the pinned inventory: "
            f"missing={sorted(expected_voice_paths - actual_voice_paths)}, "
            f"unexpected={sorted(actual_voice_paths - expected_voice_paths)}"
        )

    return inventory, checkpoints, language_model_inventory


def download_verified(url: str, destination: Path, expected_sha256: str, size: int) -> None:
    if destination.is_file() and destination.stat().st_size == size:
        if sha256(destination) == expected_sha256:
            return
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(destination.suffix + ".download")
    with requests.get(url, stream=True, timeout=120) as response:
        response.raise_for_status()
        with temporary.open("wb") as output:
            for block in response.iter_content(1024 * 1024):
                if block:
                    output.write(block)
    if temporary.stat().st_size != size or sha256(temporary) != expected_sha256:
        raise RuntimeError(f"Downloaded file verification failed: {url}")
    temporary.replace(destination)


def sync_phonemis() -> list[dict[str, Any]]:
    inventory = []
    for source_path, (target_path, expected_sha, expected_size) in PHONEMIS_FILES.items():
        url = (
            "https://media.githubusercontent.com/media/"
            f"{PHONEMIS_REPOSITORY}/{PHONEMIS_REVISION}/{source_path}"
        )
        destination = ROOT / target_path
        download_verified(url, destination, expected_sha, expected_size)
        inventory.append(
            {
                "sourcePath": source_path,
                "sourceSizeBytes": expected_size,
                "sourceSha256": expected_sha,
                "artifactPath": target_path,
            }
        )
    return inventory


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=Path(tempfile.gettempdir()) / "kokoro-kmp-hf-cache",
    )
    parser.add_argument("--workers", type=int, default=8)
    args = parser.parse_args()
    if not 1 <= args.workers <= 16:
        raise SystemExit("--workers must be in 1..16")

    voices, checkpoints, language_models = sync_voices(args.cache_dir, args.workers)
    phonemis = sync_phonemis()
    manifest = {
        "sourceRepositories": {
            "kokoroBase": {
                "repository": BASE_REPOSITORY,
                "revision": BASE_REVISION,
            },
            "kokoroChinese": {
                "repository": ZH_REPOSITORY,
                "revision": ZH_REVISION,
            },
            "languageModels": {
                "repository": LANGUAGE_MODELS_REPOSITORY,
                "revision": LANGUAGE_MODELS_REVISION,
            },
            "phonemis": {
                "repository": PHONEMIS_REPOSITORY,
                "revision": PHONEMIS_REVISION,
            },
        },
        "voiceCount": len(voices),
        "catalogVoiceCount": len(voices) + len(LANGUAGE_VOICE_FILES),
        "voices": voices,
        "languageModels": language_models,
        "phonemis": phonemis,
    }
    source_dir = ROOT / "sources"
    source_dir.mkdir(parents=True, exist_ok=True)
    (source_dir / "upstream-assets.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    checkpoint_manifest = Path(tempfile.gettempdir()) / "kokoro-kmp-checkpoints.json"
    checkpoint_manifest.write_text(
        json.dumps(checkpoints, indent=2) + "\n", encoding="utf-8"
    )
    print(f"Converted {len(voices)} upstream voices")
    print(f"Checkpoint paths: {checkpoint_manifest}")


if __name__ == "__main__":
    main()