"""LLMLingua-2 based context compressor."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from context_manager.distill.base import Compressor

# Avoid importing heavy dependencies at module level
if TYPE_CHECKING:
    from llmlingua import PromptCompressor


class LLMLinguaCompressor(Compressor):
    """Context compressor using Microsoft's LLMLingua-2.

    LLMLingua-2 uses a small, fast BERT-based token classifier (e.g.
    XLM-RoBERTa) to identify and retain essential tokens, achieving
    3-6x speedup over the original perplexity-based LLMLingua.

    This class lazily initializes the heavy model on first use.

    Args:
        model_name: The HuggingFace model identifier. Defaults to
            "microsoft/llmlingua-2-xlm-roberta-large-meetingbank".
        device_map: Device to load the model on (e.g. "cpu", "cuda",
            "mps"). Defaults to "cpu" for broad compatibility.
    """

    def __init__(
        self,
        model_name: str = "microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
        device_map: str = "cpu",
        use_llmlingua2: bool = True,
    ) -> None:
        self._model_name = model_name
        self._device_map = device_map
        self._use_llmlingua2 = use_llmlingua2
        self._compressor: PromptCompressor | None = None

    def _get_compressor(self) -> PromptCompressor:
        """Lazily load the PromptCompressor."""
        if self._compressor is None:
            try:
                from llmlingua import PromptCompressor
            except ImportError as exc:
                raise ImportError(
                    "llmlingua is required for LLMLinguaCompressor. "
                    "Install it with: pip install llm-ctx-mgr[distill]"
                ) from exc

            # Initialize the model. This might download weights (~2GB).
            self._compressor = PromptCompressor(
                model_name=self._model_name,
                device_map=self._device_map,
                use_llmlingua2=self._use_llmlingua2,
            )
        return self._compressor

    def compress(
        self,
        text: str,
        target_token_count: int | None = None,
        rate: float = 0.5,
        **kwargs: Any,
    ) -> str:
        """Compress text using LLMLingua-2.

        Args:
            text: The text to compress.
            target_token_count: Specific target token count (optional).
            rate: Target compression rate (0.0 to 1.0) if token count
                is not specified. Default is 0.5 (50% reduction).
            **kwargs: Additional arguments passed to
                `compress_prompt`.

        Returns:
            The compressed text.
        """
        if not text.strip():
            return text

        compressor = self._get_compressor()

        # If target tokens provided, use that; otherwise use rate.
        # LLMLingua's API is a bit flexible here.
        result = compressor.compress_prompt(
            context=[text],
            target_token=target_token_count if target_token_count else -1,
            rate=rate if not target_token_count else 0.5,
            **kwargs,
        )

        return result["compressed_prompt"]

    def model_name(self) -> str:
        return f"LLMLinguaCompressor({self._model_name})"
