"""Tests for langchain-condensa integration."""

import pytest
from unittest.mock import MagicMock
from uuid import uuid4

from langchain_core.messages import (
    HumanMessage,
    AIMessage,
    SystemMessage,
    ToolMessage,
)

from langchain_condensa import (
    CondensaTransformer,
    CondensaCallbackHandler,
    compress_message,
    decompress_message,
    compress_messages,
)


# ── CondensaTransformer tests ────────────────────────────────────────

class TestCondensaTransformer:
    """Test the core transformer."""

    def setup_method(self):
        self.t = CondensaTransformer(min_tokens=5)

    def test_compress_long_message(self):
        msg = HumanMessage(content="Search the web for SpaceX news and return the top 5 results formatted as bullet points")
        result = self.t.compress_message(msg)
        assert "[CDN]" in result.content
        assert isinstance(result, HumanMessage)

    def test_compress_short_message_unchanged(self):
        msg = HumanMessage(content="Hi")
        result = self.t.compress_message(msg)
        assert result.content == "Hi"

    def test_compress_preserves_system_by_default(self):
        msg = SystemMessage(content="You are a helpful assistant that searches the web for news")
        result = self.t.compress_message(msg)
        assert result.content == msg.content
        assert "[CDN]" not in result.content

    def test_compress_system_when_disabled(self):
        t = CondensaTransformer(preserve_system=False, min_tokens=5)
        msg = SystemMessage(content="Search the web for SpaceX news and return the top 5 results")
        result = t.compress_message(msg)
        # May or may not compress depending on savings, but should try
        assert isinstance(result, SystemMessage)

    def test_compress_ai_message(self):
        msg = AIMessage(content="Here are the results: I found 5 articles about SpaceX Starship news from the last week")
        result = self.t.compress_message(msg)
        assert isinstance(result, AIMessage)

    def test_compress_stores_metadata(self):
        msg = HumanMessage(content="Search the web for SpaceX news and return the top 5 results formatted as bullet points")
        result = self.t.compress_message(msg)
        meta = result.additional_kwargs.get("condensa")
        if "[CDN]" in result.content:
            assert meta is not None
            assert "original" in meta
            assert "original_tokens" in meta
            assert "compressed_tokens" in meta
            assert "reduction_pct" in meta
            assert meta["compressed_tokens"] < meta["original_tokens"]

    def test_decompress_restores_original(self):
        original = "Search the web for SpaceX news and return the top 5 results formatted as bullet points"
        msg = HumanMessage(content=original)
        compressed = self.t.compress_message(msg)
        decompressed = self.t.decompress_message(compressed)
        # If compression happened, original should be restored from metadata
        if "[CDN]" in compressed.content:
            assert decompressed.content == original

    def test_decompress_uncompressed_passthrough(self):
        msg = HumanMessage(content="Hello world")
        result = self.t.decompress_message(msg)
        assert result.content == "Hello world"

    def test_decompress_tagged_without_metadata(self):
        """Messages tagged [CDN] but without metadata should still decode."""
        msg = HumanMessage(content="[CDN] !:srch SpaceX /n:5")
        result = self.t.decompress_message(msg)
        assert "[CDN]" not in result.content
        assert isinstance(result.content, str)
        assert len(result.content) > 0

    def test_compress_messages_list(self):
        messages = [
            SystemMessage(content="You are a helpful assistant"),
            HumanMessage(content="Search the web for SpaceX news and return the top 5 results"),
            AIMessage(content="Here are the results from the search"),
        ]
        results = self.t.compress_messages(messages)
        assert len(results) == 3
        # System should be preserved
        assert results[0].content == messages[0].content

    def test_decompress_messages_list(self):
        messages = [
            HumanMessage(content="Search the web for SpaceX news and return the top 5 results"),
        ]
        compressed = self.t.compress_messages(messages)
        decompressed = self.t.decompress_messages(compressed)
        assert len(decompressed) == 1

    def test_non_string_content_passthrough(self):
        """Messages with non-string content (e.g., multimodal) pass through."""
        msg = HumanMessage(content=[{"type": "text", "text": "hello"}])
        result = self.t.compress_message(msg)
        assert result.content == msg.content

    def test_no_tag_mode(self):
        t = CondensaTransformer(tag=False, min_tokens=5)
        msg = HumanMessage(content="Search the web for SpaceX news and return the top 5 results formatted as bullet points")
        result = t.compress_message(msg)
        # Should not have [CDN] prefix
        if result.content != msg.content:
            assert not result.content.startswith("[CDN]")

    def test_preserves_message_name(self):
        msg = HumanMessage(content="Search the web for SpaceX news and return the top 5 results", name="AgentA")
        result = self.t.compress_message(msg)
        assert result.name == "AgentA"

    def test_compress_only_if_saves_tokens(self):
        """If compression doesn't save tokens, original is preserved."""
        msg = HumanMessage(content="OK")
        result = self.t.compress_message(msg)
        assert result.content == "OK"


