"""
Condition Groups — flat UI-friendly representation of ConditionExpression trees.

Provides bidirectional conversion between recursive ConditionExpression ASTs
and flat ConditionGroup arrays suitable for visual condition builder UIs.

Also provides:
- expression_to_cedar(): render any AST node to valid Cedar condition text
- extract_context_fields(): collect all context field names from an AST
"""

import itertools
import threading
from typing import Any, Literal, TypedDict, Union

from .builder import _condition_to_cedar, _sanitize_identifier, _is_valid_raw_condition
from .rule import ConditionExpression, PolicyCondition

# Sentinel field name for raw (unparseable) conditions
RAW_CONDITION_FIELD = "__raw"

GroupLogic = Literal["and", "or"]


class ConditionGroup(TypedDict):
    """A flat, UI-friendly condition group.

    Each group maps 1:1 to a visual block in the condition builder.
    Groups are implicitly combined with AND at the top level.
    """

    id: str
    logic: GroupLogic
    conditions: list[PolicyCondition]
    negated: bool


# Thread-safe group counter using a lock-protected itertools.count.
# itertools.count itself is not thread-safe for reset, so we wrap it.
_group_counter_lock = threading.Lock()
_group_counter_iter = itertools.count(1)


def _next_group_id() -> str:
    with _group_counter_lock:
        return f"group-{next(_group_counter_iter)}"


def reset_group_counter() -> None:
    """Reset the group counter (for testing)."""
    global _group_counter_iter
    with _group_counter_lock:
        _group_counter_iter = itertools.count(1)


# ---------------------------------------------------------------------------
# expression_to_groups
# ---------------------------------------------------------------------------


def expression_to_groups(expr: ConditionExpression) -> list[ConditionGroup]:
    """Convert a ConditionExpression AST into a flat array of ConditionGroups.

    The top-level AND is split into separate groups. Each OR subtree becomes
    a single group with ``logic: 'or'``. NOT wrappers set ``negated`` to True.
    """
    if expr.get("kind") == "and":
        groups: list[ConditionGroup] = []
        for child in expr.get("children", []):
            groups.extend(_node_to_groups(child))
        return groups
    return _node_to_groups(expr)


def _node_to_groups(expr: ConditionExpression) -> list[ConditionGroup]:
    kind = expr.get("kind")
    if kind == "and":
        conditions = _collect_leaf_conditions(expr.get("children", []))
        return [ConditionGroup(id=_next_group_id(), logic="and", conditions=conditions, negated=False)]
    elif kind == "or":
        conditions = _collect_leaf_conditions(expr.get("children", []))
        return [ConditionGroup(id=_next_group_id(), logic="or", conditions=conditions, negated=False)]
    elif kind == "not":
        child = expr.get("child")
        if child:
            return _not_to_groups(child)
        return []
    elif kind == "raw":
        return [ConditionGroup(
            id=_next_group_id(),
            logic="and",
            conditions=[PolicyCondition(field=RAW_CONDITION_FIELD, operator="eq", value=expr.get("text", ""))],
            negated=False,
        )]
    else:
        # Leaf node
        cond = _leaf_to_condition(expr)
        return [ConditionGroup(id=_next_group_id(), logic="and", conditions=[cond], negated=False)]


def _not_to_groups(inner: ConditionExpression) -> list[ConditionGroup]:
    kind = inner.get("kind")
    if kind == "or":
        conditions = _collect_leaf_conditions(inner.get("children", []))
        return [ConditionGroup(id=_next_group_id(), logic="or", conditions=conditions, negated=True)]
    elif kind == "and":
        conditions = _collect_leaf_conditions(inner.get("children", []))
        return [ConditionGroup(id=_next_group_id(), logic="and", conditions=conditions, negated=True)]
    elif kind == "not":
        child = inner.get("child")
        if child:
            return _node_to_groups(child)
        return []
    elif kind == "raw":
        return [ConditionGroup(
            id=_next_group_id(),
            logic="and",
            conditions=[PolicyCondition(field=RAW_CONDITION_FIELD, operator="eq", value=inner.get("text", ""))],
            negated=True,
        )]
    else:
        cond = _leaf_to_condition(inner)
        return [ConditionGroup(id=_next_group_id(), logic="and", conditions=[cond], negated=True)]


def _is_leaf(expr: ConditionExpression) -> bool:
    return expr.get("kind") in ("comparison", "contains", "like", "has")


def _collect_leaf_conditions(children: list[ConditionExpression]) -> list[PolicyCondition]:
    conditions = []
    for child in children:
        if _is_leaf(child):
            conditions.append(_leaf_to_condition(child))
        else:
            conditions.append(PolicyCondition(
                field=RAW_CONDITION_FIELD,
                operator="eq",
                value=expression_to_cedar(child),
            ))
    return conditions


