"""Text embedding using Qwen3 embedding model.

This module provides a text embedding interface using the Qwen3-Embedding-8B model
from HuggingFace. Embeddings are truncated to 256 dimensions using Matryoshka
representation learning for efficient storage and computation.
"""

from typing import Union

import torch
from transformers import AutoModel, AutoTokenizer


class TextEmbedder:
    """Embeds text using Qwen3-Embedding-8B with Matryoshka truncation.

    This embedder uses a generic classification prompt suitable for general-purpose
    text encoding tasks.

    Attributes:
        model_name: HuggingFace model identifier for the embedding model.
        embedding_dim: Output embedding dimension after truncation.
        device: Compute device (cuda/cpu).
    """

    DEFAULT_MODEL = "Qwen/Qwen3-Embedding-8B"
    DEFAULT_PROMPT = "Encode this text for classification: "
    OUTPUT_DIM = 256

    def __init__(
        self,
        model_name: str | None = None,
        device: str | None = None,
        load_in_8bit: bool = False,
    ) -> None:
        """Initialize the text embedder.

        Args:
            model_name: HuggingFace model identifier. Defaults to Qwen3 embedding.
            device: Compute device. Defaults to CUDA if available.
            load_in_8bit: Whether to load model in 8-bit quantization for memory
                efficiency.
        """
        self.model_name = model_name or self.DEFAULT_MODEL
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._load_in_8bit = load_in_8bit

        self._tokenizer: AutoTokenizer | None = None
        self._model: AutoModel | None = None

    def _ensure_loaded(self) -> None:
        """Lazily load the model and tokenizer on first use."""
        if self._model is not None:
            return

        self._tokenizer = AutoTokenizer.from_pretrained(
            self.model_name,
            trust_remote_code=True,
        )

        load_kwargs: dict = {
            "trust_remote_code": True,
        }
        if self._load_in_8bit:
            load_kwargs["load_in_8bit"] = True
            load_kwargs["device_map"] = "auto"
        else:
            load_kwargs["torch_dtype"] = torch.float16

        self._model = AutoModel.from_pretrained(self.model_name, **load_kwargs)

        if not self._load_in_8bit:
            self._model = self._model.to(self.device)

        self._model.eval()

    def embed(self, text: str) -> torch.Tensor:
        """Embed a single text string.

        Args:
            text: Input text to embed.

        Returns:
            Embedding tensor of shape (OUTPUT_DIM,).
        """
        return self.embed_batch([text])[0]

    def embed_batch(self, texts: list[str]) -> torch.Tensor:
        """Embed a batch of text strings.

        Args:
            texts: List of input texts to embed.

        Returns:
            Embedding tensor of shape (batch_size, OUTPUT_DIM).
        """
        self._ensure_loaded()
        assert self._tokenizer is not None
        assert self._model is not None

        # Prepend generic classification prompt
        prompted_texts = [f"{self.DEFAULT_PROMPT}{t}" for t in texts]

        inputs = self._tokenizer(
            prompted_texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt",
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self._model(**inputs)
            # Use last hidden state with mean pooling
            attention_mask = inputs["attention_mask"]
            hidden_states = outputs.last_hidden_state
            mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size())
            sum_embeddings = torch.sum(hidden_states * mask_expanded, dim=1)
            sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
            embeddings = sum_embeddings / sum_mask

        # Matryoshka truncation to OUTPUT_DIM
        embeddings = embeddings[:, : self.OUTPUT_DIM]

        # L2 normalize
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=-1)

        return embeddings

    def get_memory_footprint_gb(self) -> float:
        """Estimate GPU memory usage in GB.

        Returns:
            Estimated memory footprint in gigabytes.
        """
        self._ensure_loaded()
        assert self._model is not None

        param_bytes = sum(
            p.numel() * p.element_size() for p in self._model.parameters()
        )
        # Add ~20% overhead for activations and gradients
        total_bytes = param_bytes * 1.2
        return total_bytes / (1024**3)

    def unload(self) -> None:
        """Unload model from memory."""
        if self._model is not None:
            del self._model
            self._model = None
        if self._tokenizer is not None:
            del self._tokenizer
            self._tokenizer = None
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