class TestTransformerStats:
    """Test compression statistics."""

    def test_stats_with_compressed_messages(self):
        t = CondensaTransformer(min_tokens=5)
        messages = [
            HumanMessage(content="Search the web for SpaceX news and return the top 5 results"),
            HumanMessage(content="Filter active users and sort by registration date descending"),
        ]
        compressed = t.compress_messages(messages)
        stats = t.get_stats(compressed)

        assert stats["messages_total"] == 2
        assert stats["total_original_tokens"] > 0
        assert stats["total_compressed_tokens"] > 0
        assert stats["total_compressed_tokens"] <= stats["total_original_tokens"]

    def test_stats_with_uncompressed(self):
        t = CondensaTransformer(min_tokens=5)
        messages = [HumanMessage(content="Hi")]
        stats = t.get_stats(messages)
        assert stats["messages_compressed"] == 0
        assert stats["reduction_pct"] == 0.0


class TestTransformerLCEL:
    """Test LCEL runnable integration."""

    def test_as_compressor_returns_runnable(self):
        t = CondensaTransformer()
        runnable = t.as_compressor()
        assert hasattr(runnable, "invoke")

    def test_compressor_runnable_with_list(self):
        t = CondensaTransformer(min_tokens=5)
        runnable = t.as_compressor()
        messages = [
            HumanMessage(content="Search the web for SpaceX news and return the top 5 results"),
        ]
        result = runnable.invoke(messages)
        assert isinstance(result, list)
        assert len(result) == 1

    def test_compressor_runnable_with_dict(self):
        t = CondensaTransformer(min_tokens=5)
        runnable = t.as_compressor()
        input_dict = {
            "messages": [HumanMessage(content="Search for news")],
            "other_key": "preserved",
        }
        result = runnable.invoke(input_dict)
        assert isinstance(result, dict)
        assert "messages" in result
        assert result["other_key"] == "preserved"

    def test_compressor_runnable_passthrough_non_message(self):
        t = CondensaTransformer()
        runnable = t.as_compressor()
        result = runnable.invoke("just a string")
        assert result == "just a string"

    def test_decompressor_runnable(self):
        t = CondensaTransformer(min_tokens=5)
        compressor = t.as_compressor()
        decompressor = t.as_decompressor()

        messages = [HumanMessage(content="Search the web for SpaceX news and return the top 5 results")]
        compressed = compressor.invoke(messages)
        decompressed = decompressor.invoke(compressed)
        assert isinstance(decompressed, list)


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

class TestConvenienceFunctions:
    """Test module-level convenience functions."""

    def test_compress_message_function(self):
        msg = HumanMessage(content="Search the web for SpaceX news and return the top 5 results formatted as bullet points")
        result = compress_message(msg)
        assert isinstance(result, HumanMessage)

    def test_decompress_message_function(self):
        msg = HumanMessage(content="Hello")
        result = decompress_message(msg)
        assert result.content == "Hello"

    def test_compress_messages_function(self):
        messages = [
            HumanMessage(content="Search for news"),
            AIMessage(content="Here are the results"),
        ]
        result = compress_messages(messages)
        assert len(result) == 2


# ── CondensaCallbackHandler tests ────────────────────────────────────

