# SPDX-License-Identifier: Apache-2.0

"""Unit tests for SamplingParams validation and construction.

These tests exercise every validation rule in SamplingParams._verify_args(),
the greedy-mode clamping applied after validation, the _make_seed helper,
and the deprecation warning for the `logprobs` parameter.

All tests are CPU-only and require no model weights or CUDA device.
"""

from __future__ import annotations

import random
import warnings

import pytest
import torch

from deepinfer.sampling_params import SAMPLING_EPS, SamplingParams


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def make(**kwargs) -> SamplingParams:
    """Construct SamplingParams with defaults overridden by kwargs."""
    return SamplingParams(**kwargs)


# ---------------------------------------------------------------------------
# Default construction
# ---------------------------------------------------------------------------

class TestDefaults:
    def test_default_construction_succeeds(self) -> None:
        sp = SamplingParams()
        assert sp.temperature == pytest.approx(1.0)
        assert sp.top_p == pytest.approx(1.0)
        assert sp.top_k == -1
        assert sp.ignore_eos is False
        assert sp.max_tokens == 16
        assert sp.prompt_logprobs is False
        assert sp.top_logprobs == 0
        assert sp.so_args is None
        assert sp.golden_token_ids is None

    def test_seq_seed_auto_generated(self) -> None:
        sp1 = SamplingParams()
        sp2 = SamplingParams()
        # With overwhelming probability two auto-seeds differ.
        # If they collide the test is still valid (just flaky by chance).
        assert 0 <= sp1.seq_seed < 2**64
        assert 0 <= sp2.seq_seed < 2**64

    def test_repr_contains_key_fields(self) -> None:
        sp = SamplingParams(temperature=0.8, max_tokens=64)
        r = repr(sp)
        assert "temperature=0.8" in r
        assert "max_tokens=64" in r
        assert "top_p=" in r
        assert "seq_seed=" in r


# ---------------------------------------------------------------------------
# temperature
# ---------------------------------------------------------------------------

class TestTemperature:
    def test_valid_positive(self) -> None:
        make(temperature=0.7)

    def test_valid_exactly_one(self) -> None:
        make(temperature=1.0)

    def test_valid_high(self) -> None:
        make(temperature=2.5)

    def test_zero_triggers_greedy_clamp(self) -> None:
        """temperature=0 should not raise; instead it is clamped to SAMPLING_EPS
        and top_p/top_k are reset to greedy defaults."""
        sp = make(temperature=0.0, top_p=0.9, top_k=50)
        assert sp.temperature == pytest.approx(SAMPLING_EPS)
        assert sp.top_p == pytest.approx(1.0)
        assert sp.top_k == -1

    def test_near_zero_below_eps_triggers_greedy_clamp(self) -> None:
        sp = make(temperature=SAMPLING_EPS / 2)
        assert sp.temperature == pytest.approx(SAMPLING_EPS)

    def test_exactly_eps_is_not_clamped(self) -> None:
        sp = make(temperature=SAMPLING_EPS)
        # SAMPLING_EPS is NOT strictly less than SAMPLING_EPS, so no clamp.
        assert sp.temperature == pytest.approx(SAMPLING_EPS)

    def test_negative_raises(self) -> None:
        with pytest.raises(ValueError, match="temperature must be non-negative"):
            make(temperature=-0.1)


# ---------------------------------------------------------------------------
# top_p
# ---------------------------------------------------------------------------

class TestTopP:
    def test_valid_one(self) -> None:
        make(top_p=1.0)

    def test_valid_fraction(self) -> None:
        make(top_p=0.9)

    def test_valid_near_zero(self) -> None:
        make(top_p=1e-9)

    def test_zero_raises(self) -> None:
        with pytest.raises(ValueError, match="top_p must be in"):
            make(top_p=0.0)

    def test_above_one_raises(self) -> None:
        with pytest.raises(ValueError, match="top_p must be in"):
            make(top_p=1.01)

    def test_negative_raises(self) -> None:
        with pytest.raises(ValueError, match="top_p must be in"):
            make(top_p=-0.5)


# ---------------------------------------------------------------------------
# top_k
# ---------------------------------------------------------------------------

class TestTopK:
    def test_disabled(self) -> None:
        make(top_k=-1)

    def test_valid_one(self) -> None:
        make(top_k=1)

    def test_valid_large(self) -> None:
        make(top_k=torch.iinfo(torch.int32).max)

    def test_zero_raises(self) -> None:
        with pytest.raises(ValueError, match="top_k must be -1"):
            make(top_k=0)

    def test_negative_other_than_minus_one_raises(self) -> None:
        with pytest.raises(ValueError, match="top_k must be -1"):
            make(top_k=-2)

    def test_above_int32_max_raises(self) -> None:
        with pytest.raises(ValueError, match="top_k must be -1"):
            make(top_k=torch.iinfo(torch.int32).max + 1)


# ---------------------------------------------------------------------------
# max_tokens
# ---------------------------------------------------------------------------

