from __future__ import annotations

import importlib.util
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .models import ErrorCode, ProviderError

DEFAULT_VAD_URL = "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/silero_vad.onnx"
DEFAULT_STT_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
    "sherpa-onnx-streaming-zipformer-ctc-zh-int8-2025-06-30.tar.bz2"
)
DEFAULT_STT_DIR = "sherpa-onnx-streaming-zipformer-ctc-zh-int8-2025-06-30"
DEFAULT_KWS_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/kws-models/"
    "sherpa-onnx-kws-zipformer-wenetspeech-3.3M-2024-01-01-mobile.tar.bz2"
)
DEFAULT_KWS_DIR = "sherpa-onnx-kws-zipformer-wenetspeech-3.3M-2024-01-01-mobile"
DEFAULT_KWS_KEYWORDS = "x iǎo j ié x iǎo j ié @小杰小杰\n"
DEFAULT_DENOISER_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/"
    "speech-enhancement-models/gtcrn_simple.onnx"
)
DEFAULT_DENOISER_PATH = "denoise/gtcrn_simple.onnx"

REQUIRED_MODEL_FILES = (
    f"wake/{DEFAULT_KWS_DIR}/tokens.txt",
    f"wake/{DEFAULT_KWS_DIR}/encoder-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
    f"wake/{DEFAULT_KWS_DIR}/decoder-epoch-12-avg-2-chunk-16-left-64.onnx",
    f"wake/{DEFAULT_KWS_DIR}/joiner-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
    "wake/keywords.txt",
    "vad/silero_vad.onnx",
    f"stt/{DEFAULT_STT_DIR}/tokens.txt",
    f"stt/{DEFAULT_STT_DIR}/model.int8.onnx",
    DEFAULT_DENOISER_PATH,
)


@dataclass(frozen=True, slots=True)
class SpeechModelStatus:
    models_dir: Path
    manifest_path: Path
    manifest_present: bool
    required_files: tuple[str, ...]
    missing_files: tuple[str, ...]
    sherpa_onnx_available: bool

    @property
    def ok(self) -> bool:
        return not self.missing_files and self.sherpa_onnx_available

    def to_json(self) -> dict[str, Any]:
        return {
            "ok": self.ok,
            "models_dir": str(self.models_dir),
            "manifest_path": str(self.manifest_path),
            "manifest_present": self.manifest_present,
            "required_files": list(self.required_files),
            "missing_files": list(self.missing_files),
            "sherpa_onnx_available": self.sherpa_onnx_available,
        }


def default_manifest() -> dict[str, Any]:
    return {
        "schema": "owner_voice_pet.speech_models.v1",
        "sample_rate": 16000,
        "sources": {
            "wake": DEFAULT_KWS_URL,
            "vad": DEFAULT_VAD_URL,
            "stt": DEFAULT_STT_URL,
            "denoiser": DEFAULT_DENOISER_URL,
        },
        "providers": {
            "wake": {
                "type": "sherpa-onnx-keyword-spotter",
                "model_dir": f"wake/{DEFAULT_KWS_DIR}",
                "tokens": f"wake/{DEFAULT_KWS_DIR}/tokens.txt",
                "encoder": f"wake/{DEFAULT_KWS_DIR}/encoder-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
                "decoder": f"wake/{DEFAULT_KWS_DIR}/decoder-epoch-12-avg-2-chunk-16-left-64.onnx",
                "joiner": f"wake/{DEFAULT_KWS_DIR}/joiner-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
                "keywords": "wake/keywords.txt",
            },
            "vad": {
                "type": "silero-vad",
                "path": "vad/silero_vad.onnx",
            },
            "stt": {
                "type": "sherpa-onnx-streaming-zipformer2-ctc",
                "model_dir": f"stt/{DEFAULT_STT_DIR}",
                "tokens": f"stt/{DEFAULT_STT_DIR}/tokens.txt",
                "model": f"stt/{DEFAULT_STT_DIR}/model.int8.onnx",
            },
            "denoiser": {
                "type": "sherpa-onnx-gtcrn",
                "path": DEFAULT_DENOISER_PATH,
            },
        },
        "required_files": list(REQUIRED_MODEL_FILES),
    }


