"""Encoder utilities for HIEROMEM.

The default backend is a lightweight, deterministic stub encoder that
emits hash-seeded vectors. It keeps the repository self-contained while
preserving compatibility with heavier transformer backends when their
optional dependencies are installed.
"""
# flake8: noqa

from __future__ import annotations

import importlib
import hashlib
import os
from dataclasses import dataclass
from types import ModuleType
from typing import Any, Optional, Sequence, Union, cast

import numpy as np
import numpy.typing as npt

from hieromem_core.circuit_breaker import CircuitBreaker

VectorArray = npt.NDArray[np.float32]


_FORCE_STUB = os.environ.get(
    "HIEROMEM_FORCE_STUB_ENCODER",
    "").lower() in {
        "1",
        "true",
    "yes"}


def _optional_import(module_name: str) -> ModuleType | None:
    """Import *module_name* if available, otherwise return ``None``."""

    if _FORCE_STUB:
        return None
    try:
        return importlib.import_module(module_name)
    except ImportError:  # pragma: no cover - optional dependency
        return None


_TORCH: ModuleType | None = _optional_import("torch")
_SENTENCE_TRANSFORMERS: ModuleType | None = _optional_import(
    "sentence_transformers")
_TRANSFORMERS: ModuleType | None = _optional_import("transformers")


def _default_device() -> str:
    torch_module = cast(Optional[Any], _TORCH)
    if torch_module is None:
        return "cpu"

    cuda = getattr(torch_module, "cuda", None)
    is_available = getattr(
        cuda,
        "is_available",
        None) if cuda is not None else None
    if callable(is_available) and bool(is_available()):
        return "cuda"
    return "cpu"


def _resolve_sentence_transformer() -> type[Any] | None:
    module = _SENTENCE_TRANSFORMERS
    if module is None:
        return None
    candidate = getattr(module, "SentenceTransformer", None)
    return cast(Optional[type[Any]], candidate)


def _resolve_transformer_attr(name: str) -> Any | None:
    module = _TRANSFORMERS
    if module is None:
        return None
    return getattr(module, name, None)


def _torch_module() -> Any | None:
    return cast(Optional[Any], _TORCH)


# --------------------------------------------------------------------------- #
# Configuration Dataclass
# --------------------------------------------------------------------------- #
@dataclass
class EncoderConfig:
    model_name: str = "stub-hash"
    device: str = _default_device()
    normalize_embeddings: bool = True
    backend: str = "stub"  # "stub", "sentence-transformers", or "huggingface"
    embedding_dim: Optional[int] = 384


