"""
Condition Groups Tests

Tests cover:
1. expression_to_groups — AST → flat ConditionGroup[]
2. groups_to_expression — ConditionGroup[] → AST
3. expression_to_cedar — AST → Cedar condition text
4. extract_context_fields — AST → unique field names
5. Round-trip consistency
"""

import pytest
from highflame_policy.condition_groups import (
    ConditionGroup,
    expression_to_groups,
    groups_to_expression,
    expression_to_cedar,
    extract_context_fields,
    reset_group_counter,
    RAW_CONDITION_FIELD,
)


@pytest.fixture(autouse=True)
def _reset_counter():
    """Reset group counter before each test for deterministic IDs."""
    reset_group_counter()


# Helpers
def comp(field: str, op: str, value) -> dict:
    return {"kind": "comparison", "field": field, "operator": op, "value": value}


def contains(field: str, value) -> dict:
    return {"kind": "contains", "field": field, "value": value}


def like(field: str, pattern: str) -> dict:
    return {"kind": "like", "field": field, "pattern": pattern}


def has(field: str) -> dict:
    return {"kind": "has", "field": field}


# =============================================================================
# expression_to_groups
# =============================================================================


class TestExpressionToGroups:
    def test_simple_and(self):
        expr = {
            "kind": "and",
            "children": [
                comp("injection_score", "gt", 50),
                comp("contains_secrets", "eq", True),
            ],
        }
        groups = expression_to_groups(expr)
        assert len(groups) == 2
        assert groups[0]["logic"] == "and"
        assert groups[0]["negated"] is False
        assert groups[0]["conditions"] == [
            {"field": "injection_score", "operator": "gt", "value": 50}
        ]
        assert groups[1]["conditions"] == [
            {"field": "contains_secrets", "operator": "eq", "value": True}
        ]

    def test_or_group(self):
        expr = {
            "kind": "or",
            "children": [
                comp("injection_score", "gt", 50),
                comp("jailbreak_score", "gt", 80),
            ],
        }
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["logic"] == "or"
        assert groups[0]["negated"] is False
        assert len(groups[0]["conditions"]) == 2

    def test_mixed_and_or(self):
        expr = {
            "kind": "and",
            "children": [
                comp("contains_secrets", "eq", True),
                {
                    "kind": "or",
                    "children": [
                        comp("injection_score", "gt", 50),
                        comp("jailbreak_score", "gt", 80),
                    ],
                },
            ],
        }
        groups = expression_to_groups(expr)
        assert len(groups) == 2
        assert groups[0]["logic"] == "and"
        assert len(groups[0]["conditions"]) == 1
        assert groups[0]["conditions"][0]["field"] == "contains_secrets"
        assert groups[1]["logic"] == "or"
        assert len(groups[1]["conditions"]) == 2

    def test_not_wrapping_leaf(self):
        expr = {"kind": "not", "child": comp("injection_score", "gt", 50)}
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["negated"] is True
        assert groups[0]["logic"] == "and"
        assert len(groups[0]["conditions"]) == 1
        assert groups[0]["conditions"][0]["field"] == "injection_score"

    def test_not_wrapping_or(self):
        expr = {
            "kind": "not",
            "child": {
                "kind": "or",
                "children": [
                    comp("injection_score", "gt", 50),
                    comp("jailbreak_score", "gt", 80),
                ],
            },
        }
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["logic"] == "or"
        assert groups[0]["negated"] is True
        assert len(groups[0]["conditions"]) == 2

    def test_not_wrapping_and(self):
        expr = {
            "kind": "and",
            "children": [
                {
                    "kind": "not",
                    "child": {
                        "kind": "and",
                        "children": [
                            comp("injection_score", "eq", 1),
                            comp("jailbreak_score", "eq", 2),
                        ],
                    },
                },
            ],
        }
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["logic"] == "and"
        assert groups[0]["negated"] is True
        assert len(groups[0]["conditions"]) == 2

    def test_raw_fallback(self):
        expr = {"kind": "raw", "text": "context.x.someCustomFn()"}
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["conditions"][0]["field"] == RAW_CONDITION_FIELD
        assert groups[0]["conditions"][0]["value"] == "context.x.someCustomFn()"

    def test_single_leaf(self):
        expr = comp("tool_risk_score", "gte", 90)
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["conditions"] == [
            {"field": "tool_risk_score", "operator": "gte", "value": 90}
        ]

    def test_contains_and_like(self):
        expr = {
            "kind": "and",
            "children": [
                contains("secret_types", "aws_access_key"),
                like("tool_name", "shell*"),
            ],
        }
        groups = expression_to_groups(expr)
        assert len(groups) == 2
        assert groups[0]["conditions"][0] == {
            "field": "secret_types",
            "operator": "contains",
            "value": "aws_access_key",
        }
        assert groups[1]["conditions"][0] == {
            "field": "tool_name",
            "operator": "like",
            "value": "shell*",
        }

    def test_has_expression(self):
        expr = has("tool_name")
        groups = expression_to_groups(expr)
        assert len(groups) == 1
        assert groups[0]["conditions"][0] == {
            "field": "tool_name",
            "operator": "eq",
            "value": True,
        }


