"""Tests for the formula parser (KTD7 security + parsing).

Test-first per U3 execution note: parser, security constraints, and cycle
detection tests are written before the engine/recalc worker.
"""

from __future__ import annotations

import pytest

from agentkit.bitable.formula.parser import (
    MAX_AST_DEPTH,
    FormulaDepthExceededError,
    FormulaParseError,
    FormulaSecurityError,
    UnknownFunctionError,
    evaluate_ast,
    parse_formula,
)

ALLOWED = {"SUM", "AVG", "COUNT", "MIN", "MAX", "ABS", "ROUND", "IF", "LEN", "CONCAT"}


# ---------------------------------------------------------------------------
# Parsing happy paths
# ---------------------------------------------------------------------------


def test_parse_simple_arithmetic() -> None:
    tree, mapping, _ = parse_formula("=1+2*3", ALLOWED)
    assert mapping == {}
    result = evaluate_ast(tree, {}, {})
    assert result == 7


def test_parse_strips_equals_prefix() -> None:
    tree1, _, _ = parse_formula("=1+1", ALLOWED)
    tree2, _, _ = parse_formula("1+1", ALLOWED)
    assert evaluate_ast(tree1, {}, {}) == evaluate_ast(tree2, {}, {}) == 2


def test_parse_field_reference() -> None:
    tree, mapping, _ = parse_formula("={field_abc} + 1", ALLOWED)
    assert "field_abc" in mapping.values()
    # Safe name is prefixed with _f_
    safe_name = next(k for k, v in mapping.items() if v == "field_abc")
    result = evaluate_ast(tree, {safe_name: 41}, {})
    assert result == 42


def test_parse_uuid_field_reference() -> None:
    """Field IDs are UUIDs with hyphens — must be substituted to safe names."""
    fid = "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
    tree, mapping, _ = parse_formula(f"={{{fid}}} * 2", ALLOWED)
    # The mapping should have a safe name → original UUID
    assert fid in mapping.values()
    # Evaluate using the safe name (prefixed with _f_)
    safe_name = next(k for k, v in mapping.items() if v == fid)
    assert safe_name.startswith("_f_")
    result = evaluate_ast(tree, {safe_name: 21}, {})
    assert result == 42


def test_parse_string_concatenation() -> None:
    tree, _, _ = parse_formula('="hello" + " " + "world"', ALLOWED)
    assert evaluate_ast(tree, {}, {}) == "hello world"


def test_parse_conditional_ifexp() -> None:
    tree, _, _ = parse_formula("=1 if True else 2", ALLOWED)
    assert evaluate_ast(tree, {}, {}) == 1


def test_parse_comparison() -> None:
    tree, mapping, _ = parse_formula("={f} > 5", ALLOWED)
    safe_name = next(k for k, v in mapping.items() if v == "f")
    assert evaluate_ast(tree, {safe_name: 10}, {}) is True
    assert evaluate_ast(tree, {safe_name: 3}, {}) is False


def test_parse_boolean_ops() -> None:
    tree, _, _ = parse_formula("=True and False", ALLOWED)
    assert evaluate_ast(tree, {}, {}) is False
    tree2, _, _ = parse_formula("=True or False", ALLOWED)
    assert evaluate_ast(tree2, {}, {}) is True


# ---------------------------------------------------------------------------
# Function calls
# ---------------------------------------------------------------------------


def test_parse_function_call_sum() -> None:
    tree, mapping, _ = parse_formula("=SUM({f1})", ALLOWED)
    safe_name = next(k for k, v in mapping.items() if v == "f1")
    result = evaluate_ast(tree, {safe_name: [1, 2, 3]}, {"SUM": sum})
    assert result == 6


def test_parse_function_call_concat() -> None:
    tree, mapping, _ = parse_formula('=CONCAT({f1}, "-", {f2})', ALLOWED)
    safe_f1 = next(k for k, v in mapping.items() if v == "f1")
    safe_f2 = next(k for k, v in mapping.items() if v == "f2")
    result = evaluate_ast(
        tree, {safe_f1: "a", safe_f2: "b"}, {"CONCAT": lambda *a: "".join(str(x) for x in a)}
    )
    assert result == "a-b"


