"""Model wrappers for Gemma 3 inference."""

import asyncio
import contextlib
from collections.abc import AsyncIterator, Generator, Sequence
from enum import Enum
from pathlib import Path
from typing import Protocol, cast

import numpy as np
import numpy.typing as npt
from typing_extensions import Self

from .backends import DeviceSelection, EmbeddingBackend, GenerationBackend, resolve_device_selection
from .backends import resolve_embedding_backend as _resolve_embedding_backend_impl
from .backends import resolve_generation_backend as _resolve_generation_backend_impl
from .config import EmbeddingConfig, GenerationConfig
from .hub import HubManager
from .hydration import ImageHydrator
from .loader import ModelLoader, auto_loader
from .telemetry import tracer
from .typing import SENTENCEPIECE_INSTALLED, _SPProcessorImpl

try:
    from . import _core
except ImportError:
    # Keep import optional for development/testing environments.
    _core = None


class _EncodedToken(Protocol):
    ids: Sequence[int]


class _Tokenizer:
    """Wrapper for sentencepiece to match internal tokenizer needs."""

    def __init__(self, model_path: str) -> None:
        if _SPProcessorImpl is None:
            msg = "sentencepiece runtime is unavailable"
            raise ModuleNotFoundError(msg)
        self.sp = _SPProcessorImpl(model_file=model_path)
        self._max_length: int | None = None

    def enable_truncation(self, *, max_length: int, **_: object) -> None:
        self._max_length = max_length

    def enable_padding(self, **_: object) -> None:
        pass

    def encode_batch(self, text: Sequence[str]) -> Sequence[_EncodedToken]:
        return [self.encode(t) for t in text]

    def encode(self, text: str) -> _EncodedToken:
        ids = self.sp.EncodeAsIds(text)
        if self._max_length is not None:
            ids = ids[: self._max_length]

        class _Result:
            def __init__(self, ids: list[int]) -> None:
                self.ids = ids

        return cast("_EncodedToken", _Result(ids))

    def decode(self, ids: Sequence[int]) -> str:
        return str(self.sp.DecodeIds(list(ids)))

    def token_to_id(self, token: str) -> int | None:
        token_id = self.sp.piece_to_id(token)
        if token_id == self.sp.unk_id() and token not in ("<unk>", " "):
            return None
        return int(token_id)


_EXPECTED_MATRIX_DIMS = 2
_BOS_TOKEN_ID = 2
_EOS_TOKEN_ID_ALIASES = ("<end_of_turn>", "</s>", "<eos>", "<|eos|>")
_INSTRUCTION_START = "<start_of_turn>"
_INSTRUCTION_END = "<end_of_turn>"


class ModelVariant(str, Enum):
    """Enumeration of supported model architectural variants."""

    STANDARD = "gemma_standard"
    NANO = "gemma_nano"


def _resolve_model_path(raw_model_path: str | Path, cache_path: str | Path | None = None) -> Path:
    """Resolve user-supplied model input consistently for all model types."""
    return HubManager(cache_path=cache_path).resolve_model(str(raw_model_path), download_if_missing=True, strict=True)


def _core_unavailable_message(model_type: str) -> str:
    if model_type == "embedding":
        return (
            "Mojo core is unavailable for embeddings. "
            "Build/install the `mogemma._core` extension before calling embed()/embed_tokens()."
        )
    return (
        "Mojo core is unavailable for text generation. "
        "Build/install the `mogemma._core` extension before calling generate()."
    )


def _detect_model_variant(metadata: dict[str, tuple[int, tuple[int, ...], str]]) -> ModelVariant:
    """Classify model variant from tensor metadata names."""
    keys = metadata.keys()
    # Nano conversion emits per-layer map and Laurel tensors absent in standard Gemma.
    if any(".per_layer_map." in name or ".laurel." in name or ".post_laurel_layernorm." in name for name in keys):
        return ModelVariant.NANO
    return ModelVariant.STANDARD


def _normalize_architecture_overrides(overrides: dict[str, int | float] | None) -> dict[str, int | float] | None:
    if overrides is None:
        return None
    normalized: dict[str, int | float] = {}
    for key, value in overrides.items():
        if isinstance(value, bool) or not isinstance(value, (int, float)):  # pyright: ignore[reportUnnecessaryIsInstance]
            msg = "architecture_overrides values must be numeric (int|float)"
            raise TypeError(msg)
        normalized[key] = value
    return normalized


