"""Unit tests for utils/llm_sampling_params.py"""

import pytest
from copy import deepcopy
from gmi_ieops.utils.llm_sampling_params import (
    update_vllm_sampling_params, update_sglang_sampling_params,
)

MC = {"eos_token_id": [2]}  # default model config


class TestVllm:

    @pytest.mark.parametrize("cfg_in,key,expected", [
        ({"max_length": 512}, "max_tokens", 512),
        ({"top_k": 0, "temperature": 0.7}, "top_k", -1),
        ({"ngram_penalty": 1.2}, "presence_penalty", 1.2),
        ({"num_results": 3}, "n", 3),
        ({}, "max_tokens", 32768),       # default
        ({"max_tokens": 0}, "max_tokens", 32768),
        ({"top_p": 0}, "top_p", 0.8),
        ({"repetition_penalty": 0}, "repetition_penalty", 1.0),
    ])
    def test_key_mappings_and_defaults(self, cfg_in, key, expected):
        result = update_vllm_sampling_params(cfg_in, MC, "062")
        assert result[key] == expected

    def test_stop_token_ids(self):
        assert update_vllm_sampling_params({}, {"eos_token_id": [2, 50256]}, "062")["stop_token_ids"] == [2, 50256]
        assert update_vllm_sampling_params({}, {"eos_token_id": 2}, "062")["stop_token_ids"] == [2]

    def test_unsupported_keys_removed(self):
        r = update_vllm_sampling_params({"custom": "v", "temperature": 0.9}, MC, "062")
        assert "custom" not in r and "temperature" in r

    def test_beam_search_062_only(self):
        assert update_vllm_sampling_params({"type": "random"}, MC, "062").get("use_beam_search") is False
        assert update_vllm_sampling_params({"type": "beam"}, MC, "062").get("use_beam_search") is True
        assert "use_beam_search" not in update_vllm_sampling_params({"type": "random"}, MC, "063")

    @pytest.mark.parametrize("version", ["062", "063", "0110", "0130"])
    def test_all_versions(self, version):
        r = update_vllm_sampling_params({"temperature": 0.7}, MC, version)
        assert "temperature" in r and "stop_token_ids" in r

    def test_unsupported_version_raises(self):
        with pytest.raises(ValueError, match="not supported"):
            update_vllm_sampling_params({}, MC, "999")

    def test_does_not_mutate_model_config(self):
        mc = {"eos_token_id": 2}
        orig = deepcopy(mc)
        update_vllm_sampling_params({}, mc, "062")
        assert mc == orig


class TestSglang:

    @pytest.mark.parametrize("cfg_in,key,expected", [
        ({}, "max_new_tokens", 1024),
        ({"max_length": 2048}, "max_new_tokens", 2048),
        ({"max_length": 500, "max_tokens": 1000}, "max_new_tokens", 1000),
        ({"top_k": 0}, "top_k", -1),
        ({"ngram_penalty": 1.5}, "presence_penalty", 1.5),
    ])
    def test_key_mappings(self, cfg_in, key, expected):
        assert update_sglang_sampling_params(cfg_in, MC, "035")[key] == expected

    def test_stop_token_ids_wraps_int(self):
        assert update_sglang_sampling_params({}, {"eos_token_id": 42}, "035")["stop_token_ids"] == [42]

    def test_unsupported_keys_removed(self):
        r = update_sglang_sampling_params({"unknown": True, "temperature": 0.5}, MC, "035")
        assert "unknown" not in r and "temperature" in r

    def test_unsupported_version_raises(self):
        with pytest.raises(ValueError, match="not supported"):
            update_sglang_sampling_params({}, MC, "999")
