"""Tests for model selection suggestion."""

from bilgeai.llm.model_selector import (
    ModelSuggestion,
    _filter_candidates,
    _get_available_vram,
    suggest_model_rule_based,
)


class TestGetAvailableVram:
    def test_single_gpu(self):
        hw = {"gpus": [{"name": "RTX 4090", "vram_gb": 24.0}]}
        assert _get_available_vram(hw) == 24.0

    def test_multi_gpu(self):
        hw = {"gpus": [{"vram_gb": 8.0}, {"vram_gb": 8.0}]}
        assert _get_available_vram(hw) == 16.0

    def test_no_gpu(self):
        assert _get_available_vram({"gpus": []}) == 0.0
        assert _get_available_vram({}) == 0.0


class TestFilterCandidates:
    def test_4bit_filters(self):
        candidates = _filter_candidates(8.0, "4bit")
        # 8GB with 4bit should fit models up to ~27GB min_vram
        assert len(candidates) > 0
        for c in candidates:
            assert c["min_vram_gb"] * 0.25 <= 8.0 * 0.85

    def test_no_vram_returns_empty(self):
        assert _filter_candidates(0.0, "4bit") == []

    def test_large_vram_fits_all(self):
        candidates = _filter_candidates(100.0, "none")
        from bilgeai.llm.model_selector import MODEL_CATALOG

        assert len(candidates) == len(MODEL_CATALOG)


class TestSuggestModelRuleBased:
    def test_very_low_vram_suggests_phi(self):
        hw = {"gpus": [{"vram_gb": 0.5}]}
        ds = {"num_rows": 1000, "columns": ["instruction", "output"], "format": "alpaca"}
        result = suggest_model_rule_based(hw, ds)
        assert isinstance(result, ModelSuggestion)
        assert "phi" in result.recommended_model.lower()
        assert result.source == "rule_engine"

    def test_medium_vram_general_task(self):
        hw = {"gpus": [{"vram_gb": 8.0}]}
        ds = {"num_rows": 5000, "columns": ["instruction", "output"], "format": "alpaca"}
        result = suggest_model_rule_based(hw, ds)
        assert result.recommended_model != ""
        assert result.reason != ""

    def test_high_vram_suggests_large_model(self):
        hw = {"gpus": [{"vram_gb": 80.0}]}
        ds = {"num_rows": 10000, "columns": ["instruction", "output"], "format": "alpaca"}
        result = suggest_model_rule_based(hw, ds)
        # Should suggest a larger model
        assert (
            any(size in result.recommended_model for size in ["70B", "72B", "34b", "8x7B"])
            or result.recommended_model != ""
        )

    def test_code_task_prefers_code_model(self):
        hw = {"gpus": [{"vram_gb": 8.0}]}
        ds = {
            "num_rows": 5000,
            "columns": ["code", "programming_language", "output"],
            "format": "custom",
        }
        result = suggest_model_rule_based(hw, ds)
        assert "code" in result.recommended_model.lower() or "code" in result.reason.lower()

    def test_small_dataset_prefers_small_model(self):
        hw = {"gpus": [{"vram_gb": 24.0}]}
        ds = {"num_rows": 200, "columns": ["instruction", "output"], "format": "alpaca"}
        result = suggest_model_rule_based(hw, ds)
        # With small dataset, should not suggest 70B+
        assert result.recommended_model != ""

    def test_alternatives_provided(self):
        hw = {"gpus": [{"vram_gb": 24.0}]}
        ds = {"num_rows": 5000, "columns": ["instruction", "output"], "format": "alpaca"}
        result = suggest_model_rule_based(hw, ds)
        # Should have at least one alternative when multiple models fit
        assert isinstance(result.alternatives, list)

    def test_no_gpu_returns_smallest(self):
        hw = {"gpus": []}
        ds = {"num_rows": 1000, "columns": ["text"], "format": "unknown"}
        result = suggest_model_rule_based(hw, ds)
        assert "phi" in result.recommended_model.lower()