def _invoke_init_model_with_options(
    _core: object,
    metadata: dict[str, tuple[int, tuple[int, ...], str]],
    overrides: dict[str, int | float],
    descriptor: dict[str, object],
) -> object:
    init_model_with_options = getattr(_core, "init_model_with_options", None)
    if callable(init_model_with_options):
        return init_model_with_options(metadata, overrides, descriptor)
    return None


def _invoke_legacy_init_model(
    _core: object,
    metadata: dict[str, tuple[int, tuple[int, ...], str]],
    overrides: dict[str, int | float] | None,
    backend: GenerationBackend | EmbeddingBackend,
) -> object:
    if overrides is None:
        return backend.init_model(metadata)

    init_model_fn = getattr(_core, "init_model", None)
    if not callable(init_model_fn):
        msg = "Mojo core does not expose init_model(metadata, architecture_overrides)"
        raise TypeError(msg)
    try:
        return init_model_fn(metadata, overrides)
    except TypeError as exc:
        msg = (
            "Mojo core init_model does not accept architecture_overrides. "
            "Rebuild/install a compatible mogemma._core extension."
        )
        raise RuntimeError(msg) from exc


def _initialize_llm(
    loader: ModelLoader,
    backend: GenerationBackend | EmbeddingBackend,
    *,
    device_selection: DeviceSelection,
    model_type: str,
    architecture_overrides: dict[str, int | float] | None = None,
) -> object:
    if _core is None:
        raise RuntimeError(_core_unavailable_message(model_type))

    metadata = loader.get_tensor_metadata()

    normalized_overrides = _normalize_architecture_overrides(architecture_overrides)
    descriptor = device_selection.as_runtime_descriptor()

    try:
        llm = _invoke_init_model_with_options(_core, metadata, normalized_overrides or {}, descriptor)
        if llm is None:
            llm = _invoke_legacy_init_model(_core, metadata, normalized_overrides, backend)
    except ValueError:
        raise
    except Exception as exc:
        msg = f"{model_type} model failed to initialize from '{loader.model_path}': {exc}"
        raise RuntimeError(msg) from exc

    if isinstance(llm, dict):
        llm.setdefault("device_selection", descriptor)
        llm.setdefault("device_backend", device_selection.backend)
        llm.setdefault("device_kind", device_selection.device_kind)
        llm.setdefault("device_index", device_selection.device_index)
        llm.setdefault("device_request", device_selection.requested)
        llm.setdefault("device_availability_source", device_selection.availability_source)
        llm.setdefault("device_strict", device_selection.strict)
        if normalized_overrides:
            llm.setdefault("architecture_overrides", normalized_overrides)
    return llm


def _resolve_generation_backend(device: str) -> GenerationBackend:
    if _core is None:
        raise RuntimeError(_core_unavailable_message("generation"))
    return _resolve_generation_backend_impl(device=device, core_module=_core)


def _resolve_embedding_backend(device: str) -> EmbeddingBackend:
    if _core is None:
        raise RuntimeError(_core_unavailable_message("embedding"))
    return _resolve_embedding_backend_impl(device=device, core_module=_core)


def _sample_next_token(logits: npt.ArrayLike, *, temperature: float, top_k: int, top_p: float) -> int:
    """Sample the next token ID from backend logits."""
    logits_array = np.asarray(logits, dtype=np.float64).reshape(-1)
    if logits_array.size == 0:
        msg = "backend returned empty logits"
        raise ValueError(msg)

    if temperature <= 0:
        return int(np.argmax(logits_array))

    filtered_logits = logits_array / temperature

    if 0 < top_k < filtered_logits.size:
        top_indices = np.argpartition(filtered_logits, -top_k)[-top_k:]
        top_mask = np.zeros(filtered_logits.size, dtype=bool)
        top_mask[top_indices] = True
        filtered_logits = np.where(top_mask, filtered_logits, -np.inf)

    if 0.0 < top_p < 1.0:
        sorted_indices = np.argsort(filtered_logits)[::-1]
        sorted_logits = filtered_logits[sorted_indices]
        finite_mask = np.isfinite(sorted_logits)

        if finite_mask.any():
            finite_logits = sorted_logits[finite_mask]
            finite_logits = finite_logits - np.max(finite_logits)
            finite_probs = np.exp(finite_logits)
            finite_probs = finite_probs / finite_probs.sum()

            cumulative = np.cumsum(finite_probs)
            overflow = cumulative > top_p
            if overflow.any():
                first_overflow = int(np.argmax(overflow))
                remove_indices = sorted_indices[finite_mask][first_overflow + 1 :]
                filtered_logits[remove_indices] = -np.inf

    if not np.isfinite(filtered_logits).any():
        return int(np.argmax(logits_array))

    stable_logits = filtered_logits - np.max(filtered_logits)
    probs = np.exp(stable_logits)
    probs_sum = probs.sum()

    if probs_sum <= 0 or not np.isfinite(probs_sum):
        return int(np.argmax(logits_array))

    probs = probs / probs_sum
    return int(np.random.default_rng().choice(probs.size, p=probs))