def _leaf_to_condition(expr: ConditionExpression) -> PolicyCondition:
    kind = expr.get("kind")
    if kind == "comparison":
        return PolicyCondition(field=expr["field"], operator=expr["operator"], value=expr["value"])
    elif kind == "contains":
        return PolicyCondition(field=expr["field"], operator="contains", value=expr["value"])
    elif kind == "like":
        return PolicyCondition(field=expr["field"], operator="like", value=expr["pattern"])
    elif kind == "has":
        return PolicyCondition(field=expr["field"], operator="eq", value=True)
    else:
        return PolicyCondition(field=RAW_CONDITION_FIELD, operator="eq", value=expression_to_cedar(expr))


# ---------------------------------------------------------------------------
# groups_to_expression
# ---------------------------------------------------------------------------


def groups_to_expression(groups: list[ConditionGroup]) -> ConditionExpression:
    """Convert a flat array of ConditionGroups back into a ConditionExpression AST.

    Each group becomes an AND/OR node. If negated, wrapped in NOT.
    Multiple groups are combined with a top-level AND.
    """
    if not groups:
        return {"kind": "and", "children": []}

    expressions = [_group_to_expression(g) for g in groups]

    if len(expressions) == 1:
        return expressions[0]

    return {"kind": "and", "children": expressions}


def _group_to_expression(group: ConditionGroup) -> ConditionExpression:
    children = [_condition_to_expr(c) for c in group["conditions"]]

    if len(children) == 1:
        inner = children[0]
    elif group["logic"] == "or":
        inner: ConditionExpression = {"kind": "or", "children": children}
    else:
        inner = {"kind": "and", "children": children}

    if group["negated"]:
        return {"kind": "not", "child": inner}
    return inner


def _condition_to_expr(cond: PolicyCondition) -> ConditionExpression:
    if cond["field"] == RAW_CONDITION_FIELD:
        return {"kind": "raw", "text": str(cond["value"])}

    op = cond["operator"]
    if op == "contains":
        return {"kind": "contains", "field": cond["field"], "value": cond["value"]}
    elif op == "like":
        return {"kind": "like", "field": cond["field"], "pattern": str(cond["value"])}
    else:
        return {"kind": "comparison", "field": cond["field"], "operator": op, "value": cond["value"]}


# ---------------------------------------------------------------------------
# expression_to_cedar
# ---------------------------------------------------------------------------


def expression_to_cedar(
    expr: ConditionExpression,
    optional_fields: set[str] | None = None,
) -> str:
    """Render any ConditionExpression node to valid Cedar condition text.

    This handles the full AST including AND, OR, NOT, and raw nodes —
    unlike ``_condition_to_cedar()`` which only handles leaf PolicyConditions.
    """
    kind = expr.get("kind")

    if kind in ("comparison", "contains", "like"):
        cond = _leaf_to_condition(expr)
        return _condition_to_cedar(cond, optional_fields)

    if kind == "has":
        field = _sanitize_identifier(expr["field"], "field")
        return f"context has {field}"

    if kind == "and":
        children = expr.get("children", [])
        if not children:
            return "true"
        if len(children) == 1:
            return expression_to_cedar(children[0], optional_fields)
        parts = []
        for c in children:
            s = expression_to_cedar(c, optional_fields)
            if c.get("kind") == "or":
                s = f"({s})"
            parts.append(s)
        return " && ".join(parts)

    if kind == "or":
        children = expr.get("children", [])
        if not children:
            return "false"
        if len(children) == 1:
            return expression_to_cedar(children[0], optional_fields)
        parts = []
        for c in children:
            s = expression_to_cedar(c, optional_fields)
            if c.get("kind") == "and":
                s = f"({s})"
            parts.append(s)
        return " || ".join(parts)

    if kind == "not":
        child = expr.get("child")
        if not child:
            return "false"
        inner = expression_to_cedar(child, optional_fields)
        if _is_leaf(child):
            return f"!{inner}"
        return f"!({inner})"

    if kind == "raw":
        text = expr.get("text", "")
        if _is_valid_raw_condition(text):
            return text
        return "/* invalid raw condition */"

    return "false"


# ---------------------------------------------------------------------------
# extract_context_fields
# ---------------------------------------------------------------------------


def extract_context_fields(expr: ConditionExpression) -> list[str]:
    """Extract all unique context field names referenced in a ConditionExpression tree.

    Returns a sorted list of unique field names.
    """
    fields: set[str] = set()
    _collect_fields(expr, fields)
    return sorted(fields)


def _collect_fields(expr: ConditionExpression, fields: set[str]) -> None:
    kind = expr.get("kind")
    if kind in ("comparison", "contains", "like", "has"):
        fields.add(expr["field"])
    elif kind in ("and", "or"):
        for child in expr.get("children", []):
            _collect_fields(child, fields)
    elif kind == "not":
        child = expr.get("child")
        if child:
            _collect_fields(child, fields)
