"""Tests for CoT acceleration functionality (Task3.5)."""

from __future__ import annotations

import pytest

from sagellm_compression.cot import (
    ChainShorteningConfig,
    CompressionResult,
    CoTAccelerationResult,
    CoTAccelerator,
    CoTDetectionResult,
    CoTDetector,
    CoTMetrics,
    CoTStrategyType,
    CoTTemplateManager,
    HeuristicPromptCompressor,
    ParallelChainsConfig,
    PromptCompressionConfig,
)

# ── Fixtures ──────────────────────────────────────────────────────────


COT_TEXT = (
    "Let's think step by step.\n"
    "Step 1: Identify the variables.\n"
    "Step 2: Set up the equation.\n"
    "Step 3: Solve for x.\n"
    "Therefore, x = 42."
)

NON_COT_TEXT = "What is the capital of France?"


def _make_token_ids(n: int = 200) -> list[int]:
    """Generate n sequential token ids."""
    return list(range(n))


# ── Config Tests ──────────────────────────────────────────────────────


@pytest.mark.cot
class TestCoTConfigs:
    """Test CoT configuration classes."""

    def test_prompt_compression_config_defaults(self) -> None:
        cfg = PromptCompressionConfig()
        assert cfg.strategy == CoTStrategyType.PROMPT_COMPRESSION
        assert cfg.enabled is True
        assert 0.0 < cfg.target_compression_ratio < 1.0
        assert cfg.preserve_last_n_tokens >= 0
        assert cfg.min_segment_tokens >= 1

    def test_chain_shortening_config_defaults(self) -> None:
        cfg = ChainShorteningConfig()
        assert cfg.strategy == CoTStrategyType.CHAIN_SHORTENING
        assert cfg.max_chain_length >= 1
        assert cfg.preserve_first_step is True
        assert cfg.preserve_last_step is True

    def test_parallel_chains_config_defaults(self) -> None:
        cfg = ParallelChainsConfig()
        assert cfg.strategy == CoTStrategyType.PARALLEL_CHAINS
        assert cfg.n_chains >= 2
        assert cfg.voting_strategy in ("majority", "weighted")

    def test_prompt_compression_config_custom(self) -> None:
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.3,
            preserve_last_n_tokens=128,
            min_segment_tokens=32,
            quality_threshold=0.9,
        )
        assert cfg.target_compression_ratio == 0.3
        assert cfg.preserve_last_n_tokens == 128
        assert cfg.min_segment_tokens == 32
        assert cfg.quality_threshold == 0.9

    def test_invalid_quality_threshold(self) -> None:
        with pytest.raises(Exception):
            PromptCompressionConfig(quality_threshold=1.5)

    def test_invalid_compression_ratio(self) -> None:
        with pytest.raises(Exception):
            PromptCompressionConfig(target_compression_ratio=0.0)

    def test_strategy_enum_values(self) -> None:
        assert CoTStrategyType.PROMPT_COMPRESSION == "prompt_compression"
        assert CoTStrategyType.CHAIN_SHORTENING == "chain_shortening"
        assert CoTStrategyType.PARALLEL_CHAINS == "parallel_chains"


# ── Detector Tests ────────────────────────────────────────────────────


