"""
Tests for the token_diet.compressor module.
"""

from __future__ import annotations

import pytest

from token_diet.compressor import (
    MIN_TOKENS_TO_COMPRESS,
    CompressionLevel,
    CompressionResult,
    Compressor,
    count_tokens,
    estimate_tokens,
    rule_based_compress,
)


# ---------------------------------------------------------------------------
# Token counting
# ---------------------------------------------------------------------------

class TestEstimateTokens:
    def test_empty_string(self) -> None:
        assert estimate_tokens("") == 1  # max(1, 0)

    def test_short_string(self) -> None:
        # "test" = 4 chars → 1 token
        assert estimate_tokens("test") == 1

    def test_longer_string(self) -> None:
        text = "a" * 400
        assert estimate_tokens(text) == 100

    def test_count_tokens_fallback(self) -> None:
        # Without tiktoken installed, should fall back to estimate
        result = count_tokens("hello world this is a test")
        assert result >= 1


# ---------------------------------------------------------------------------
# Rule-based compression
# ---------------------------------------------------------------------------

class TestRuleBasedCompress:
    def test_removes_please(self) -> None:
        text = "Please summarize this document for me."
        compressed, techniques = rule_based_compress(text, CompressionLevel.BALANCED)
        assert "Please" not in compressed
        assert len(techniques) > 0

    def test_removes_could_you_please(self) -> None:
        text = "Could you please help me write a report."
        compressed, _ = rule_based_compress(text, CompressionLevel.BALANCED)
        assert "Could you please" not in compressed

    def test_in_order_to(self) -> None:
        text = "In order to achieve the goal, we need to act."
        compressed, _ = rule_based_compress(text, CompressionLevel.BALANCED)
        assert "in order to" not in compressed.lower()
        assert "to achieve" in compressed

    def test_at_this_point_in_time(self) -> None:
        text = "At this point in time we should proceed."
        compressed, _ = rule_based_compress(text, CompressionLevel.BALANCED)
        assert "now" in compressed.lower()

    def test_due_to_the_fact_that(self) -> None:
        text = "Due to the fact that it rained, the match was cancelled."
        compressed, _ = rule_based_compress(text, CompressionLevel.BALANCED)
        assert "because" in compressed.lower()

    def test_gentle_does_not_remove_hedging_words(self) -> None:
        """'very', 'quite' etc. should NOT be stripped in gentle mode."""
        text = "Please explain this concept clearly and very simply."
        compressed_gentle, _ = rule_based_compress(text, CompressionLevel.GENTLE)
        compressed_balanced, _ = rule_based_compress(text, CompressionLevel.BALANCED)
        # 'very' should survive gentle but be removed in balanced
        assert "very" in compressed_gentle
        assert "very" not in compressed_balanced

    def test_aggressive_removes_more(self) -> None:
        text = (
            "Can you help me write a summary of the following document "
            "in order to create a comprehensive overview?"
        )
        gentle, _ = rule_based_compress(text, CompressionLevel.GENTLE)
        aggressive, _ = rule_based_compress(text, CompressionLevel.AGGRESSIVE)
        assert len(aggressive) <= len(gentle)

    def test_returns_techniques_list(self) -> None:
        text = (
            "Please summarize the document. In order to do this, "
            "make sure to cover all key points."
        )
        _, techniques = rule_based_compress(text, CompressionLevel.BALANCED)
        assert isinstance(techniques, list)
        assert len(techniques) > 0

    def test_normalises_double_spaces(self) -> None:
        text = "Please   help   me  with   this."
        compressed, _ = rule_based_compress(text, CompressionLevel.GENTLE)
        assert "  " not in compressed

    def test_no_change_when_clean(self) -> None:
        text = "Summarize this document covering main points."
        compressed, _ = rule_based_compress(text, CompressionLevel.GENTLE)
        # Text is already clean; it should be essentially unchanged
        assert compressed.strip() == text.strip()


# ---------------------------------------------------------------------------
# Compressor class
# ---------------------------------------------------------------------------

class TestCompressor:
    def _long_prompt(self, extra: str = "") -> str:
        base = (
            "Please take the following document and create a comprehensive summary "
            "that covers all of the main points and key details. I would like you to "
            "make sure to include every important fact. In order to do this effectively, "
            "please go through the text very carefully. It is important that you do not "
            "miss anything. Could you please also make sure to keep the summary concise? "
            "Due to the fact that the document is long, a good summary is essential. "
        )
        return (base + extra) * 3  # repeat to exceed token threshold

    def test_returns_compression_result(self) -> None:
        c = Compressor()
        result = c.compress(self._long_prompt())
        assert isinstance(result, CompressionResult)

    def test_short_prompt_not_compressed(self) -> None:
        c = Compressor()
        short = "Summarize this."
        result = c.compress(short)
        assert result.compressed_text == short
        assert result.tokens_saved == 0
        assert "nothing to compress" in result.techniques_applied[0].lower()

    def test_reduces_tokens(self) -> None:
        c = Compressor()
        result = c.compress(self._long_prompt())
        assert result.compressed_tokens < result.original_tokens

    def test_percent_saved_positive(self) -> None:
        c = Compressor()
        result = c.compress(self._long_prompt())
        assert result.percent_saved > 0

    def test_original_text_preserved(self) -> None:
        c = Compressor()
        prompt = self._long_prompt()
        result = c.compress(prompt)
        assert result.original_text == prompt

    def test_compression_levels(self) -> None:
        prompt = self._long_prompt()
        gentle = Compressor(level=CompressionLevel.GENTLE).compress(prompt)
        aggressive = Compressor(level=CompressionLevel.AGGRESSIVE).compress(prompt)
        assert aggressive.compressed_tokens <= gentle.compressed_tokens

    def test_no_api_key_uses_rule_based(self) -> None:
        c = Compressor(api_key=None)
        result = c.compress(self._long_prompt())
        assert not result.used_api

    def test_cost_estimate_non_negative(self) -> None:
        c = Compressor()
        result = c.compress(self._long_prompt())
        assert result.cost_saved_gpt4o >= 0

    def test_chars_count(self) -> None:
        c = Compressor()
        prompt = self._long_prompt()
        result = c.compress(prompt)
        assert result.original_chars == len(prompt)
        assert result.compressed_chars == len(result.compressed_text)


# ---------------------------------------------------------------------------
# CompressionResult properties
# ---------------------------------------------------------------------------

class TestCompressionResult:
    def _make_result(self, orig_tokens: int, comp_tokens: int) -> CompressionResult:
        return CompressionResult(
            original_text="original " * orig_tokens,
            compressed_text="compressed " * comp_tokens,
            original_tokens=orig_tokens,
            compressed_tokens=comp_tokens,
        )

    def test_tokens_saved(self) -> None:
        r = self._make_result(100, 60)
        assert r.tokens_saved == 40

    def test_percent_saved(self) -> None:
        r = self._make_result(100, 60)
        assert abs(r.percent_saved - 40.0) < 0.01

    def test_percent_saved_zero_tokens(self) -> None:
        r = self._make_result(0, 0)
        assert r.percent_saved == 0.0

    def test_cost_saved_gpt4o(self) -> None:
        r = self._make_result(1_000_000, 500_000)
        # 500k tokens saved at $5/1M → $2.50
        assert abs(r.cost_saved_gpt4o - 2.50) < 0.01