def test_parse_nested_function_calls() -> None:
    tree, _, _ = parse_formula("=ABS(-5) + ROUND(3.7, 0)", ALLOWED)
    funcs = {"ABS": abs, "ROUND": round}
    result = evaluate_ast(tree, {}, funcs)
    assert result == 9  # 5 + 4


# ---------------------------------------------------------------------------
# KTD7 Security — disallowed nodes
# ---------------------------------------------------------------------------


def test_security_rejects_attribute_access() -> None:
    """__import__('os') is rejected — it's a Call to an unregistered function.
    (Attribute access like os.system would be caught by the Attribute node check,
    but __import__ is caught earlier as an unknown function.)"""
    with pytest.raises((FormulaSecurityError, UnknownFunctionError)):
        parse_formula("=__import__('os')", ALLOWED)


def test_security_rejects_attribute_chain() -> None:
    """Attribute access like ''.join([]) is rejected by the Attribute node check."""
    with pytest.raises(FormulaSecurityError):
        parse_formula("=''.join([])", ALLOWED)


def test_security_rejects_lambda() -> None:
    with pytest.raises(FormulaSecurityError):
        parse_formula("=(lambda: 1)()", ALLOWED)


def test_security_rejects_subscript() -> None:
    with pytest.raises(FormulaSecurityError):
        parse_formula("=[1,2,3][0]", ALLOWED)


def test_security_rejects_assignment() -> None:
    """Assignment is a statement, not an expression — rejected at parse stage."""
    with pytest.raises((FormulaSecurityError, FormulaParseError)):
        parse_formula("=x = 1", ALLOWED)


def test_unknown_function_rejected() -> None:
    with pytest.raises(UnknownFunctionError):
        parse_formula("=UNKNOWN(1)", ALLOWED)


def test_eval_function_rejected_if_not_registered() -> None:
    """eval is not in the registry → UnknownFunctionError."""
    with pytest.raises(UnknownFunctionError):
        parse_formula("=eval('1+1')", ALLOWED)


# ---------------------------------------------------------------------------
# Error paths
# ---------------------------------------------------------------------------


def test_parse_error_unbalanced_parens() -> None:
    with pytest.raises(FormulaParseError):
        parse_formula("=(1+2", ALLOWED)


def test_parse_error_empty_formula() -> None:
    with pytest.raises(FormulaParseError):
        parse_formula("=", ALLOWED)


def test_parse_error_empty_string() -> None:
    with pytest.raises(FormulaParseError):
        parse_formula("", ALLOWED)


def test_evaluate_unknown_field_ref_raises() -> None:
    tree, _, _ = parse_formula("={nonexistent} + 1", ALLOWED)
    with pytest.raises(FormulaParseError, match="Unknown field reference"):
        evaluate_ast(tree, {}, {})


# ---------------------------------------------------------------------------
# Mixed aggregate + row context
# ---------------------------------------------------------------------------


def test_mixed_aggregate_and_row_context() -> None:
    """={f1} + SUM({f2}) — row f1 + column f2 sum."""
    tree, mapping, _ = parse_formula("={f1} + SUM({f2})", ALLOWED)
    safe_f1 = next(k for k, v in mapping.items() if v == "f1")
    safe_f2 = next(k for k, v in mapping.items() if v == "f2")
    # f1 is a row value (scalar), f2 is a column value (list)
    result = evaluate_ast(tree, {safe_f1: 10, safe_f2: [1, 2, 3]}, {"SUM": sum})
    assert result == 16  # 10 + 6


# ---------------------------------------------------------------------------
# KTD8: AST depth limit
# ---------------------------------------------------------------------------


def test_depth_limit_allows_normal_formula() -> None:
    """Normal formulas should pass the depth check."""
    tree, _, _ = parse_formula("=1+2+3+4+5", ALLOWED)
    assert tree is not None


def test_depth_limit_rejects_deeply_nested() -> None:
    """Deeply nested expressions should be rejected (KTD8)."""
    # Build a formula deeper than MAX_AST_DEPTH: nested additions
    # Each "x + " adds 1 level of BinOp depth
    deep_formula = "=" + " + ".join(["1"] * (MAX_AST_DEPTH + 5))
    with pytest.raises(FormulaDepthExceededError, match="depth"):
        parse_formula(deep_formula, ALLOWED)