# =============================================================================
# groups_to_expression
# =============================================================================


class TestGroupsToExpression:
    def test_empty_groups(self):
        expr = groups_to_expression([])
        assert expr == {"kind": "and", "children": []}

    def test_single_and_group(self):
        expr = groups_to_expression(
            [
                {
                    "id": "g1",
                    "logic": "and",
                    "conditions": [
                        {"field": "a", "operator": "gt", "value": 50},
                        {"field": "b", "operator": "eq", "value": True},
                    ],
                    "negated": False,
                }
            ]
        )
        assert expr["kind"] == "and"
        assert len(expr.get("children", [])) == 2

    def test_single_or_group(self):
        expr = groups_to_expression(
            [
                {
                    "id": "g1",
                    "logic": "or",
                    "conditions": [
                        {"field": "a", "operator": "gt", "value": 50},
                        {"field": "b", "operator": "gt", "value": 80},
                    ],
                    "negated": False,
                }
            ]
        )
        assert expr["kind"] == "or"

    def test_negated_group(self):
        expr = groups_to_expression(
            [
                {
                    "id": "g1",
                    "logic": "and",
                    "conditions": [{"field": "a", "operator": "gt", "value": 50}],
                    "negated": True,
                }
            ]
        )
        assert expr["kind"] == "not"

    def test_raw_sentinel_roundtrip(self):
        expr = groups_to_expression(
            [
                {
                    "id": "g1",
                    "logic": "and",
                    "conditions": [
                        {
                            "field": RAW_CONDITION_FIELD,
                            "operator": "eq",
                            "value": "context.x.foo()",
                        }
                    ],
                    "negated": False,
                }
            ]
        )
        assert expr["kind"] == "raw"
        assert expr.get("text") == "context.x.foo()"


# =============================================================================
# expression_to_cedar
# =============================================================================