class TestCondensaCallbackHandler:
    """Test the callback handler."""

    def test_initial_stats(self):
        handler = CondensaCallbackHandler()
        stats = handler.stats
        assert stats["total_original_tokens"] == 0
        assert stats["calls"] == 0

    def test_on_chat_model_start_tracks_stats(self):
        handler = CondensaCallbackHandler()
        messages = [[
            HumanMessage(content="Search the web for SpaceX news and return the top 5 results formatted as bullet points"),
        ]]
        handler.on_chat_model_start(
            serialized={},
            messages=messages,
            run_id=uuid4(),
        )
        stats = handler.stats
        assert stats["calls"] == 1
        assert stats["total_original_tokens"] > 0

    def test_multiple_calls_accumulate(self):
        handler = CondensaCallbackHandler()
        messages = [[HumanMessage(content="Search the web for SpaceX news and summarize the results")]]

        handler.on_chat_model_start(serialized={}, messages=messages, run_id=uuid4())
        handler.on_chat_model_start(serialized={}, messages=messages, run_id=uuid4())

        stats = handler.stats
        assert stats["calls"] == 2

    def test_reset(self):
        handler = CondensaCallbackHandler()
        messages = [[HumanMessage(content="Search for SpaceX news")]]
        handler.on_chat_model_start(serialized={}, messages=messages, run_id=uuid4())
        handler.reset()
        assert handler.stats["calls"] == 0
        assert handler.stats["total_original_tokens"] == 0

    def test_already_compressed_detection(self):
        handler = CondensaCallbackHandler()
        messages = [[
            HumanMessage(
                content="[CDN] !:srch SpaceX /n:5",
                additional_kwargs={"condensa": {"original_tokens": 20}},
            ),
        ]]
        handler.on_chat_model_start(serialized={}, messages=messages, run_id=uuid4())
        stats = handler.stats
        assert stats["messages_already_compressed"] == 1

    def test_non_string_content_skipped(self):
        handler = CondensaCallbackHandler()
        messages = [[
            HumanMessage(content=[{"type": "image", "url": "http://example.com/img.png"}]),
        ]]
        handler.on_chat_model_start(serialized={}, messages=messages, run_id=uuid4())
        stats = handler.stats
        assert stats["total_original_tokens"] == 0

    def test_verbose_mode(self, caplog):
        import logging
        handler = CondensaCallbackHandler(verbose=True, log_level=logging.WARNING)
        messages = [[
            HumanMessage(content="Search the web for SpaceX Starship news and return the top 5 results formatted as bullet points"),
        ]]
        with caplog.at_level(logging.WARNING, logger="langchain_condensa"):
            handler.on_chat_model_start(serialized={}, messages=messages, run_id=uuid4())
        # Should have logged per-message detail
        assert any("Condensa" in record.message for record in caplog.records)


# ── Round-trip tests ─────────────────────────────────────────────────

class TestRoundTrip:
    """Test compress -> decompress round-trip preserves meaning."""

    @pytest.mark.parametrize("content", [
        "Search the web for SpaceX news and return the top 5 results",
        "Filter active users, group by region, sort descending",
        "Error: the API returned a 429 rate limit error",
        "Cancel all running tasks and discard partial results",
        "Delegate to AgentC: review the code for security issues",
    ])
    def test_roundtrip_preserves_original(self, content):
        t = CondensaTransformer(min_tokens=5)
        msg = HumanMessage(content=content)
        compressed = t.compress_message(msg)
        decompressed = t.decompress_message(compressed)
        # Original should be restored from metadata
        if compressed.content != content:
            assert decompressed.content == content

    @pytest.mark.parametrize("content", [
        "Search the web for SpaceX news and return the top 5 results",
        "Filter active users, group by region, sort descending",
    ])
    def test_compressed_is_shorter(self, content):
        t = CondensaTransformer(min_tokens=5)
        msg = HumanMessage(content=content)
        compressed = t.compress_message(msg)
        if compressed.content != content:
            from cdn import count_tokens
            assert count_tokens(compressed.content) < count_tokens(content)
