"""Lazy, offline NeMo adapter for an explicit local Parakeet archive."""

from __future__ import annotations

from copy import deepcopy
import hashlib
import os
from pathlib import Path
from typing import Protocol, Sequence

MAX_TERMS_FILE_BYTES = 64 * 1024
MAX_CODING_TERMS = 512
MAX_TERM_BYTES = 128
PINNED_MODEL_SHA256 = "3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d"


class TranscriptionModel(Protocol):
    def load(self) -> None: ...
    def transcribe(self, samples: Sequence[float]) -> str: ...


def load_coding_terms(path: str | os.PathLike[str]) -> list[str]:
    """Load a bounded host-owned phrase list, ignoring comments and duplicates."""

    terms_path = Path(path).expanduser()
    if not terms_path.is_file():
        raise ValueError("PI_DICTATE_TERMS_PATH must name an existing local dictionary")
    if terms_path.stat().st_size > MAX_TERMS_FILE_BYTES:
        raise ValueError(f"coding dictionary exceeds {MAX_TERMS_FILE_BYTES} bytes")

    terms: list[str] = []
    seen: set[str] = set()
    for raw_line in terms_path.read_text(encoding="utf-8").splitlines():
        term = raw_line.strip()
        if not term or term.startswith("#"):
            continue
        if len(term.encode("utf-8")) > MAX_TERM_BYTES:
            raise ValueError(f"coding dictionary terms may not exceed {MAX_TERM_BYTES} UTF-8 bytes")
        if term not in seen:
            terms.append(term)
            seen.add(term)
        if len(terms) > MAX_CODING_TERMS:
            raise ValueError(f"coding dictionary may not exceed {MAX_CODING_TERMS} terms")
    return terms


class NemoParakeetModel:
    """Load NeMo only when requested so core tests require neither CUDA nor NeMo."""

    def __init__(
        self,
        model_path: str | os.PathLike[str],
        *,
        terms_path: str | os.PathLike[str] | None = None,
        boost_alpha: float = 1.0,
        expected_sha256: str = PINNED_MODEL_SHA256,
    ) -> None:
        path = Path(model_path).expanduser()
        if path.suffix != ".nemo" or not path.is_file():
            raise ValueError("PI_DICTATE_MODEL_PATH must name an existing local .nemo file")
        if boost_alpha <= 0:
            raise ValueError("PI_DICTATE_BOOST_ALPHA must be positive")
        self._model_path = path.resolve()
        self._terms = load_coding_terms(terms_path) if terms_path is not None else []
        self._boost_alpha = boost_alpha
        self._expected_sha256 = expected_sha256
        self._model: object | None = None
        self._numpy: object | None = None

    @property
    def loaded(self) -> bool:
        return self._model is not None

    def load(self) -> None:
        if self._model is not None:
            return
        with self._model_path.open("rb") as model_file:
            actual_sha256 = hashlib.file_digest(model_file, "sha256").hexdigest()
        if actual_sha256 != self._expected_sha256:
            raise ValueError(
                "model checksum mismatch: expected the pinned Parakeet revision"
            )
        os.environ["HF_HUB_OFFLINE"] = "1"
        os.environ["TRANSFORMERS_OFFLINE"] = "1"
        os.environ["HF_DATASETS_OFFLINE"] = "1"
        os.environ["NEMO_OFFLINE"] = "1"

        import numpy
        import torch
        from nemo.collections.asr.models import EncDecRNNTBPEModel
        from omegaconf import open_dict

        model = EncDecRNNTBPEModel.restore_from(
            restore_path=str(self._model_path),
            map_location=torch.device("cuda"),
        )
        model.to("cuda")
        model.eval()
        if self._terms:
            decoding = deepcopy(model.cfg.decoding)
            with open_dict(decoding):
                greedy = decoding["greedy"]
                greedy["boosting_tree"] = {
                    "key_phrases_list": list(self._terms),
                    "context_score": 1.0,
                    "depth_scaling": 2.0,
                    "unk_score": 0.0,
                    "final_eos_score": 1.0,
                    "score_per_phrase": 0.0,
                    "source_lang": "en",
                    "use_triton": True,
                    "uniform_weights": False,
                    "use_bpe_dropout": False,
                }
                greedy["boosting_tree_alpha"] = self._boost_alpha
            model.change_decoding_strategy(decoding, verbose=False)
        self._numpy = numpy
        self._model = model

    def transcribe(self, samples: Sequence[float]) -> str:
        if self._model is None or self._numpy is None:
            raise RuntimeError("model is not loaded")
        numpy = self._numpy
        audio = numpy.asarray(samples, dtype=numpy.float32)  # type: ignore[attr-defined]
        output: Sequence[object] = self._model.transcribe(  # type: ignore[attr-defined]
            audio=[audio],
            batch_size=1,
            num_workers=0,
            verbose=False,
        )
        if not output:
            return ""
        hypothesis = output[0]
        if isinstance(hypothesis, str):
            return hypothesis
        text = getattr(hypothesis, "text", None)
        if isinstance(text, str):
            return text
        raise RuntimeError("NeMo returned an unsupported transcription result")