class TestMaxTokens:
    def test_valid_minimum(self) -> None:
        sp = make(max_tokens=1)
        assert sp.max_tokens == 1

    def test_valid_default(self) -> None:
        sp = make()
        assert sp.max_tokens == 16

    def test_valid_large(self) -> None:
        make(max_tokens=131072)

    def test_zero_raises(self) -> None:
        """max_tokens=0 must raise ValueError rather than being silently coerced to 1."""
        with pytest.raises(ValueError, match="max_tokens must be >= 1"):
            make(max_tokens=0)

    def test_negative_raises(self) -> None:
        with pytest.raises(ValueError, match="max_tokens must be >= 1"):
            make(max_tokens=-1)

    def test_above_int32_max_raises(self) -> None:
        with pytest.raises(ValueError, match="max_tokens must be >= 1"):
            make(max_tokens=torch.iinfo(torch.int32).max + 1)


# ---------------------------------------------------------------------------
# top_logprobs
# ---------------------------------------------------------------------------

class TestTopLogprobs:
    def test_valid_zero(self) -> None:
        make(top_logprobs=0)

    def test_valid_boundary(self) -> None:
        make(top_logprobs=1000)

    def test_negative_raises(self) -> None:
        with pytest.raises(ValueError, match="top_logprobs must be in"):
            make(top_logprobs=-1)

    def test_above_1000_raises(self) -> None:
        with pytest.raises(ValueError, match="top_logprobs must be in"):
            make(top_logprobs=1001)


# ---------------------------------------------------------------------------
# seq_seed
# ---------------------------------------------------------------------------

class TestSeqSeed:
    def test_none_generates_random_seed(self) -> None:
        sp = make(seq_seed=None)
        assert 0 <= sp.seq_seed < 2**64

    def test_explicit_seed_stored(self) -> None:
        sp = make(seq_seed=42)
        assert sp.seq_seed == 42

    def test_seed_clamped_to_uint64(self) -> None:
        # A value larger than uint64 is masked to the lower 64 bits.
        big = (2**64) + 5
        sp = make(seq_seed=big)
        assert sp.seq_seed == 5

    def test_zero_seed(self) -> None:
        sp = make(seq_seed=0)
        assert sp.seq_seed == 0

    def test_max_uint64_seed(self) -> None:
        max_seed = 2**64 - 1
        sp = make(seq_seed=max_seed)
        assert sp.seq_seed == max_seed

    def test_non_int_seed_raises(self) -> None:
        with pytest.raises(TypeError):
            make(seq_seed="not-an-int")  # type: ignore[arg-type]

    def test_make_seed_deterministic_for_same_value(self) -> None:
        assert SamplingParams._make_seed(123) == SamplingParams._make_seed(123)

    def test_make_seed_none_is_random(self) -> None:
        random.seed(0)
        s1 = SamplingParams._make_seed(None)
        random.seed(0)
        s2 = SamplingParams._make_seed(None)
        # Same random state → same result.
        assert s1 == s2


# ---------------------------------------------------------------------------
# logprobs deprecation
# ---------------------------------------------------------------------------

class TestLogprobsDeprecation:
    def test_logprobs_false_no_warning(self) -> None:
        """The default logprobs=False must not emit any warning."""
        with warnings.catch_warnings():
            warnings.simplefilter("error")
            make(logprobs=False)

    def test_logprobs_true_emits_deprecation_warning(self) -> None:
        with pytest.warns(DeprecationWarning, match="logprobs.*deprecated"):
            make(logprobs=True)

    def test_logprobs_true_does_not_affect_top_logprobs(self) -> None:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            sp = make(logprobs=True, top_logprobs=5)
        assert sp.top_logprobs == 5


# ---------------------------------------------------------------------------
# Optional / passthrough fields
# ---------------------------------------------------------------------------

class TestOptionalFields:
    def test_golden_token_ids_stored(self) -> None:
        ids = [1, 2, 3]
        sp = make(golden_token_ids=ids)
        assert sp.golden_token_ids == ids

    def test_golden_token_ids_none_by_default(self) -> None:
        assert make().golden_token_ids is None

    def test_so_args_stored(self) -> None:
        so = {"schema_type": "json", "schema": "{}"}
        sp = make(so_args=so)
        assert sp.so_args == so

    def test_ignore_eos_stored(self) -> None:
        sp = make(ignore_eos=True)
        assert sp.ignore_eos is True

    def test_first_n_token_no_eos_stored(self) -> None:
        sp = make(first_n_token_no_eos=3, eos_id=2)
        assert sp.first_n_token_no_eos == 3
        assert sp.eos_id == 2

    def test_prompt_logprobs_stored(self) -> None:
        sp = make(prompt_logprobs=True)
        assert sp.prompt_logprobs is True


# ---------------------------------------------------------------------------
# Greedy mode interaction
# ---------------------------------------------------------------------------

class TestGreedyMode:
    def test_greedy_overrides_top_k(self) -> None:
        sp = make(temperature=0.0, top_k=100)
        assert sp.top_k == -1

    def test_non_greedy_preserves_top_k(self) -> None:
        sp = make(temperature=0.5, top_k=50)
        assert sp.top_k == 50

    def test_greedy_overrides_top_p(self) -> None:
        sp = make(temperature=0.0, top_p=0.95)
        assert sp.top_p == pytest.approx(1.0)

    def test_non_greedy_preserves_top_p(self) -> None:
        sp = make(temperature=0.5, top_p=0.95)
        assert sp.top_p == pytest.approx(0.95)