"""tests/test_healing_query_rewriter.py — Unit tests for QueryRewriter."""
import pytest
import warnings
from unittest.mock import MagicMock, patch

pytestmark = pytest.mark.unit


def _make_trust_obj(risk=0.75, tier="LOW", failure_mode=None, should_rewrite=True):
    from llm_guard.trust.trust_object import A2ATrustObject
    return A2ATrustObject(
        answer="Some answer",
        risk_score=risk,
        confidence_tier=tier,
        failure_mode=failure_mode,
        step_count=3,
        judge_label=None,
        downstream_hint="rewrite_query",
        should_rewrite=should_rewrite,
    )


class TestImport:
    def test_import_rewriter(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        assert QueryRewriter is not None

    def test_import_result(self):
        from llm_guard.healing.query_rewriter import RewriteResult
        assert RewriteResult is not None


class TestInstantiation:
    def test_no_key_construction(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        r = QueryRewriter(api_key="")
        assert r._risk_threshold == 0.70

    def test_custom_threshold(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        r = QueryRewriter(api_key="", risk_threshold=0.60)
        assert r._risk_threshold == 0.60


class TestShouldRewrite:
    def test_high_risk_should_rewrite(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        r = QueryRewriter(api_key="")
        trust = _make_trust_obj(risk=0.80, tier="LOW", should_rewrite=True)
        assert r.should_rewrite(trust) is True

    def test_low_risk_no_rewrite(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        r = QueryRewriter(api_key="")
        trust = _make_trust_obj(risk=0.30, tier="HIGH", should_rewrite=False)
        assert r.should_rewrite(trust) is False


class TestRewriteIfNeeded:
    def test_high_confidence_returns_empty(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        r = QueryRewriter(api_key="")
        trust = _make_trust_obj(risk=0.30, tier="HIGH", should_rewrite=False)
        result = r.rewrite_if_needed("What is Paris?", trust)
        assert result == []

    def test_no_key_rewrite_returns_list(self):
        from llm_guard.healing.query_rewriter import QueryRewriter
        r = QueryRewriter(api_key="")
        trust = _make_trust_obj(risk=0.80, tier="LOW", should_rewrite=True)
        # With no API key, rewrite() returns fallback variants or empty list
        result = r.rewrite_if_needed("What is the capital?", trust)
        assert isinstance(result, list)


class TestRewriteResultDataclass:
    def test_paraphrase_property(self):
        from llm_guard.healing.query_rewriter import RewriteResult
        r = RewriteResult(
            original_question="Q",
            variants=["para", "decomp", "alt"],
            failure_mode_hint="retrieval_fail",
            triggered_by_tier="LOW",
        )
        assert r.paraphrase == "para"
        assert r.decomposed == "decomp"
        assert r.alternative == "alt"

    def test_empty_variants_fallback(self):
        from llm_guard.healing.query_rewriter import RewriteResult
        r = RewriteResult(
            original_question="Q",
            variants=[],
            failure_mode_hint="",
            triggered_by_tier="LOW",
        )
        assert r.paraphrase == "Q"

    def test_partial_variants_fallback(self):
        from llm_guard.healing.query_rewriter import RewriteResult
        r = RewriteResult(
            original_question="Q",
            variants=["para"],
            failure_mode_hint="",
            triggered_by_tier="LOW",
        )
        assert r.decomposed == "Q"
        assert r.alternative == "Q"