@pytest.mark.cot
class TestCoTDetector:
    """Test CoT pattern detection."""

    def test_detects_cot_text(self) -> None:
        detector = CoTDetector()
        result = detector.detect(COT_TEXT)
        assert result.is_cot is True
        assert result.confidence > 0.0
        assert len(result.matched_patterns) >= 2
        assert result.estimated_steps >= 3

    def test_rejects_non_cot_text(self) -> None:
        detector = CoTDetector()
        result = detector.detect(NON_COT_TEXT)
        assert result.is_cot is False

    def test_empty_text(self) -> None:
        detector = CoTDetector()
        result = detector.detect("")
        assert result.is_cot is False
        assert result.confidence == 0.0

    def test_whitespace_only(self) -> None:
        detector = CoTDetector()
        result = detector.detect("   \n\t  ")
        assert result.is_cot is False

    def test_custom_thresholds(self) -> None:
        strict = CoTDetector(min_confidence=0.9, min_patterns=5)
        result = strict.detect(COT_TEXT)
        # With very strict thresholds, detection may fail even for CoT text
        assert isinstance(result, CoTDetectionResult)

    def test_detection_result_fields(self) -> None:
        detector = CoTDetector()
        result = detector.detect(COT_TEXT)
        assert isinstance(result.matched_patterns, list)
        assert isinstance(result.step_boundaries, list)
        assert isinstance(result.estimated_steps, int)
        assert 0.0 <= result.confidence <= 1.0

    def test_invalid_min_confidence(self) -> None:
        with pytest.raises(ValueError, match="min_confidence"):
            CoTDetector(min_confidence=-0.1)

    def test_invalid_min_patterns(self) -> None:
        with pytest.raises(ValueError, match="min_patterns"):
            CoTDetector(min_patterns=0)

    def test_various_cot_markers(self) -> None:
        """Test different CoT-style texts are detected."""
        texts = [
            "Let's think step by step. First, we need to... Therefore, the answer is 5.",
            "Step 1: Read the problem. Step 2: Analyze. Thus we conclude...",
            "Reasoning: Because A implies B, and B implies C, so the answer is C.",
        ]
        detector = CoTDetector()
        for text in texts:
            result = detector.detect(text)
            assert result.is_cot is True, f"Failed to detect CoT in: {text!r}"


# ── Prompt Compressor Tests ──────────────────────────────────────────


@pytest.mark.cot
class TestPromptCompressor:
    """Test prompt compression."""

    def test_heuristic_compressor_reduces_tokens(self) -> None:
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.5,
            preserve_last_n_tokens=0,
            min_segment_tokens=8,
        )
        compressor = HeuristicPromptCompressor(cfg)
        tokens = _make_token_ids(200)
        result = compressor.compress(tokens)
        assert result.compressed_tokens < result.original_tokens
        assert result.compression_ratio <= 0.7  # some flexibility
        assert result.removed_segment_count > 0
        assert result.tokens_saved > 0

    def test_short_prompt_no_compression(self) -> None:
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.5,
            min_segment_tokens=8,
        )
        compressor = HeuristicPromptCompressor(cfg)
        tokens = _make_token_ids(5)
        result = compressor.compress(tokens)
        # 5 tokens * 0.5 = 2.5 → target 2 tokens
        # But with min_segment_tokens=8, the 5-token segment is one unit
        assert isinstance(result, CompressionResult)

    def test_empty_tokens_raises(self) -> None:
        cfg = PromptCompressionConfig()
        compressor = HeuristicPromptCompressor(cfg)
        with pytest.raises(ValueError, match="empty"):
            compressor.compress([])

    def test_preserve_last_n_tokens(self) -> None:
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.3,
            preserve_last_n_tokens=32,
            min_segment_tokens=8,
        )
        compressor = HeuristicPromptCompressor(cfg)
        tokens = _make_token_ids(200)
        result = compressor.compress(tokens)
        # Tail tokens should be preserved
        compressed_flat: list[int] = []
        for seg in result.compressed_segments:
            compressed_flat.extend(seg)
        # The last 32 tokens of original should appear at the end
        assert compressed_flat[-32:] == tokens[-32:]

    def test_compression_result_properties(self) -> None:
        r = CompressionResult(
            original_tokens=100,
            compressed_tokens=60,
            compression_ratio=0.6,
            compressed_segments=[[1, 2, 3]],
            removed_segment_count=5,
        )
        assert r.tokens_saved == 40

    def test_heuristic_score_segment(self) -> None:
        cfg = PromptCompressionConfig()
        compressor = HeuristicPromptCompressor(cfg)
        # Later segment should score higher (position bias)
        score_early = compressor.score_segment([1, 2, 3], position=0, total=10)
        score_late = compressor.score_segment([1, 2, 3], position=9, total=10)
        assert score_late > score_early


# ── Metrics Tests ─────────────────────────────────────────────────────