def _normalize_eos_token_id(tokenizer: _Tokenizer) -> int:
    """Return an EOS token id from tokenizer aliases."""
    for eos_token in _EOS_TOKEN_ID_ALIASES:
        eos_token_id = tokenizer.token_to_id(eos_token)
        if eos_token_id is not None:
            return int(eos_token_id)
    return 0


def _is_instruction_tuned_model(model_path: Path, config_model_path: str | Path) -> bool:
    """Return True when the configured model looks like an instruction-tuned Gemma variant."""
    configured = str(config_model_path).lower()
    resolved = model_path.name.lower()
    return configured.endswith("-it") or resolved.endswith("-it")


def _format_instruction_prompt(prompt: str) -> str:
    """Wrap plain user prompts in Gemma instruction-turn format."""
    if _INSTRUCTION_START in prompt or _INSTRUCTION_END in prompt:
        return prompt
    return f"{_INSTRUCTION_START}user\n{prompt}\n{_INSTRUCTION_END}\n{_INSTRUCTION_START}model\n"


def _reset_llm_session_state(llm: object) -> None:
    """Reset mutable session fields before a new generation run."""
    if not isinstance(llm, dict):
        return

    llm["pos"] = 0
    if _core is not None and hasattr(_core, "reset_cache"):
        _core.reset_cache(llm)
    else:
        # Fallback for legacy core
        for cache_key in ("k_cache", "v_cache"):
            cache = llm.get(cache_key)
            if cache is None:
                continue
            if isinstance(cache, (int, float)):
                continue
            if hasattr(cache, "fill"):
                cache.fill(0.0)
                continue
            if isinstance(cache, list):
                for i in range(len(cache)):
                    cache[i] = 0.0
                continue
            try:
                np.asarray(cache).fill(0.0)
            except (TypeError, ValueError, AttributeError, NotImplementedError):
                continue


