"""Tests for LLM-based error analysis."""

from bilgeai.llm.error_analyzer import (
    ErrorAnalysis,
    analyze_error_rule_based,
)


class TestAnalyzeErrorRuleBased:
    def test_oom_error(self):
        result = analyze_error_rule_based("CUDA out of memory. Tried to allocate 2.00 GiB")
        assert result.category == "VRAM"
        assert len(result.suggestions) > 0
        assert result.source == "rule_engine"

    def test_oom_torch_error(self):
        result = analyze_error_rule_based(
            "torch.cuda.OutOfMemoryError",
            traceback_str="RuntimeError: OutOfMemoryError: CUDA out of memory",
        )
        assert result.category == "VRAM"

    def test_nccl_error(self):
        result = analyze_error_rule_based("RuntimeError: NCCL error in: /pytorch/torch/lib")
        assert result.category == "distributed"
        assert any("NCCL" in s for s in result.suggestions)

    def test_device_mismatch(self):
        result = analyze_error_rule_based(
            "RuntimeError: Expected all tensors to be on the same device"
        )
        assert result.category == "device"

    def test_tokenizer_error(self):
        result = analyze_error_rule_based("ValueError: Tokenizer class not found")
        assert result.category == "tokenizer"

    def test_import_error(self):
        result = analyze_error_rule_based("ImportError: No module named 'peft'")
        assert result.category == "dependency"

    def test_key_error_data(self):
        result = analyze_error_rule_based("KeyError: 'instruction'")
        assert result.category == "data"

    def test_nan_loss(self):
        result = analyze_error_rule_based("Training stopped: loss is nan at step 50")
        assert result.category == "training"
        assert any("learning rate" in s.lower() for s in result.suggestions)

    def test_zero_loss(self):
        result = analyze_error_rule_based("Warning: loss is 0.0 for 10 consecutive steps")
        assert result.category == "training"

    def test_deepspeed_error(self):
        result = analyze_error_rule_based("DeepSpeed ZeRO stage 3 error")
        assert result.category == "deepspeed"

    def test_unknown_error(self):
        result = analyze_error_rule_based("Something completely unexpected happened")
        assert result.category == "unknown"
        assert len(result.suggestions) > 0

    def test_returns_error_analysis_type(self):
        result = analyze_error_rule_based("any error")
        assert isinstance(result, ErrorAnalysis)

    def test_traceback_also_searched(self):
        result = analyze_error_rule_based(
            "RuntimeError occurred",
            traceback_str="File train.py line 42\nCUDA out of memory",
        )
        assert result.category == "VRAM"

    def test_case_insensitive_matching(self):
        result = analyze_error_rule_based("cuda OUT OF MEMORY error")
        assert result.category == "VRAM"
