"""tests/test_scoring_deep_chain_scorer.py — Unit tests for llm_guard.scoring.deep_chain_scorer."""
import pytest
from unittest.mock import MagicMock
import numpy as np

pytestmark = pytest.mark.unit

SHORT_STEPS = [
    {"action_type": "Search", "action_arg": "test query", "observation": "test obs",
     "thought": "thinking"},
    {"action_type": "Finish", "action_arg": "ans", "observation": "", "thought": "done"},
]
LONG_STEPS = [
    {"action_type": "Search", "action_arg": f"query {i}", "observation": f"obs {i}",
     "thought": f"think {i}"}
    for i in range(6)
]


class TestImport:
    def test_import_deep_chain_scorer(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        assert DeepChainScorer is not None

    def test_import_retrieval_cascade(self):
        from llm_guard.scoring.deep_chain_scorer import RetrievalCascadeScorer
        assert RetrievalCascadeScorer is not None

    def test_import_result(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainResult
        assert DeepChainResult is not None


class TestRetrievalCascadeScorer:
    def test_default_construction(self):
        from llm_guard.scoring.deep_chain_scorer import RetrievalCascadeScorer
        scorer = RetrievalCascadeScorer()
        assert scorer.min_search_steps == 2

    def test_score_returns_neutral_for_short_chain(self):
        from llm_guard.scoring.deep_chain_scorer import RetrievalCascadeScorer
        scorer = RetrievalCascadeScorer(min_search_steps=5)
        score = scorer.score(SHORT_STEPS)
        assert score == 0.5

    def test_score_returns_float_for_long_chain(self):
        from llm_guard.scoring.deep_chain_scorer import RetrievalCascadeScorer
        scorer = RetrievalCascadeScorer(min_search_steps=2)
        score = scorer.score(LONG_STEPS)
        assert 0.0 <= score <= 1.0

    def test_score_chain_wrapper(self):
        from llm_guard.scoring.deep_chain_scorer import RetrievalCascadeScorer
        scorer = RetrievalCascadeScorer()
        chain = {"steps": SHORT_STEPS}
        score = scorer.score_chain(chain)
        assert isinstance(score, float)

    def test_empty_steps_returns_neutral(self):
        from llm_guard.scoring.deep_chain_scorer import RetrievalCascadeScorer
        scorer = RetrievalCascadeScorer()
        score = scorer.score([])
        assert score == 0.5


class TestDeepChainScorer:
    def test_default_construction(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        scorer = DeepChainScorer()
        assert scorer.chain_length_threshold == 4

    def test_score_long_chain_uses_cascade(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        scorer = DeepChainScorer(chain_length_threshold=4)
        result = scorer.score("Q", LONG_STEPS, "A")
        assert result.scorer_used == "retrieval_cascade"

    def test_score_result_has_correct_fields(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        scorer = DeepChainScorer(chain_length_threshold=4)
        result = scorer.score("Q", LONG_STEPS, "A")
        assert hasattr(result, "risk_score")
        assert hasattr(result, "scorer_used")
        assert hasattr(result, "n_search_steps")
        assert hasattr(result, "n_steps")

    def test_score_with_mamba_scorer(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        mock_mamba = MagicMock()
        mock_mamba.predict_proba.return_value = [0.7]
        scorer = DeepChainScorer(mamba_scorer=mock_mamba, mamba_threshold=4)
        result = scorer.score("Q", LONG_STEPS, "A")
        assert result.scorer_used == "mamba"
        assert result.risk_score == 0.7

    def test_score_risk_in_range(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        scorer = DeepChainScorer(chain_length_threshold=4)
        result = scorer.score("Q", LONG_STEPS, "A")
        assert 0.0 <= result.risk_score <= 1.0

    def test_empty_steps(self):
        from llm_guard.scoring.deep_chain_scorer import DeepChainScorer
        scorer = DeepChainScorer(chain_length_threshold=4)
        # Empty steps: n_search = 0, goes to SC_OLD path
        result = scorer.score("Q", [], "A")
        assert isinstance(result.risk_score, float)