class EmbeddingModel:
    """Python interface for the Gemma 3 embedding engine."""

    def __init__(self, config: EmbeddingConfig | str | None = None, tokenizer: _Tokenizer | None = None) -> None:
        """Initialize the embedding model.

        Args:
            config: Model ID string, ``EmbeddingConfig``, or ``None`` for defaults.
            tokenizer: Optional pre-built tokenizer instance.
        """
        if config is None:
            config = EmbeddingConfig()
        elif isinstance(config, str):
            config = EmbeddingConfig(model_path=config)
        self.config = config
        self._tokenizer = tokenizer

        self._device_selection: DeviceSelection = resolve_device_selection(config.device)
        # Resolve model path (Hub or local)
        self.model_path = _resolve_model_path(config.model_path, config.cache_path)
        self._loader = auto_loader(self.model_path)
        self._backend = _resolve_embedding_backend(self._device_selection.effective_device)
        self._backend_id = self._backend.backend_id

        # Initialize Mojo core
        self._llm: object | None = _initialize_llm(
            self._loader,
            self._backend,
            device_selection=self._device_selection,
            model_type="embedding",
            architecture_overrides=config.architecture_overrides,
        )

    def _ensure_tokenizer(self) -> _Tokenizer:
        if self._tokenizer is not None:
            return self._tokenizer
        if not SENTENCEPIECE_INSTALLED:
            msg = (
                "Text tokenization requires 'sentencepiece' dependency. "
                "Install with: pip install 'mogemma[llm]' or call embed_tokens(...) with pre-tokenized inputs."
            )
            raise ModuleNotFoundError(msg)

        sp_path = self.model_path / "tokenizer.model"
        if sp_path.exists():
            self._tokenizer = _Tokenizer(str(sp_path))
            return self._tokenizer

        msg = f"No tokenizer.model found in {self.model_path}"
        raise FileNotFoundError(msg)

    def _embed_token_array(self, tokens: Sequence[Sequence[int]], input_count: int) -> npt.NDArray[np.float32]:
        if self._llm is None:
            raise RuntimeError(_core_unavailable_message("embedding"))

        if input_count > 0:
            expected_len = len(tokens[0])
            if any(len(row) != expected_len for row in tokens):
                msg = (
                    "all token sequences in a batch must have the exact same length "
                    "(padding is required for batch inference)"
                )
                raise ValueError(msg)

        raw_embeddings = self._backend.generate_embeddings(self._llm, tokens)
        embeddings = np.asarray(raw_embeddings, dtype=np.float32)

        if embeddings.ndim != _EXPECTED_MATRIX_DIMS:
            msg = f"expected 2D embeddings from backend, got {embeddings.ndim}D"
            raise ValueError(msg)
        if embeddings.shape[0] != input_count:
            msg = f"backend returned {embeddings.shape[0]} embedding rows for {input_count} inputs"
            raise ValueError(msg)
        return embeddings

    def embed(self, text: str | list[str]) -> npt.NDArray[np.float32]:
        """Generate embeddings for text by tokenizing in Python, then running Mojo inference."""
        with tracer.start_as_current_span("EmbeddingModel.embed") as span:
            if isinstance(text, str):
                text = [text]
            if not text:
                msg = "text input must contain at least one value"
                raise ValueError(msg)

            span.set_attribute("text_count", len(text))

        tokenizer = self._ensure_tokenizer()
        tokenizer.enable_truncation(max_length=self.config.max_sequence_length)
        tokenizer.enable_padding()
        encoded = tokenizer.encode_batch(text)
        tokens = [e.ids for e in encoded]
        return self._embed_token_array(tokens, len(text))

    def embed_tokens(self, tokens: Sequence[Sequence[int]] | npt.NDArray[np.int32]) -> npt.NDArray[np.float32]:
        """Generate embeddings directly from pre-tokenized IDs using Mojo inference."""
        token_list = []
        if isinstance(tokens, np.ndarray):
            if tokens.ndim != _EXPECTED_MATRIX_DIMS:
                msg = f"tokens must be a 2D array of token IDs, got shape {tokens.shape}"
                raise ValueError(msg)
            token_list = tokens.tolist()
        else:
            token_list = list(tokens)

        if len(token_list) == 0:
            msg = "tokens must contain at least one row"
            raise ValueError(msg)
        return self._embed_token_array(token_list, len(token_list))

    @property
    def tokenizer(self) -> _Tokenizer:
        """Access the tokenizer (loads lazily)."""
        return self._ensure_tokenizer()

    def close(self) -> None:
        """Release underlying resources and free the memory arena."""
        if hasattr(self, "_llm") and self._llm is not None and _core is not None and hasattr(_core, "free_arena"):
            _core.free_arena(self._llm)
            self._llm = None
        if hasattr(self, "_loader"):
            self._loader.close()

    def __del__(self) -> None:
        """Cleanup on garbage collection."""
        self.close()

    def __enter__(self) -> Self:
        """Enter the context manager."""
        return self

    def __exit__(self, *args: object) -> None:
        """Exit the context manager and release resources."""
        self.close()


