"""Condensa callback handler for LangChain.

Monitors LLM calls, logs compression opportunities, and tracks
cumulative token savings across a session.

Usage:
    from langchain_condensa import CondensaCallbackHandler

    handler = CondensaCallbackHandler()
    result = chain.invoke(input, config={"callbacks": [handler]})

    # After execution
    print(handler.stats)
    # {'total_original_tokens': 1847, 'total_compressed_tokens': 523,
    #  'reduction_pct': 71.7, 'calls': 5}
"""

from __future__ import annotations

import logging
from typing import Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.messages import BaseMessage

from cdn import count_tokens
from cdn.encoder import CondensaEncoder

logger = logging.getLogger("langchain_condensa")


class CondensaCallbackHandler(BaseCallbackHandler):
    """Callback handler that tracks Condensa compression metrics.

    Attaches to any LangChain chain/agent via the callbacks config.
    Analyzes messages sent to LLMs and reports how much Condensa
    would save (or did save, if messages are already compressed).

    Parameters
    ----------
    log_level : int
        Logging level for compression reports. Default: logging.INFO.
    verbose : bool
        If True, log per-message compression details.
        If False, only log per-call summaries. Default: False.
    """

    def __init__(self, *, log_level: int = logging.INFO, verbose: bool = False):
        super().__init__()
        self.log_level = log_level
        self.verbose = verbose
        self._encoder = CondensaEncoder()

        # Cumulative stats
        self._total_original = 0
        self._total_compressed = 0
        self._total_already_compressed = 0
        self._calls = 0

    @property
    def stats(self) -> dict:
        """Return cumulative compression stats."""
        return {
            "total_original_tokens": self._total_original,
            "total_compressed_tokens": self._total_compressed,
            "tokens_saved": self._total_original - self._total_compressed,
            "reduction_pct": round(
                (1 - self._total_compressed / self._total_original) * 100, 1
            ) if self._total_original > 0 else 0.0,
            "calls": self._calls,
            "messages_already_compressed": self._total_already_compressed,
        }

    def reset(self):
        """Reset cumulative stats."""
        self._total_original = 0
        self._total_compressed = 0
        self._total_already_compressed = 0
        self._calls = 0

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[BaseMessage]],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Analyze messages being sent to a chat model."""
        self._calls += 1

        call_original = 0
        call_compressed = 0
        already_compressed = 0

        for message_list in messages:
            for msg in message_list:
                if not isinstance(msg.content, str):
                    continue

                content = msg.content
                tokens = count_tokens(content)
                call_original += tokens

                # Check if already Condensa-compressed
                if content.startswith("[CDN] ") or msg.additional_kwargs.get("condensa"):
                    call_compressed += tokens
                    already_compressed += 1
                    continue

                # Simulate compression
                encoded = self._encoder.encode(content)
                encoded_tokens = count_tokens(encoded)

                if encoded_tokens < tokens:
                    call_compressed += encoded_tokens
                    if self.verbose:
                        saving = round((1 - encoded_tokens / tokens) * 100, 1)
                        logger.log(
                            self.log_level,
                            "Condensa: %s -> %s tokens (-%s%%) | %s",
                            tokens, encoded_tokens, saving,
                            content[:60] + "..." if len(content) > 60 else content,
                        )
                else:
                    call_compressed += tokens

        self._total_original += call_original
        self._total_compressed += call_compressed
        self._total_already_compressed += already_compressed

        if call_original > 0:
            saving_pct = round((1 - call_compressed / call_original) * 100, 1)
            logger.log(
                self.log_level,
                "Condensa: call %d — %d -> %d tokens (-%s%% potential savings)",
                self._calls, call_original, call_compressed, saving_pct,
            )