# --------------------------------------------------------------------------- #
# Encoder Wrapper
# --------------------------------------------------------------------------- #
class Encoder:
    """Unified interface for text encoding."""

    def __init__(
        self,
        config: Optional[EncoderConfig] = None,
        *,
        model_name: Optional[str] = None,
        backend: Optional[str] = None,
        embedding_dim: Optional[int] = None,
    ):
        config = config or EncoderConfig()

        if model_name is not None:
            config.model_name = model_name
        if backend is not None:
            config.backend = backend
        if embedding_dim is not None:
            config.embedding_dim = embedding_dim

        self.config = config
        self.device = config.device
        self.backend = config.backend.lower()
        self.model: Any | None = None
        self.tokenizer: Any | None = None

        if self.backend == "stub":
            self._init_stub()
        elif self.backend == "sentence-transformers":
            self._load_sentence_transformer()
        elif self.backend == "huggingface":
            self._load_huggingface_encoder()
        else:
            raise ValueError(f"Unknown backend: {self.backend}")

        self._breaker = CircuitBreaker(failure_threshold=5, recovery_timeout=30.0)

    # ------------------------------------------------------------------ #
    # Loading functions
    # ------------------------------------------------------------------ #
    def _init_stub(self) -> None:
        if not self.config.embedding_dim:
            self.config.embedding_dim = 384

    def _load_sentence_transformer(self) -> None:
        sentence_transformer_cls = _resolve_sentence_transformer()
        if sentence_transformer_cls is None:
            raise ImportError(
                "SentenceTransformers not installed. Run: pip install sentence-transformers"
            )
        model = sentence_transformer_cls(
            self.config.model_name, device=self.device)
        get_dim = getattr(model, "get_sentence_embedding_dimension", None)
        if not callable(get_dim):  # pragma: no cover - defensive branch
            raise TypeError(
                "SentenceTransformer missing get_sentence_embedding_dimension()")
        self.model = model
        self.config.embedding_dim = int(get_dim())
        print(
            f"[Encoder] Loaded SentenceTransformer '{self.config.model_name}' "
            f"on {self.device} (dim={self.config.embedding_dim})."
        )

    def _load_huggingface_encoder(self) -> None:
        auto_tokenizer = _resolve_transformer_attr("AutoTokenizer")
        auto_model = _resolve_transformer_attr("AutoModel")
        torch_module = _torch_module()
        if auto_tokenizer is None or auto_model is None or torch_module is None:
            raise ImportError(
                "Transformers backend requires the 'transformers' and 'torch' packages."
            )
        tokenizer = auto_tokenizer.from_pretrained(self.config.model_name)
        model = auto_model.from_pretrained(self.config.model_name)
        to_method = getattr(model, "to", None)
        if callable(to_method):
            model = to_method(self.device)
        eval_method = getattr(model, "eval", None)
        if callable(eval_method):
            eval_method()
        self.tokenizer = tokenizer
        self.model = model

        no_grad = getattr(torch_module, "no_grad", None)
        dummy: VectorArray
        if callable(no_grad):
            with no_grad():
                dummy = self.encode(["test"])
        else:  # pragma: no cover - fallback for unusual torch builds
            dummy = self.encode(["test"])

        if dummy.ndim < 2:
            raise ValueError(
                "Expected huggingface encoder to return batched embeddings")
        self.config.embedding_dim = int(dummy.shape[1])
        print(
            f"[Encoder] Loaded HuggingFace model '{self.config.model_name}' "
            f"on {self.device} (dim={self.config.embedding_dim})."
        )

    # ------------------------------------------------------------------ #
    # Encoding interface
    # ------------------------------------------------------------------ #
    def encode(self, texts: Union[str, Sequence[str]]) -> VectorArray:
        """Encode one or more text strings into dense embeddings."""
        if isinstance(texts, str):
            text_sequence: Sequence[str] = [texts]
        else:
            text_sequence = list(texts)

        @self._breaker
        def _execute_backend(seq: Sequence[str]) -> VectorArray:
            if self.backend == "stub":
                return np.asarray([self._encode_stub(text)
                                     for text in seq], dtype=np.float32)
            elif self.backend == "sentence-transformers":
                model = self._require_model()
                encode_fn = getattr(model, "encode", None)
                if not callable(encode_fn):
                    raise TypeError("SentenceTransformer model missing encode()")
                return np.asarray(
                    encode_fn(
                        seq,
                        convert_to_numpy=True,
                        show_progress_bar=False),
                    dtype=np.float32,
                )
            elif self.backend == "huggingface":
                return self._encode_huggingface(seq)
            else:
                raise RuntimeError("Encoder backend not initialized properly.")

        result = _execute_backend(text_sequence)

        if self.config.normalize_embeddings:
            result = self._normalize(result)
        return result

    def _encode_huggingface(self, texts: Sequence[str]) -> VectorArray:
        """Mean-pool last hidden state for HuggingFace models."""
        tokenizer = self._require_tokenizer()
        model = self._require_model()
        torch_module = _torch_module()
        if torch_module is None:
            raise RuntimeError("Torch is required for the huggingface backend")

        encoding: Any = tokenizer(
            texts,
            padding=True,
            truncation=True,
            return_tensors="pt")
        to_method = getattr(encoding, "to", None)
        if callable(to_method):
            encoding = to_method(self.device)

        no_grad = getattr(torch_module, "no_grad", None)
        if not callable(no_grad):  # pragma: no cover - defensive branch
            raise TypeError("torch.no_grad context manager is unavailable")

        with no_grad():
            output = model(**encoding)

        hidden = getattr(output, "last_hidden_state")
        attention_mask = encoding["attention_mask"].unsqueeze(-1)
        pooled = (hidden * attention_mask).sum(dim=1) / \
            attention_mask.sum(dim=1)
        array = pooled.cpu().numpy()
        return np.asarray(array, dtype=np.float32)

    def _encode_stub(self, text: str) -> VectorArray:
        """Generate a deterministic hash-based embedding."""
        dim = int(self.config.embedding_dim or 384)
        digest = hashlib.sha256(text.encode("utf-8")).digest()
        seed = int.from_bytes(digest[:8], "big", signed=False)
        rng = np.random.default_rng(seed)
        vec = rng.standard_normal(dim, dtype=np.float32)
        return vec

    # ------------------------------------------------------------------ #
    # Utilities
    # ------------------------------------------------------------------ #
    @staticmethod
    def _normalize(vecs: VectorArray) -> VectorArray:
        """L2 normalize embeddings."""
        norm = np.linalg.norm(vecs, axis=1, keepdims=True)
        clipped = np.clip(norm, 1e-8, None)
        normalised = (vecs / clipped).astype(np.float32)
        return cast(VectorArray, normalised)

    def encode_single(self, text: str) -> VectorArray:
        """Shortcut for single input (returns 1D vector)."""
        single = np.asarray(self.encode([text])[0], dtype=np.float32)
        return single

    @property
    def embedding_dim(self) -> int:
        """Expose resolved embedding dimensionality for compatibility callers."""

        return int(self.config.embedding_dim or 0)

    def info(self) -> dict[str, Any]:
        """Return encoder configuration details."""
        return {
            "backend": self.backend,
            "model_name": self.config.model_name,
            "device": self.device,
            "embedding_dim": self.config.embedding_dim,
            "normalize_embeddings": self.config.normalize_embeddings,
        }

    def _require_model(self) -> Any:
        if self.model is None:
            raise RuntimeError("Encoder backend model has not been loaded")
        return self.model

    def _require_tokenizer(self) -> Any:
        if self.tokenizer is None:
            raise RuntimeError("Encoder tokenizer has not been loaded")
        return self.tokenizer