def test_depth_limit_allows_at_boundary() -> None:
    """Formula at exactly the limit should pass."""
    # Build a formula with depth close to but not exceeding MAX_AST_DEPTH
    # A simple sum of N numbers has depth ~N, but ast groups them differently.
    # A safer test: shallow formula always passes.
    tree, _, _ = parse_formula("=SUM({f1}) + IF({f2}, {f3}, {f4})", ALLOWED)
    assert tree is not None


def test_list_literal_evaluated() -> None:
    """List literals [1, 2, 3] should be evaluated as Python lists."""
    tree, _, _ = parse_formula("=[1, 2, 3]", ALLOWED)
    result = evaluate_ast(tree, {}, {})
    assert result == [1, 2, 3]


def test_tuple_literal_evaluated() -> None:
    """Tuple literals (1, 2, 3) should be evaluated as Python tuples."""
    tree, _, _ = parse_formula("=(1, 2, 3)", ALLOWED)
    result = evaluate_ast(tree, {}, {})
    assert result == (1, 2, 3)


# ---------------------------------------------------------------------------
# U3: Cross-table references {rel.target}
# ---------------------------------------------------------------------------


def test_parse_cross_table_reference() -> None:
    """{rel_field.target_field} is parsed into a single Name node (KTD-2)."""
    tree, field_mapping, cross_table_mapping = parse_formula("={f_rel.f_target}", ALLOWED)
    # Single-field mapping should be empty
    assert field_mapping == {}
    # Cross-table mapping should have one entry
    assert len(cross_table_mapping) == 1
    safe_name, (rel_id, target_id) = next(iter(cross_table_mapping.items()))
    assert rel_id == "f_rel"
    assert target_id == "f_target"
    assert safe_name.startswith("_ct_")
    assert "__" in safe_name  # separator between rel and target


def test_parse_cross_table_with_aggregate() -> None:
    """=SUM({rel.target}) — cross-table ref inside an aggregate function."""
    tree, _, cross_table_mapping = parse_formula("=SUM({f_rel.f_target})", ALLOWED)
    assert len(cross_table_mapping) == 1
    safe_name, (rel_id, target_id) = next(iter(cross_table_mapping.items()))
    assert rel_id == "f_rel"
    assert target_id == "f_target"
    # Should evaluate correctly when given a list of values
    result = evaluate_ast(tree, {safe_name: [1, 2, 3]}, {"SUM": sum})
    assert result == 6


def test_parse_mixed_single_and_cross_table_refs() -> None:
    """={f_local} + SUM({f_rel.f_target}) — mixed row + cross-table aggregate."""
    tree, field_mapping, cross_table_mapping = parse_formula(
        "={f_local} + SUM({f_rel.f_target})", ALLOWED
    )
    assert "f_local" in field_mapping.values()
    assert len(cross_table_mapping) == 1
    safe_local = next(k for k, v in field_mapping.items() if v == "f_local")
    safe_cross = next(iter(cross_table_mapping.keys()))
    result = evaluate_ast(
        tree, {safe_local: 10, safe_cross: [1, 2, 3]}, {"SUM": sum}
    )
    assert result == 16


def test_parse_cross_table_uuid_field_ids() -> None:
    """Cross-table refs work with UUID-style field IDs containing hyphens."""
    rel_id = "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
    target_id = "11111111-2222-3333-4444-555555555555"
    tree, _, cross_table_mapping = parse_formula(f"={{{rel_id}.{target_id}}}", ALLOWED)
    assert len(cross_table_mapping) == 1
    _, (parsed_rel, parsed_target) = next(iter(cross_table_mapping.items()))
    assert parsed_rel == rel_id
    assert parsed_target == target_id


def test_cross_table_ref_does_not_match_single_field() -> None:
    """{rel.target} must not be partially matched by _FIELD_REF_RE."""
    # The rel part should not leak into field_mapping
    tree, field_mapping, cross_table_mapping = parse_formula("={f_rel.f_target}", ALLOWED)
    assert field_mapping == {}
    assert len(cross_table_mapping) == 1