"""Tests for speculative decoding functionality (Task3.3)."""

from __future__ import annotations

from dataclasses import dataclass

import pytest

from sagellm_compression.speculative import (
    DraftModel,
    SpeculativeController,
    VerificationResult,
    Verifier,
)


@dataclass
class SequentialDraftModel(DraftModel):
    """Draft model that generates consecutive ids."""

    def propose_tokens(self, context_ids: list[int], n_tokens: int) -> list[int]:
        start = context_ids[-1] if context_ids else 0
        return [start + i + 1 for i in range(n_tokens)]


@dataclass
class PrefixVerifier(Verifier):
    """Verifier that accepts a fixed prefix length."""

    accept_len: int

    def verify_tokens(self, context_ids: list[int], draft_tokens: list[int]) -> VerificationResult:
        accepted = draft_tokens[: self.accept_len]
        if self.accept_len < len(draft_tokens):
            return VerificationResult(
                accepted_tokens=accepted,
                rejected_token=draft_tokens[self.accept_len] + 100,
            )
        return VerificationResult(accepted_tokens=accepted, rejected_token=None)


@dataclass
class EmptyDraftModel(DraftModel):
    """Invalid draft model for fail-fast tests."""

    def propose_tokens(self, context_ids: list[int], n_tokens: int) -> list[int]:
        return []


@dataclass
class InvalidVerifier(Verifier):
    """Invalid verifier for fail-fast tests."""

    def verify_tokens(self, context_ids: list[int], draft_tokens: list[int]) -> VerificationResult:
        wrong_prefix = list(reversed(draft_tokens))
        return VerificationResult(accepted_tokens=wrong_prefix, rejected_token=None)


@pytest.mark.speculative
class TestSpeculativeController:
    """Speculative draft-verify controller tests."""

    def test_generate_all_accepted(self) -> None:
        controller = SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=PrefixVerifier(accept_len=8),
            n_speculative_tokens=4,
        )

        result = controller.generate(input_ids=[10], max_new_tokens=8)

        assert result.generated_tokens == [11, 12, 13, 14, 15, 16, 17, 18]
        assert result.metrics.accepted_tokens_total == 8
        assert result.metrics.rejected_tokens_total == 0
        assert result.metrics.accept_rate == 1.0

    def test_generate_partial_accept_with_fallback(self) -> None:
        controller = SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=PrefixVerifier(accept_len=2),
            n_speculative_tokens=4,
        )

        result = controller.generate(input_ids=[1], max_new_tokens=6)

        assert result.generated_tokens == [2, 3, 104, 105, 106, 207]
        assert result.metrics.accepted_tokens_total == 4
        assert result.metrics.rejected_tokens_total == 2
        assert result.metrics.draft_tokens_total == 7
        assert result.metrics.accept_rate == pytest.approx(4 / 7, abs=1e-9)

    def test_fail_fast_empty_draft(self) -> None:
        controller = SpeculativeController(
            draft_model=EmptyDraftModel(),
            verifier=PrefixVerifier(accept_len=1),
            n_speculative_tokens=3,
        )

        with pytest.raises(ValueError, match="empty token list"):
            controller.generate(input_ids=[1], max_new_tokens=3)

    def test_fail_fast_invalid_accept_prefix(self) -> None:
        controller = SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=InvalidVerifier(),
            n_speculative_tokens=3,
        )

        with pytest.raises(ValueError, match="prefix"):
            controller.generate(input_ids=[1], max_new_tokens=3)

    def test_fail_fast_all_accepted_with_rejected_token(self) -> None:
        class BadVerifier(Verifier):
            def verify_tokens(
                self,
                context_ids: list[int],
                draft_tokens: list[int],
            ) -> VerificationResult:
                return VerificationResult(accepted_tokens=draft_tokens, rejected_token=42)

        controller = SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=BadVerifier(),
            n_speculative_tokens=2,
        )

        with pytest.raises(ValueError, match="must be None"):
            controller.generate(input_ids=[1], max_new_tokens=2)

    def test_metrics_serialization(self) -> None:
        controller = SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=PrefixVerifier(accept_len=1),
            n_speculative_tokens=2,
        )

        result = controller.generate(input_ids=[1], max_new_tokens=3)
        payload = result.to_dict()

        assert "metrics" in payload
        assert payload["metrics"]["accept_rate"] <= 1.0
        assert payload["metrics"]["spec_accept_rate"] == payload["metrics"]["accept_rate"]
        assert payload["metrics"]["estimated_speedup"] >= 1.0


@pytest.mark.speculative
def test_invalid_controller_arguments() -> None:
    """Controller should validate constructor arguments."""
    with pytest.raises(ValueError, match="n_speculative_tokens"):
        SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=PrefixVerifier(accept_len=1),
            n_speculative_tokens=0,
        )

    with pytest.raises(ValueError, match="max_iterations"):
        SpeculativeController(
            draft_model=SequentialDraftModel(),
            verifier=PrefixVerifier(accept_len=1),
            n_speculative_tokens=1,
            max_iterations=0,
        )


@pytest.mark.speculative
def test_invalid_generate_arguments() -> None:
    """Generate should validate runtime arguments."""
    controller = SpeculativeController(
        draft_model=SequentialDraftModel(),
        verifier=PrefixVerifier(accept_len=1),
        n_speculative_tokens=2,
    )

    with pytest.raises(ValueError, match="max_new_tokens"):
        controller.generate(input_ids=[1], max_new_tokens=0)