class SyncGemmaModel:
    """Python interface for the Gemma 3 text generation engine."""

    def __init__(self, config: GenerationConfig | str | None = None, tokenizer: _Tokenizer | None = None) -> None:
        """Initialize the text model.

        Args:
            config: Model ID string, ``GenerationConfig``, or ``None`` for defaults.
            tokenizer: Optional pre-built tokenizer instance.
        """
        if config is None:
            config = GenerationConfig()
        elif isinstance(config, str):
            config = GenerationConfig(model_path=config)
        self.config = config
        self._tokenizer = tokenizer

        self._device_selection: DeviceSelection = resolve_device_selection(config.device)
        # Resolve model path (Hub or local)
        self.model_path = _resolve_model_path(config.model_path, config.cache_path)
        self._loader = auto_loader(self.model_path)
        self._instruction_tuned = _is_instruction_tuned_model(self.model_path, config.model_path)
        self._backend = _resolve_generation_backend(self._device_selection.effective_device)
        self._backend_id = self._backend.backend_id

        # Initialize Mojo core
        self._llm: object | None = _initialize_llm(
            self._loader,
            self._backend,
            device_selection=self._device_selection,
            model_type="generation",
            architecture_overrides=config.architecture_overrides,
        )

    def _ensure_tokenizer(self) -> _Tokenizer:
        if self._tokenizer is not None:
            return self._tokenizer
        if not SENTENCEPIECE_INSTALLED:
            msg = "Text generation requires 'sentencepiece' dependency. Install with: pip install 'mogemma[llm]'"
            raise ModuleNotFoundError(msg)

        sp_path = self.model_path / "tokenizer.model"
        if sp_path.exists():
            self._tokenizer = _Tokenizer(str(sp_path))
            return self._tokenizer

        msg = f"No tokenizer.model found in {self.model_path}"
        raise FileNotFoundError(msg)

    def generate(
        self, prompt: str, images: Sequence[str | Path | bytes | npt.NDArray[np.generic]] | None = None
    ) -> str:
        """Generate text from the given prompt."""
        return "".join(list(self.generate_stream(prompt, images=images)))

    def generate_stream(
        self, prompt: str, images: Sequence[str | Path | bytes | npt.NDArray[np.generic]] | None = None
    ) -> Generator[str, None, None]:
        """Generate text as a stream of tokens."""
        tokenizer = self._ensure_tokenizer()
        prompt_to_encode = _format_instruction_prompt(prompt) if self._instruction_tuned else prompt
        with tracer.start_as_current_span("SyncGemmaModel.generate_stream") as span:
            span.set_attribute("prompt_length", len(prompt))
            tokenizer.enable_truncation(max_length=self.config.max_sequence_length)
            tokenizer.enable_padding()
            encoded = tokenizer.encode(prompt_to_encode)
            tokens = encoded.ids
            if not tokens or tokens[0] != _BOS_TOKEN_ID:
                tokens = [_BOS_TOKEN_ID, *list(tokens)]

        if self._llm is None:
            raise RuntimeError(_core_unavailable_message("generation"))

        _reset_llm_session_state(self._llm)

        if images is not None:
            hydrated = ImageHydrator().hydrate(images)
            self._backend.process_images(self._llm, hydrated)

        for t in tokens[:-1]:
            self._backend.step(self._llm, int(t), self.config.temperature, self.config.top_k, self.config.top_p)

        if tokens:
            current_token = int(tokens[-1])
        else:
            eos_token_id = _normalize_eos_token_id(tokenizer)
            current_token = int(eos_token_id)

        eos_token_id = _normalize_eos_token_id(tokenizer)
        for _ in range(self.config.max_tokens):
            logits = self._backend.step(
                self._llm, current_token, self.config.temperature, self.config.top_k, self.config.top_p
            )

            next_token = _sample_next_token(
                logits, temperature=self.config.temperature, top_k=self.config.top_k, top_p=self.config.top_p
            )
            if next_token == eos_token_id:
                return

            decoded = tokenizer.decode([next_token])
            if not isinstance(decoded, str):  # pyright: ignore[reportUnnecessaryIsInstance]
                msg = "tokenizer.decode returned non-string output"
                raise TypeError(msg)
            if not decoded:
                return

            yield decoded
            current_token = next_token

    @property
    def tokenizer(self) -> _Tokenizer:
        """Access the tokenizer (loads lazily)."""
        return self._ensure_tokenizer()

    def close(self) -> None:
        """Release underlying resources and free the memory arena."""
        if hasattr(self, "_llm") and self._llm is not None and _core is not None and hasattr(_core, "free_arena"):
            _core.free_arena(self._llm)
            self._llm = None
        if hasattr(self, "_loader"):
            self._loader.close()

    def __del__(self) -> None:
        """Cleanup on garbage collection."""
        self.close()

    def __enter__(self) -> Self:
        """Enter the context manager."""
        return self

    def __exit__(self, *args: object) -> None:
        """Exit the context manager and release resources."""
        self.close()


class AsyncGemmaModel:
    """Asynchronous wrapper for SyncGemmaModel."""

    def __init__(self, config: GenerationConfig | str | None = None) -> None:
        """Initialize the async model.

        Args:
            config: Model ID string, ``GenerationConfig``, or ``None`` for defaults.
        """
        self._model = SyncGemmaModel(config)

    async def generate(
        self, prompt: str, images: Sequence[str | Path | bytes | npt.NDArray[np.generic]] | None = None
    ) -> str:
        """Generate text asynchronously."""
        return await asyncio.to_thread(self._model.generate, prompt, images)

    async def generate_stream(
        self, prompt: str, images: Sequence[str | Path | bytes | npt.NDArray[np.generic]] | None = None
    ) -> AsyncIterator[str]:
        """Generate text as an async stream of tokens."""
        generator = self._model.generate_stream(prompt, images=images)

        def get_next() -> str | None:
            try:
                return next(generator)
            except StopIteration:
                return None

        try:
            while True:
                token = await asyncio.to_thread(get_next)
                if token is None:
                    break
                yield token
        except asyncio.CancelledError:
            generator.close()
            raise
        finally:
            with contextlib.suppress(RuntimeError):
                generator.close()

    @property
    def tokenizer(self) -> _Tokenizer:
        """Access to the underlying tokenizer."""
        return self._model.tokenizer