def write_default_manifest(models_dir: str | Path) -> Path:
    root = Path(models_dir)
    root.mkdir(parents=True, exist_ok=True)
    manifest_path = root / "manifest.json"
    manifest_path.write_text(
        json.dumps(default_manifest(), ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    return manifest_path


def load_manifest(models_dir: str | Path) -> dict[str, Any]:
    manifest_path = Path(models_dir) / "manifest.json"
    if not manifest_path.exists():
        return default_manifest()
    return json.loads(manifest_path.read_text(encoding="utf-8"))


def required_model_files(models_dir: str | Path) -> tuple[str, ...]:
    manifest = load_manifest(models_dir)
    files = manifest.get("required_files", list(REQUIRED_MODEL_FILES))
    return tuple(str(item) for item in files)


def vad_model_path(models_dir: str | Path) -> Path:
    root = Path(models_dir)
    manifest = load_manifest(root)
    path = manifest.get("providers", {}).get("vad", {}).get("path", "vad/silero_vad.onnx")
    return root / str(path)


def denoiser_model_path(models_dir: str | Path) -> Path:
    root = Path(models_dir)
    manifest = load_manifest(root)
    path = manifest.get("providers", {}).get("denoiser", {}).get("path", DEFAULT_DENOISER_PATH)
    return root / str(path)


def wake_model_paths(models_dir: str | Path) -> dict[str, Path]:
    root = Path(models_dir)
    manifest = load_manifest(root)
    wake = manifest.get("providers", {}).get("wake", {})
    return {
        "model_dir": root / str(wake.get("model_dir", f"wake/{DEFAULT_KWS_DIR}")),
        "tokens": root / str(wake.get("tokens", f"wake/{DEFAULT_KWS_DIR}/tokens.txt")),
        "encoder": root
        / str(wake.get("encoder", f"wake/{DEFAULT_KWS_DIR}/encoder-epoch-12-avg-2-chunk-16-left-64.int8.onnx")),
        "decoder": root
        / str(wake.get("decoder", f"wake/{DEFAULT_KWS_DIR}/decoder-epoch-12-avg-2-chunk-16-left-64.onnx")),
        "joiner": root
        / str(wake.get("joiner", f"wake/{DEFAULT_KWS_DIR}/joiner-epoch-12-avg-2-chunk-16-left-64.int8.onnx")),
        "keywords": root / str(wake.get("keywords", "wake/keywords.txt")),
    }


def stt_model_paths(model_path: str | Path) -> dict[str, Any]:
    root = Path(model_path)
    if (root / "manifest.json").exists() or (root / "stt").exists():
        manifest = load_manifest(root)
        stt = manifest.get("providers", {}).get("stt", {})
        stt_type = str(stt.get("type", "sherpa-onnx-streaming-transducer"))
        if stt_type == "sherpa-onnx-streaming-zipformer2-ctc":
            return {
                "type": stt_type,
                "model_dir": root / str(stt.get("model_dir", f"stt/{DEFAULT_STT_DIR}")),
                "tokens": root / str(stt.get("tokens", f"stt/{DEFAULT_STT_DIR}/tokens.txt")),
                "model": root / str(stt.get("model", f"stt/{DEFAULT_STT_DIR}/model.int8.onnx")),
            }
        return {
            "type": stt_type,
            "model_dir": root / str(stt.get("model_dir", f"stt/{DEFAULT_STT_DIR}")),
            "tokens": root / str(stt.get("tokens", f"stt/{DEFAULT_STT_DIR}/tokens.txt")),
            "encoder": root / str(stt.get("encoder", f"stt/{DEFAULT_STT_DIR}/encoder-epoch-99-avg-1.int8.onnx")),
            "decoder": root / str(stt.get("decoder", f"stt/{DEFAULT_STT_DIR}/decoder-epoch-99-avg-1.onnx")),
            "joiner": root / str(stt.get("joiner", f"stt/{DEFAULT_STT_DIR}/joiner-epoch-99-avg-1.int8.onnx")),
        }
    return {
        "type": "sherpa-onnx-streaming-transducer",
        "model_dir": root,
        "tokens": root / "tokens.txt",
        "encoder": root / "encoder-epoch-99-avg-1.int8.onnx",
        "decoder": root / "decoder-epoch-99-avg-1.onnx",
        "joiner": root / "joiner-epoch-99-avg-1.int8.onnx",
    }


def check_speech_models(models_dir: str | Path, require_sherpa: bool = True) -> SpeechModelStatus:
    root = Path(models_dir)
    manifest_path = root / "manifest.json"
    required = required_model_files(root)
    missing = tuple(relative for relative in required if not (root / relative).exists())
    sherpa_available = True
    if require_sherpa:
        sherpa_available = importlib.util.find_spec("sherpa_onnx") is not None
    return SpeechModelStatus(
        models_dir=root,
        manifest_path=manifest_path,
        manifest_present=manifest_path.exists(),
        required_files=required,
        missing_files=missing,
        sherpa_onnx_available=sherpa_available,
    )


def model_status_errors(status: SpeechModelStatus) -> list[ProviderError]:
    errors: list[ProviderError] = []
    if not status.sherpa_onnx_available:
        errors.append(
            ProviderError(
                ErrorCode.STT_TRANSCRIBE_FAILED,
                "sherpa_onnx is not installed in the active Python environment",
                False,
                "sherpa-onnx",
                "model-check",
            )
        )
    denoise_missing = tuple(item for item in status.missing_files if item.startswith("denoise/"))
    speech_missing = tuple(item for item in status.missing_files if not item.startswith("denoise/"))
    if speech_missing:
        errors.append(
            ProviderError(
                ErrorCode.STT_MODEL_MISSING,
                "missing speech model files: " + ", ".join(speech_missing),
                False,
                "speech-models",
                "model-check",
            )
        )
    if denoise_missing:
        errors.append(
            ProviderError(
                ErrorCode.NOISE_FILTER_MODEL_MISSING,
                "missing denoiser model files: " + ", ".join(denoise_missing),
                False,
                "speech-models",
                "model-check",
            )
        )
    return errors