#!/usr/bin/env python3
"""Validate catalog integrity and exercise every model and every voice on CPU."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any, Iterator

import numpy as np
import onnx
import onnxruntime as ort
import soundfile as sf
from jsonschema import Draft202012Validator


ROOT = Path(__file__).resolve().parents[1]
VOICE_SHAPE = (510, 256)


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 artifacts(value: Any) -> Iterator[dict[str, Any]]:
    if isinstance(value, list):
        for item in value:
            yield from artifacts(item)
    elif isinstance(value, dict):
        if {"path", "sizeBytes", "sha256"} <= value.keys():
            yield value
        for item in value.values():
            yield from artifacts(item)


def unique_by_id(values: list[dict[str, Any]], description: str) -> dict[str, dict[str, Any]]:
    result = {value["id"]: value for value in values}
    if len(result) != len(values):
        raise RuntimeError(f"Duplicate {description} id")
    return result


catalog = json.loads((ROOT / "catalog.json").read_text(encoding="utf-8"))
source_lock = json.loads((ROOT / "SOURCES.lock.json").read_text(encoding="utf-8"))
schema = json.loads(
    (ROOT / "schema/kokoro-catalog.schema.json").read_text(encoding="utf-8")
)
Draft202012Validator.check_schema(schema)
Draft202012Validator(schema).validate(catalog)

artifact_by_path: dict[str, dict[str, Any]] = {}
for artifact in artifacts(catalog):
    previous = artifact_by_path.setdefault(artifact["path"], artifact)
    if previous != artifact:
        raise RuntimeError(f"Conflicting metadata for artifact {artifact['path']}")
for artifact in artifact_by_path.values():
    path = ROOT / artifact["path"]
    if not path.is_file():
        raise FileNotFoundError(path)
    if path.stat().st_size != artifact["sizeBytes"]:
        raise RuntimeError(f"Size mismatch: {path}")
    if sha256(path) != artifact["sha256"]:
        raise RuntimeError(f"SHA-256 mismatch: {path}")

tokenizers = unique_by_id(catalog["tokenizers"], "tokenizer")
models = unique_by_id(catalog["models"], "model")
languages = unique_by_id(catalog["languages"], "language")
locked_model_sources = {
    entry["repository"]: entry
    for key in ("kokoroBase", "kokoroChinese", "languageModels")
    for entry in [source_lock[key]]
}
configs: dict[str, dict[str, Any]] = {}
for tokenizer_id, tokenizer in tokenizers.items():
    configs[tokenizer_id] = json.loads(
        (ROOT / tokenizer["artifact"]["path"]).read_text(encoding="utf-8")
    )
for model in models.values():
    if model["tokenizerId"] not in tokenizers:
        raise RuntimeError(
            f"{model['id']}: unknown tokenizer {model['tokenizerId']}"
        )
    graph = onnx.load(ROOT / model["artifact"]["path"], load_external_data=False)
    metadata = {item.key: item.value for item in graph.metadata_props}
    if metadata.get("license") != "Apache-2.0":
        raise RuntimeError(f"{model['id']}: ONNX license metadata is not Apache-2.0")
    if metadata.get("source_repository") != model["source"]["repository"]:
        raise RuntimeError(f"{model['id']}: source repository metadata mismatch")
    if metadata.get("source_revision") != model["source"]["revision"]:
        raise RuntimeError(f"{model['id']}: source revision metadata mismatch")
    if metadata.get("source_checkpoint") != model["source"]["checkpoint"]:
        raise RuntimeError(f"{model['id']}: source checkpoint metadata mismatch")
    locked_source = locked_model_sources[model["source"]["repository"]]
    expected_checkpoint_sha = locked_source["files"][model["source"]["checkpoint"]]
    if metadata.get("source_checkpoint_sha256") != expected_checkpoint_sha:
        raise RuntimeError(f"{model['id']}: source checkpoint checksum metadata mismatch")
    validation = json.loads(
        (ROOT / model["validation"]["path"]).read_text(encoding="utf-8")
    )
    if validation["modelSha256"] != model["artifact"]["sha256"]:
        raise RuntimeError(f"{model['id']}: validation model checksum mismatch")

voices: dict[tuple[str, str], dict[str, Any]] = {}
voice_pack_cache: dict[str, np.ndarray] = {}
for language in languages.values():
    language_voices = unique_by_id(language["voices"], f"voice in {language['id']}")
    if language["defaultVoiceId"] not in language_voices:
        raise RuntimeError(f"{language['id']}: unknown default voice")
    if language["textFrontend"]["status"] == "bundled":
        roles = [item["role"] for item in language["textFrontend"]["artifacts"]]
        if roles.count("neural-model") != 1 or len(roles) != len(set(roles)):
            raise RuntimeError(f"{language['id']}: invalid frontend artifact roles")
    for voice_id, voice in language_voices.items():
        if voice["modelId"] not in models:
            raise RuntimeError(f"{language['id']}/{voice_id}: unknown model")
        path = ROOT / voice["artifact"]["path"]
        if path.stat().st_size != VOICE_SHAPE[0] * VOICE_SHAPE[1] * 4:
            raise RuntimeError(f"{language['id']}/{voice_id}: invalid voice byte size")
        voice_pack = np.fromfile(path, dtype="<f4").reshape(VOICE_SHAPE)
        if not np.isfinite(voice_pack).all():
            raise RuntimeError(f"{language['id']}/{voice_id}: non-finite voice values")
        voice_pack_cache[voice["artifact"]["path"]] = voice_pack
        voices[(language["id"], voice_id)] = voice

source_inventory = json.loads(
    (ROOT / "sources/upstream-assets.json").read_text(encoding="utf-8")
)
locked_manifest = source_lock["upstreamAssetManifest"]
if (
    locked_manifest["path"] != "sources/upstream-assets.json"
    or locked_manifest["sha256"] != sha256(ROOT / locked_manifest["path"])
):
    raise RuntimeError("SOURCES.lock.json does not match the upstream asset manifest")
if source_inventory["voiceCount"] != 157 or source_inventory["catalogVoiceCount"] != 159:
    raise RuntimeError("Pinned upstream inventory must contain exactly 157 voices")

phonemis_onnx_lock = source_lock["phonemisOnnx"]
phonemis_onnx_manifest_path = ROOT / phonemis_onnx_lock["manifestPath"]
if phonemis_onnx_lock["manifestSha256"] != sha256(phonemis_onnx_manifest_path):
    raise RuntimeError("SOURCES.lock.json does not match the Phonemis ONNX manifest")
phonemis_onnx_manifest = json.loads(
    phonemis_onnx_manifest_path.read_text(encoding="utf-8")
)
if len(phonemis_onnx_manifest["artifacts"]) != phonemis_onnx_lock["artifactCount"]:
    raise RuntimeError("Phonemis ONNX manifest artifact count mismatch")
phonemis_sources = {
    item["artifactPath"]: item for item in source_inventory["phonemis"]
}
for converted in phonemis_onnx_manifest["artifacts"]:
    source = phonemis_sources.get(converted["sourcePath"])
    if source is None or source["sourceSha256"] != converted["sourceSha256"]:
        raise RuntimeError(
            f"Phonemis ONNX source mismatch for {converted['sourcePath']}"
        )
    artifact = artifact_by_path.get(converted["path"])
    if artifact is None or any(
        artifact[field] != converted[field]
        for field in ("path", "sizeBytes", "sha256")
    ):
        raise RuntimeError(
            f"Phonemis ONNX catalog mismatch for {converted['path']}"
        )
for source_voice in source_inventory["voices"]:
    key = (source_voice["languageId"], source_voice["id"])
    if key not in voices:
        raise RuntimeError(f"Source inventory voice missing from catalog: {key}")
    artifact = voices[key]["artifact"]
    if (
        artifact["path"] != source_voice["artifactPath"]
        or artifact["sizeBytes"] != source_voice["artifactSizeBytes"]
        or artifact["sha256"] != source_voice["artifactSha256"]
        or voices[key]["modelId"] != source_voice["modelId"]
    ):
        raise RuntimeError(f"Source inventory mismatch for voice {key}")
for language_asset in source_inventory["languageModels"]:
    if language_asset["role"] != "voice":
        continue
    artifact = artifact_by_path[language_asset["artifactPath"]]
    if (
        artifact["sizeBytes"] != language_asset["artifactSizeBytes"]
        or artifact["sha256"] != language_asset["artifactSha256"]
    ):
        raise RuntimeError(
            f"Language model inventory mismatch for {language_asset['artifactPath']}"
        )
if len(voices) != 159:
    raise RuntimeError(f"Expected 159 catalog voices, found {len(voices)}")

frontend_validation = json.loads(
    (ROOT / "validation/frontend-validation.json").read_text(encoding="utf-8")
)
frontend_cases = [
    *frontend_validation["phonemis"]["cases"],
    *frontend_validation["misaki"]["cases"],
]
bundled_languages = {
    language["textFrontend"]["language"]: language
    for language in languages.values()
    if language["textFrontend"]["status"] == "bundled"
}
frontend_onnx_results = []
frontend_session_options = ort.SessionOptions()
frontend_session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
for language_id, language in bundled_languages.items():
    neural_model = next(
        item["artifact"]
        for item in language["textFrontend"]["artifacts"]
        if item["role"] == "neural-model"
    )
    if neural_model.get("mediaType") != "application/onnx":
        raise RuntimeError(f"{language_id}: neural frontend is not a prebuilt ONNX model")
    frontend_session = ort.InferenceSession(
        str(ROOT / neural_model["path"]),
        sess_options=frontend_session_options,
        providers=["CPUExecutionProvider"],
    )
    logits = frontend_session.run(
        ["logits"],
        {"input_ids": np.asarray([[2]], dtype=np.int64)},
    )[0]
    if logits.ndim != 3 or logits.shape[0] != 1 or logits.shape[2] != 116:
        raise RuntimeError(f"{language_id}: unexpected frontend logits shape {logits.shape}")
    if not np.isfinite(logits).all():
        raise RuntimeError(f"{language_id}: non-finite frontend logits")
    frontend_onnx_results.append(
        {
            "language": language_id,
            "path": neural_model["path"],
            "sha256": neural_model["sha256"],
            "logitsShape": list(logits.shape),
        }
    )
phonemis_languages = {
    case["language"] for case in frontend_validation["phonemis"]["cases"]
}
if phonemis_languages != set(bundled_languages):
    raise RuntimeError("Frontend validation does not cover every bundled Phonemis language")
for case in frontend_cases:
    if "tokenizerId" in case:
        tokenizer_ids = {case["tokenizerId"]}
    else:
        language = bundled_languages[case["language"]]
        tokenizer_ids = {
            models[voice["modelId"]]["tokenizerId"] for voice in language["voices"]
        }
    for tokenizer_id in tokenizer_ids:
        vocab = configs[tokenizer_id][catalog["runtime"]["tokenEncoding"]["vocabularyField"]]
        missing = sorted(set(case["phonemes"]) - set(vocab))
        if missing or not case["tokenizerCompatible"]:
            raise RuntimeError(
                f"Frontend case {case['language']} is incompatible with {tokenizer_id}: {missing}"
            )

session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sessions: dict[str, ort.InferenceSession] = {}


def session_for(model_id: str) -> ort.InferenceSession:
    if model_id not in sessions:
        sessions[model_id] = ort.InferenceSession(
            str(ROOT / models[model_id]["artifact"]["path"]),
            sess_options=session_options,
            providers=["CPUExecutionProvider"],
        )
    return sessions[model_id]


def infer(
    language_id: str,
    voice_id: str,
    phonemes: str,
) -> tuple[np.ndarray, np.ndarray, int]:
    voice = voices[(language_id, voice_id)]
    model = models[voice["modelId"]]
    config = configs[model["tokenizerId"]]
    vocab = config[catalog["runtime"]["tokenEncoding"]["vocabularyField"]]
    missing = sorted({phoneme for phoneme in phonemes if phoneme not in vocab})
    if missing:
        raise RuntimeError(
            f"{language_id}/{voice_id}: phonemes missing from {model['tokenizerId']}: {missing}"
        )
    if not 1 <= len(phonemes) <= catalog["runtime"]["maxPhonemeCodePoints"]:
        raise RuntimeError(f"{language_id}/{voice_id}: invalid phoneme length")
    token_ids = [vocab[phoneme] for phoneme in phonemes]
    input_ids = np.asarray([[0, *token_ids, 0]], dtype=np.int64)
    voice_pack = voice_pack_cache[voice["artifact"]["path"]]
    waveform, duration = session_for(model["id"]).run(
        None,
        {
            "input_ids": input_ids,
            "style": voice_pack[len(phonemes) - 1].reshape(1, 256),
            "speed": np.asarray([1.0], dtype=np.float32),
        },
    )
    if not np.isfinite(waveform).all():
        raise RuntimeError(f"{language_id}/{voice_id}: non-finite waveform")
    if duration.shape != (input_ids.shape[1],):
        raise RuntimeError(
            f"{language_id}/{voice_id}: unexpected duration shape {duration.shape}"
        )
    expected_samples = int(duration.sum()) * catalog["runtime"]["durationToSamplesFactor"]
    if waveform.shape != (expected_samples,):
        raise RuntimeError(
            f"{language_id}/{voice_id}: expected {expected_samples} samples, got {waveform.shape}"
        )
    return waveform, duration, int(input_ids.shape[1])


cases = json.loads((ROOT / "validation/cases.json").read_text(encoding="utf-8"))
case_results = []
for case in cases:
    if (case["language"], case["voice"]) not in voices:
        raise RuntimeError(f"{case['id']}: unknown language or voice")
    waveform, _duration, input_tokens = infer(
        case["language"], case["voice"], case["phonemes"]
    )
    if "previewPath" in case:
        preview_path = ROOT / case["previewPath"]
        preview_path.parent.mkdir(parents=True, exist_ok=True)
        sf.write(preview_path, waveform, catalog["audio"]["sampleRateHz"], subtype="PCM_16")
    case_results.append(
        {
            "id": case["id"],
            "language": case["language"],
            "voice": case["voice"],
            "modelId": voices[(case["language"], case["voice"])]["modelId"],
            "phonemeCodePoints": len(case["phonemes"]),
            "inputTokens": input_tokens,
            "outputSamples": int(waveform.size),
            "durationSeconds": waveform.size / catalog["audio"]["sampleRateHz"],
            "peakAbsoluteAmplitude": float(np.abs(waveform).max(initial=0.0)),
        }
    )

voice_results = []
for language in languages.values():
    for voice in language["voices"]:
        waveform, _duration, input_tokens = infer(language["id"], voice["id"], "a")
        voice_results.append(
            {
                "language": language["id"],
                "voice": voice["id"],
                "modelId": voice["modelId"],
                "inputTokens": input_tokens,
                "outputSamples": int(waveform.size),
                "peakAbsoluteAmplitude": float(np.abs(waveform).max(initial=0.0)),
            }
        )

boundary_results = []
for model_id, model in models.items():
    compatible = next(
        (language_id, voice_id)
        for (language_id, voice_id), voice in voices.items()
        if voice["modelId"] == model_id
    )
    config = configs[model["tokenizerId"]]
    vocab = config[catalog["runtime"]["tokenEncoding"]["vocabularyField"]]
    if "a" not in vocab:
        raise RuntimeError(f"{model_id}: tokenizer has no boundary-test token 'a'")
    for phoneme_count in (1, 128, catalog["runtime"]["maxPhonemeCodePoints"]):
        waveform, _duration, input_tokens = infer(
            compatible[0], compatible[1], "a" * phoneme_count
        )
        boundary_results.append(
            {
                "modelId": model_id,
                "phonemeCodePoints": phoneme_count,
                "inputTokens": input_tokens,
                "outputSamples": int(waveform.size),
            }
        )

report = {
    "catalogVersion": catalog["catalogVersion"],
    "catalogSchemaValidated": True,
    "onnxRuntimeVersion": ort.__version__,
    "artifactsVerified": len(artifact_by_path),
    "modelsValidated": len(models),
    "languageProfilesValidated": len(languages),
    "voicesValidated": len(voice_results),
    "frontendCasesValidated": len(frontend_cases),
    "frontendOnnxModelsValidated": len(frontend_onnx_results),
    "frontendOnnxSmokeTests": frontend_onnx_results,
    "cases": case_results,
    "voiceSmokeTests": voice_results,
    "boundaryCases": boundary_results,
}
(ROOT / "validation/package-validation.json").write_text(
    json.dumps(report, ensure_ascii=False, indent=2) + "\n",
    encoding="utf-8",
)
print(
    json.dumps(
        {
            "catalogVersion": report["catalogVersion"],
            "artifactsVerified": report["artifactsVerified"],
            "modelsValidated": report["modelsValidated"],
            "languageProfilesValidated": report["languageProfilesValidated"],
            "voicesValidated": report["voicesValidated"],
            "frontendCasesValidated": report["frontendCasesValidated"],
            "frontendOnnxModelsValidated": report["frontendOnnxModelsValidated"],
            "naturalCasesValidated": len(case_results),
            "boundaryCasesValidated": len(boundary_results),
        },
        indent=2,
    )
)