class TestExpressionToCedar:
    def test_comparison(self):
        assert (
            expression_to_cedar(comp("injection_score", "gt", 80))
            == "context.injection_score > 80"
        )

    def test_contains(self):
        assert (
            expression_to_cedar(contains("secret_types", "aws_key"))
            == 'context.secret_types.contains("aws_key")'
        )

    def test_like(self):
        assert (
            expression_to_cedar(like("tool_name", "shell*"))
            == 'context.tool_name like "shell*"'
        )

    def test_has(self):
        assert expression_to_cedar(has("tool_name")) == "context has tool_name"

    def test_and(self):
        expr = {
            "kind": "and",
            "children": [
                comp("a", "gt", 50),
                comp("b", "eq", True),
            ],
        }
        assert expression_to_cedar(expr) == "context.a > 50 && context.b == true"

    def test_or(self):
        expr = {
            "kind": "or",
            "children": [
                comp("a", "gt", 50),
                comp("b", "gt", 80),
            ],
        }
        assert expression_to_cedar(expr) == "context.a > 50 || context.b > 80"

    def test_and_with_nested_or(self):
        expr = {
            "kind": "and",
            "children": [
                comp("a", "eq", True),
                {
                    "kind": "or",
                    "children": [
                        comp("b", "gt", 50),
                        comp("c", "gt", 80),
                    ],
                },
            ],
        }
        assert (
            expression_to_cedar(expr)
            == "context.a == true && (context.b > 50 || context.c > 80)"
        )

    def test_not_leaf(self):
        expr = {"kind": "not", "child": comp("a", "gt", 50)}
        assert expression_to_cedar(expr) == "!context.a > 50"

    def test_not_compound(self):
        expr = {
            "kind": "not",
            "child": {
                "kind": "or",
                "children": [
                    comp("a", "gt", 50),
                    comp("b", "gt", 80),
                ],
            },
        }
        assert expression_to_cedar(expr) == "!(context.a > 50 || context.b > 80)"

    def test_empty_and(self):
        assert expression_to_cedar({"kind": "and", "children": []}) == "true"

    def test_empty_or(self):
        assert expression_to_cedar({"kind": "or", "children": []}) == "false"

    def test_raw_valid(self):
        assert (
            expression_to_cedar({"kind": "raw", "text": "context.x == 1"})
            == "context.x == 1"
        )

    def test_with_optional_fields(self):
        optional = {"tool_name"}
        assert (
            expression_to_cedar(comp("tool_name", "eq", "shell"), optional)
            == 'context has tool_name && context.tool_name == "shell"'
        )

    def test_string_escaping(self):
        assert (
            expression_to_cedar(comp("path", "eq", 'he said "hello"'))
            == 'context.path == "he said \\"hello\\""'
        )


# =============================================================================
# extract_context_fields
# =============================================================================


class TestExtractContextFields:
    def test_collects_fields_from_nested_ast(self):
        expr = {
            "kind": "and",
            "children": [
                comp("injection_score", "gt", 80),
                {
                    "kind": "or",
                    "children": [
                        comp("jailbreak_score", "gt", 70),
                        contains("secret_types", "aws_key"),
                    ],
                },
                {"kind": "not", "child": has("tool_name")},
                like("pattern_type", "exfil*"),
            ],
        }
        assert extract_context_fields(expr) == [
            "injection_score",
            "jailbreak_score",
            "pattern_type",
            "secret_types",
            "tool_name",
        ]

    def test_deduplicates(self):
        expr = {
            "kind": "and",
            "children": [
                comp("injection_score", "gt", 80),
                comp("injection_score", "lt", 100),
            ],
        }
        assert extract_context_fields(expr) == ["injection_score"]

    def test_raw_returns_empty(self):
        expr = {"kind": "raw", "text": "context.something == 1"}
        assert extract_context_fields(expr) == []

    def test_empty_and(self):
        assert extract_context_fields({"kind": "and", "children": []}) == []


# =============================================================================
# Round-trip
# =============================================================================


class TestRoundTrip:
    def test_simple_and_roundtrip(self):
        original = {
            "kind": "and",
            "children": [
                comp("a", "gt", 50),
                comp("b", "eq", True),
            ],
        }
        groups = expression_to_groups(original)
        rebuilt = groups_to_expression(groups)
        assert rebuilt["kind"] == "and"
        assert len(rebuilt.get("children", [])) == 2
        assert rebuilt["children"][0]["kind"] == "comparison"
        assert rebuilt["children"][1]["kind"] == "comparison"

    def test_or_roundtrip(self):
        original = {
            "kind": "or",
            "children": [
                comp("a", "gt", 50),
                comp("b", "gt", 80),
            ],
        }
        groups = expression_to_groups(original)
        rebuilt = groups_to_expression(groups)
        assert rebuilt["kind"] == "or"
        assert len(rebuilt.get("children", [])) == 2

    def test_cedar_equivalent_after_roundtrip(self):
        original = {
            "kind": "and",
            "children": [
                comp("injection_score", "gt", 80),
                {
                    "kind": "or",
                    "children": [
                        comp("violence_score", "gt", 90),
                        comp("hate_speech_score", "gt", 90),
                    ],
                },
            ],
        }
        original_cedar = expression_to_cedar(original)
        groups = expression_to_groups(original)
        rebuilt = groups_to_expression(groups)
        rebuilt_cedar = expression_to_cedar(rebuilt)
        assert rebuilt_cedar == original_cedar
