"""tests/test_features_step_transformer.py — Unit tests for StepTransformerVerifier / encode_step."""
import pytest
import warnings

pytestmark = pytest.mark.unit

RUNS = [
    {"question": f"Q{i}", "steps": [
        {"thought": "thinking", "action_type": "Search", "action_arg": f"q{i}",
         "observation": f"Result about topic {i}"},
    ], "final_answer": f"A{i}", "correct": i % 2 == 0}
    for i in range(20)
]


class TestImport:
    def test_import_verifier(self):
        from llm_guard.features.step_transformer import StepTransformerVerifier
        assert StepTransformerVerifier is not None

    def test_import_encode_step(self):
        from llm_guard.features.step_transformer import encode_step
        assert encode_step is not None


class TestEncodeStep:
    def test_react_step(self):
        from llm_guard.features.step_transformer import encode_step
        step = {"thought": "thinking", "action_arg": "query", "observation": "result"}
        text = encode_step(step)
        assert isinstance(text, str)
        assert len(text) > 0

    def test_empty_step(self):
        from llm_guard.features.step_transformer import encode_step
        text = encode_step({})
        assert isinstance(text, str)

    def test_langchain_step(self):
        from llm_guard.features.step_transformer import encode_step
        step = {"log": "thinking", "tool_input": "search query", "observation": "result"}
        text = encode_step(step)
        assert isinstance(text, str)


class TestStepTransformerInstantiation:
    def test_default_construction(self):
        from llm_guard.features.step_transformer import StepTransformerVerifier
        v = StepTransformerVerifier()
        assert not v._fitted

    def test_score_raises_unfitted(self):
        from llm_guard.features.step_transformer import StepTransformerVerifier
        v = StepTransformerVerifier()
        with pytest.raises(RuntimeError):
            v.score("Q", [], "A")

    def test_max_steps_stored(self):
        from llm_guard.features.step_transformer import StepTransformerVerifier
        v = StepTransformerVerifier(max_steps=5)
        assert v.max_steps == 5


class TestStepTransformerFit:
    def test_fit_marks_fitted(self):
        """Fit using the sklearn fallback (no torch required)."""
        from llm_guard.features.step_transformer import StepTransformerVerifier
        from unittest.mock import patch, MagicMock
        import numpy as np

        v = StepTransformerVerifier()

        # Patch the _embed_chain methods to avoid real sentence-transformer calls
        embed_matrix = np.random.randn(v.max_steps, 384).astype(float)
        embed_mean = np.random.randn(384).astype(float)

        with patch.object(v, "_embed_chain", return_value=(embed_matrix, 1)):
            with patch.object(v, "_embed_chain_mean", return_value=embed_mean):
                v.fit(RUNS)

        assert v._fitted

    def test_score_after_fit(self):
        from llm_guard.features.step_transformer import StepTransformerVerifier
        from unittest.mock import patch
        import numpy as np

        v = StepTransformerVerifier()
        embed_matrix = np.random.randn(v.max_steps, 384).astype(float)
        embed_mean = np.random.randn(384).astype(float)

        with patch.object(v, "_embed_chain", return_value=(embed_matrix, 1)):
            with patch.object(v, "_embed_chain_mean", return_value=embed_mean):
                v.fit(RUNS)
                score = v.score(RUNS[0]["question"], RUNS[0]["steps"], RUNS[0]["final_answer"])

        assert 0.0 <= score <= 1.0

    def test_save_load(self, tmp_path):
        from llm_guard.features.step_transformer import StepTransformerVerifier
        from unittest.mock import patch
        import numpy as np

        v = StepTransformerVerifier()
        embed_matrix = np.random.randn(v.max_steps, 384).astype(float)
        embed_mean = np.random.randn(384).astype(float)

        with patch.object(v, "_embed_chain", return_value=(embed_matrix, 1)):
            with patch.object(v, "_embed_chain_mean", return_value=embed_mean):
                v.fit(RUNS)

        path = str(tmp_path / "transformer.pkl")
        v.save(path)
        v2 = StepTransformerVerifier.load(path)
        assert v2._fitted