# --------------------------------------------------------------------------- #
# Simplified Facade for Text-only Encoding
# --------------------------------------------------------------------------- #
class TextEncoder:
    """Lightweight wrapper for quick access to text embeddings."""

    def __init__(
        self,
        model_name: str | None = None,
        backend: str = "stub",
        embedding_dim: int | None = None,
    ):
        cfg = EncoderConfig(
            model_name=model_name or "stub-hash",
            backend=backend,
            embedding_dim=embedding_dim or 384,
        )
        self.encoder = Encoder(cfg)

    def encode(self, text: Union[str, Sequence[str]]) -> VectorArray:
        """Return embeddings for the provided text input."""
        if isinstance(text, str):
            return self.encoder.encode([text])
        return self.encoder.encode(list(text))

    def encode_single(self, text: str) -> VectorArray:
        """Return an embedding for a single string."""
        return self.encoder.encode_single(text)

    def info(self) -> dict[str, Any]:
        """Expose backend details."""
        return self.encoder.info()


class MultiEncoderRegistry:
    """Registry for managing multiple encoder instances."""

    def __init__(self) -> None:
        self._encoders: dict[str, Encoder] = {}

    def register(self, encoder_id: str, encoder: Encoder) -> None:
        self._encoders[encoder_id] = encoder

    def get(self, encoder_id: str) -> Encoder | None:
        return self._encoders.get(encoder_id)

    def default(self) -> Encoder | None:
        if not self._encoders:
            return None
        # Deterministic default
        return self._encoders[sorted(self._encoders.keys())[0]]