@pytest.mark.cot
class TestCoTMetrics:
    """Test CoT metrics."""

    def test_defaults(self) -> None:
        m = CoTMetrics()
        assert m.tokens_saved == 0
        assert m.step_reduction_ratio == 0.0

    def test_tokens_saved(self) -> None:
        m = CoTMetrics(original_prompt_tokens=100, compressed_prompt_tokens=60)
        assert m.tokens_saved == 40

    def test_step_reduction(self) -> None:
        m = CoTMetrics(original_reasoning_steps=10, retained_reasoning_steps=5)
        assert m.step_reduction_ratio == pytest.approx(0.5)

    def test_step_reduction_zero_original(self) -> None:
        m = CoTMetrics(original_reasoning_steps=0, retained_reasoning_steps=0)
        assert m.step_reduction_ratio == 0.0

    def test_to_dict(self) -> None:
        m = CoTMetrics(
            original_prompt_tokens=100,
            compressed_prompt_tokens=60,
            original_reasoning_steps=10,
            retained_reasoning_steps=5,
            compression_ratio=0.6,
            estimated_speedup=1.5,
            quality_score=0.95,
        )
        d = m.to_dict()
        assert d["tokens_saved"] == 40
        assert d["step_reduction_ratio"] == pytest.approx(0.5)
        assert d["compression_ratio"] == 0.6
        assert d["estimated_speedup"] == 1.5


# ── Accelerator Tests ─────────────────────────────────────────────────


@pytest.mark.cot
class TestCoTAccelerator:
    """Test CoT accelerator orchestration."""

    def test_prompt_compression_end_to_end(self) -> None:
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.5,
            preserve_last_n_tokens=0,
            min_segment_tokens=8,
        )
        accel = CoTAccelerator(config=cfg)
        tokens = _make_token_ids(200)
        result = accel.accelerate(tokens, text=COT_TEXT)
        assert not result.skipped
        assert result.strategy_applied == CoTStrategyType.PROMPT_COMPRESSION
        assert result.compressed_token_ids is not None
        assert len(result.compressed_token_ids) < len(tokens)
        assert result.metrics.estimated_speedup >= 1.0
        assert result.metrics.compression_ratio < 1.0

    def test_acceleration_skipped_when_not_cot(self) -> None:
        cfg = PromptCompressionConfig()
        accel = CoTAccelerator(config=cfg)
        tokens = _make_token_ids(200)
        result = accel.accelerate(tokens, text=NON_COT_TEXT)
        assert result.skipped is True
        assert result.compressed_token_ids is None
        assert result.strategy_applied is None

    def test_acceleration_skipped_when_disabled(self) -> None:
        cfg = PromptCompressionConfig(enabled=False)
        accel = CoTAccelerator(config=cfg)
        tokens = _make_token_ids(100)
        result = accel.accelerate(tokens, text=COT_TEXT)
        assert result.skipped is True

    def test_empty_tokens_raises(self) -> None:
        cfg = PromptCompressionConfig()
        accel = CoTAccelerator(config=cfg)
        with pytest.raises(ValueError, match="empty"):
            accel.accelerate([], text=COT_TEXT)

    def test_none_config_raises(self) -> None:
        with pytest.raises(ValueError, match="config"):
            CoTAccelerator(config=None)  # type: ignore[arg-type]

    def test_parallel_chains_strategy(self) -> None:
        cfg = ParallelChainsConfig(n_chains=4)
        accel = CoTAccelerator(config=cfg)
        tokens = _make_token_ids(100)
        result = accel.accelerate(tokens, text=COT_TEXT)
        assert not result.skipped
        assert result.strategy_applied == CoTStrategyType.PARALLEL_CHAINS
        assert result.metrics.estimated_speedup >= 1.0

    def test_chain_shortening_no_op_when_few_steps(self) -> None:
        cfg = ChainShorteningConfig(max_chain_length=20)
        accel = CoTAccelerator(config=cfg)
        tokens = _make_token_ids(100)
        # COT_TEXT has ~4 steps, max_chain_length=20 → no shortening
        result = accel.accelerate(tokens, text=COT_TEXT)
        assert not result.skipped
        assert result.strategy_applied == CoTStrategyType.CHAIN_SHORTENING
        # Should keep all tokens
        assert result.compressed_token_ids is not None
        assert len(result.compressed_token_ids) == len(tokens)

    def test_custom_detector(self) -> None:
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.5,
            preserve_last_n_tokens=0,
            min_segment_tokens=8,
        )
        detector = CoTDetector(min_confidence=0.0, min_patterns=1)
        accel = CoTAccelerator(config=cfg, detector=detector)
        tokens = _make_token_ids(200)
        # Even semi-CoT text should trigger with low thresholds
        result = accel.accelerate(tokens, text="First, let's think. Therefore the answer is 5.")
        assert isinstance(result, CoTAccelerationResult)

    def test_acceleration_result_detection_field(self) -> None:
        cfg = PromptCompressionConfig()
        accel = CoTAccelerator(config=cfg)
        tokens = _make_token_ids(100)
        result = accel.accelerate(tokens, text=COT_TEXT)
        assert isinstance(result.detection, CoTDetectionResult)
        assert isinstance(result.metrics, CoTMetrics)


