"""Condensa message transformer for LangChain.

Compresses and decompresses LangChain messages using Condensa protocol.
Works with LCEL chains, LangGraph agents, and standalone.

Usage:
    from langchain_condensa import CondensaTransformer

    transformer = CondensaTransformer()

    # Compress a list of messages
    compressed = transformer.compress_messages(messages)

    # In an LCEL chain
    chain = transformer.as_compressor() | prompt | llm

    # Pure functions
    from langchain_condensa import compress_message, decompress_message
    msg = HumanMessage(content="Search for SpaceX news, top 5 results")
    compressed = compress_message(msg)
"""

from __future__ import annotations

from typing import Any, Sequence

from langchain_core.messages import (
    BaseMessage,
    HumanMessage,
    AIMessage,
    SystemMessage,
)
from langchain_core.runnables import RunnableLambda

from cdn import encode, decode, count_tokens
from cdn.encoder import CondensaEncoder
from cdn.decoder import CondensaDecoder


# Metadata key used to store compression info
_META_KEY = "condensa"


class CondensaTransformer:
    """Compresses and decompresses LangChain messages using Condensa.

    Parameters
    ----------
    min_tokens : int
        Only compress messages with at least this many tokens.
        Messages below this threshold pass through unchanged.
        Default: 10 (Condensa adds ~3 tokens of structure, so very
        short messages may not benefit).
    preserve_system : bool
        If True, never compress SystemMessage content.
        System messages often contain instructions that must remain
        in natural language for the LLM. Default: True.
    tag : bool
        If True, add a ``[CDN]`` prefix to compressed content so
        receiving agents can identify Condensa-encoded messages.
        Default: True.
    """

    def __init__(
        self,
        *,
        min_tokens: int = 10,
        max_compression_pct: float = 85.0,
        preserve_system: bool = True,
        tag: bool = True,
    ):
        self.min_tokens = min_tokens
        self.max_compression_pct = max_compression_pct
        self.preserve_system = preserve_system
        self.tag = tag
        self._encoder = CondensaEncoder()
        self._decoder = CondensaDecoder()

    # ── Core operations ──────────────────────────────────────────────

    def compress(self, text: str) -> dict:
        """Compress text and return result with metadata.

        Returns a dict with:
            - content: the compressed string
            - original_tokens: token count of original
            - compressed_tokens: token count of compressed
            - reduction_pct: percentage reduction
            - original: the original text (for decompression)
        """
        original_tokens = count_tokens(text)

        if original_tokens < self.min_tokens:
            return {
                "content": text,
                "original_tokens": original_tokens,
                "compressed_tokens": original_tokens,
                "reduction_pct": 0.0,
                "compressed": False,
            }

        compressed = self._encoder.encode(text)
        if self.tag:
            tagged = f"[CDN] {compressed}"
        else:
            tagged = compressed

        compressed_tokens = count_tokens(tagged)
        reduction = round((1 - compressed_tokens / original_tokens) * 100, 1) if original_tokens > 0 else 0.0

        # Only use compressed version if it actually saves tokens
        # Also reject suspiciously high compression (>max_compression_pct) —
        # the regex encoder likely destroyed content rather than compressed it.
        # Legitimate Condensa compression tops out around 70-80%.
        if compressed_tokens >= original_tokens or reduction > self.max_compression_pct:
            return {
                "content": text,
                "original_tokens": original_tokens,
                "compressed_tokens": original_tokens,
                "reduction_pct": 0.0,
                "compressed": False,
            }

        return {
            "content": tagged,
            "original_tokens": original_tokens,
            "compressed_tokens": compressed_tokens,
            "reduction_pct": reduction,
            "compressed": True,
        }

    def decompress(self, text: str) -> str:
        """Decompress a Condensa-encoded string back to natural language."""
        if text.startswith("[CDN] "):
            cdn_msg = text[6:]
        else:
            cdn_msg = text
        return self._decoder.decode(cdn_msg)

    # ── Message-level operations ─────────────────────────────────────

    def compress_message(self, message: BaseMessage) -> BaseMessage:
        """Compress a single LangChain message.

        Returns a new message with compressed content and metadata
        stored in additional_kwargs under the 'condensa' key.
        """
        if self.preserve_system and isinstance(message, SystemMessage):
            return message

        if not isinstance(message.content, str):
            return message

        result = self.compress(message.content)
        if not result["compressed"]:
            return message

        metadata = {
            "original": message.content,
            "original_tokens": result["original_tokens"],
            "compressed_tokens": result["compressed_tokens"],
            "reduction_pct": result["reduction_pct"],
        }

        kwargs = dict(message.additional_kwargs)
        kwargs[_META_KEY] = metadata

        return message.__class__(
            content=result["content"],
            additional_kwargs=kwargs,
            **({"name": message.name} if message.name else {}),
        )

    def decompress_message(self, message: BaseMessage) -> BaseMessage:
        """Decompress a Condensa-compressed message back to natural language.

        If the message has stored original text in metadata, uses that.
        Otherwise falls back to the decoder.
        """
        meta = message.additional_kwargs.get(_META_KEY)
        if meta and "original" in meta:
            return message.__class__(
                content=meta["original"],
                **({"name": message.name} if message.name else {}),
            )

        if isinstance(message.content, str) and message.content.startswith("[CDN] "):
            decoded = self.decompress(message.content)
            return message.__class__(
                content=decoded,
                **({"name": message.name} if message.name else {}),
            )

        return message

    def compress_messages(
        self, messages: Sequence[BaseMessage]
    ) -> list[BaseMessage]:
        """Compress a list of messages."""
        return [self.compress_message(m) for m in messages]

    def decompress_messages(
        self, messages: Sequence[BaseMessage]
    ) -> list[BaseMessage]:
        """Decompress a list of messages."""
        return [self.decompress_message(m) for m in messages]

    # ── Stats ────────────────────────────────────────────────────────

    def get_stats(self, messages: Sequence[BaseMessage]) -> dict:
        """Get compression stats for a list of messages.

        Analyzes messages that have Condensa metadata attached.
        """
        total_original = 0
        total_compressed = 0
        compressed_count = 0

        for msg in messages:
            meta = msg.additional_kwargs.get(_META_KEY)
            if meta:
                total_original += meta["original_tokens"]
                total_compressed += meta["compressed_tokens"]
                compressed_count += 1
            elif isinstance(msg.content, str):
                tokens = count_tokens(msg.content)
                total_original += tokens
                total_compressed += tokens

        return {
            "total_original_tokens": total_original,
            "total_compressed_tokens": total_compressed,
            "tokens_saved": total_original - total_compressed,
            "reduction_pct": round(
                (1 - total_compressed / total_original) * 100, 1
            ) if total_original > 0 else 0.0,
            "messages_compressed": compressed_count,
            "messages_total": len(messages),
        }

    # ── LCEL integration ─────────────────────────────────────────────

    def as_compressor(self) -> RunnableLambda:
        """Return a RunnableLambda that compresses messages.

        Use in LCEL chains:
            chain = transformer.as_compressor() | prompt | llm
        """
        def _compress(input_val: Any) -> Any:
            if isinstance(input_val, list):
                return self.compress_messages(input_val)
            if isinstance(input_val, dict) and "messages" in input_val:
                return {
                    **input_val,
                    "messages": self.compress_messages(input_val["messages"]),
                }
            return input_val

        return RunnableLambda(_compress)

    def as_decompressor(self) -> RunnableLambda:
        """Return a RunnableLambda that decompresses messages."""
        def _decompress(input_val: Any) -> Any:
            if isinstance(input_val, list):
                return self.decompress_messages(input_val)
            if isinstance(input_val, dict) and "messages" in input_val:
                return {
                    **input_val,
                    "messages": self.decompress_messages(input_val["messages"]),
                }
            return input_val

        return RunnableLambda(_decompress)


# ── Convenience functions ────────────────────────────────────────────

_default_transformer = None


def _get_default() -> CondensaTransformer:
    global _default_transformer
    if _default_transformer is None:
        _default_transformer = CondensaTransformer()
    return _default_transformer


def compress_message(message: BaseMessage) -> BaseMessage:
    """Compress a single message using default settings."""
    return _get_default().compress_message(message)


def decompress_message(message: BaseMessage) -> BaseMessage:
    """Decompress a single message."""
    return _get_default().decompress_message(message)


def compress_messages(messages: Sequence[BaseMessage]) -> list[BaseMessage]:
    """Compress a list of messages using default settings."""
    return _get_default().compress_messages(messages)