# ── Template Tests ────────────────────────────────────────────────────


@pytest.mark.cot
class TestCoTTemplateManager:
    """Test CoT template manager."""

    def test_load_nonexistent_raises(self) -> None:
        mgr = CoTTemplateManager(template_dir="/nonexistent")
        with pytest.raises(FileNotFoundError):
            mgr.load_template("anything")

    def test_render_simple(self) -> None:
        mgr = CoTTemplateManager()
        rendered = mgr.render("Hello {name}!", name="World")
        assert rendered == "Hello World!"

    def test_render_missing_var_raises(self) -> None:
        mgr = CoTTemplateManager()
        with pytest.raises(KeyError, match="Missing template variable"):
            mgr.render("{missing_var}")


# ── Integration / Import Tests ────────────────────────────────────────


@pytest.mark.cot
class TestCoTIntegration:
    """Integration tests for the CoT module."""

    def test_top_level_imports(self) -> None:
        """All CoT symbols should be importable from the top-level package."""
        import sagellm_compression

        for name in [
            "CoTAcceleratorConfig",
            "CoTStrategyType",
            "PromptCompressionConfig",
            "ChainShorteningConfig",
            "ParallelChainsConfig",
            "CoTDetector",
            "CoTDetectionResult",
            "PromptCompressor",
            "HeuristicPromptCompressor",
            "CompressionResult",
            "CoTAccelerator",
            "CoTAccelerationResult",
            "CoTMetrics",
            "CoTTemplateManager",
        ]:
            assert hasattr(sagellm_compression, name), f"Missing: {name}"

    def test_submodule_imports(self) -> None:
        """All CoT symbols should be importable from the cot submodule."""
        from sagellm_compression import cot  # noqa: F401

        for name in [
            "ChainShorteningConfig",
            "CompressionResult",
            "CoTAccelerationResult",
            "CoTAccelerator",
            "CoTAcceleratorConfig",
            "CoTDetectionResult",
            "CoTDetector",
            "CoTMetrics",
            "CoTStrategyType",
            "CoTTemplateManager",
            "HeuristicPromptCompressor",
            "ParallelChainsConfig",
            "PromptCompressionConfig",
            "PromptCompressor",
        ]:
            assert hasattr(cot, name), f"Missing from cot submodule: {name}"

    def test_speedup_at_least_1_3x_on_long_prompt(self) -> None:
        """Acceptance criterion: CoT scenario speedup >= 1.3x."""
        cfg = PromptCompressionConfig(
            target_compression_ratio=0.3,
            preserve_last_n_tokens=64,
            min_segment_tokens=16,
        )
        accel = CoTAccelerator(config=cfg)
        # Simulate a long 4K-token prompt with CoT reasoning text
        tokens = _make_token_ids(4096)
        text = (
            "Let's think step by step.\n"
            + "\n".join(f"Step {i}: Analyzing factor {i}." for i in range(1, 21))
            + "\nTherefore, the final answer is 42."
        )
        result = accel.accelerate(tokens, text=text)
        assert not result.skipped, "Should not skip — text is CoT"
        assert result.metrics.estimated_speedup >= 1.3, (
            f"Speedup {result.metrics.estimated_speedup:.2f}x < 1.3x acceptance threshold"
        )
