"""
Highflame Policy Rule Types

Canonical types for policy rules used across all Highflame services.
These types match the TypeScript and Go implementations exactly.
"""

from typing import Any, Literal, TypedDict, Union

from .annotations import (
    CustomAnnotations,
    PolicyAnnotations,
    PolicySeverity,
    generate_rule_id,
)


# Type aliases for policy components
PolicyEffect = Literal["permit", "forbid"]
ConditionOperator = Literal["eq", "neq", "lt", "lte", "gt", "gte", "contains", "in", "like"]


# ---------------------------------------------------------------------------
# Condition Expression Tree
# ---------------------------------------------------------------------------
# Recursive discriminated union representing a parsed Cedar condition body.
# Produced by parseCedarToRules() from the Cedar JSON AST.
# Binary && / || chains are flattened into n-ary and/or with children[].
#
# Since Python TypedDicts can't be recursive, use a plain dict type alias.
#
# Variants (discriminated by "kind"):
#   {"kind": "comparison", "field": str, "operator": ConditionOperator, "value": ...}
#   {"kind": "contains", "field": str, "value": str|int|bool}
#   {"kind": "like", "field": str, "pattern": str}
#   {"kind": "has", "field": str}
#   {"kind": "and", "children": list[ConditionExpression]}
#   {"kind": "or", "children": list[ConditionExpression]}
#   {"kind": "not", "child": ConditionExpression}
#   {"kind": "raw", "text": str}
ConditionExpression = dict[str, Any]


class PolicyCondition(TypedDict):
    """A single condition in a policy rule."""

    # The context key or attribute path (e.g., "tool_name", "threat_count")
    field: str
    # The comparison operator (eq, neq, lt, gt, gte, lte, in, contains, like)
    operator: ConditionOperator
    # The value to compare against (string, number, boolean, or string array)
    value: Union[str, int, float, bool, list[str]]


class PolicyEntity(TypedDict, total=False):
    """
    Principal or resource entity constraint.

    The ``operator`` field controls Cedar scope syntax:

    - ``'eq'`` (default): ``resource == Type::"id"`` — exact match
    - ``'in'``: ``resource in Type::"id"`` — hierarchy match (descendants)

    Use ``'in'`` for container resource types (Account, Project, App) to match
    all descendant entities. Use ``'eq'`` for leaf types (Session, Tool).
    """

    # Entity type (e.g., "Agent", "Tool", "FilePath", "User")
    type: str
    # Optional specific entity ID. If omitted, matches any entity of this type.
    id: str
    # Scope operator: 'eq' for exact match (==), 'in' for hierarchy match (in). Default: 'eq'
    operator: str


class PolicyJSON(TypedDict, total=False):
    """
    Base JSON representation of a policy for storage and editing.
    This matches the TypeScript PolicyJSON interface.
    """

    # Unique identifier for this policy
    id: str
    # Human-readable name/description
    name: str
    # Policy effect (permit or forbid)
    effect: PolicyEffect
    # Principal constraint (optional)
    principal: PolicyEntity | None
    # Action constraint - single action or array of actions
    action: Union[str, list[str]]
    # Resource constraint (optional)
    resource: PolicyEntity | None
    # When clause conditions
    conditions: list[PolicyCondition]
    # Raw condition string (for advanced users)
    rawCondition: str


class PolicyRule(TypedDict, total=False):
    """
    A policy rule with full Cedar annotation support.

    This is the canonical type used across all Highflame services:
    - highflame-studio (TypeScript UI)
    - highflame-authz (Go backend)
    - Any Python services

    Each PolicyRule maps 1:1 to a Cedar policy statement with proper annotations.

    Annotations are embedded in Cedar text:
        @id("rule-001")
        @name("Block critical threats")
        @severity("high")
        @tags("security,baseline")
        @compliance("SOC2")
        forbid(...) when {...};
    """

    # Predefined annotations (embedded in Cedar text)
    annotations: PolicyAnnotations
    # Custom user-defined annotations (embedded in Cedar text)
    customAnnotations: CustomAnnotations | None

    # Policy effect (permit or forbid)
    effect: PolicyEffect
    # Principal constraint (optional)
    principal: PolicyEntity | None
    # Action to which the rule applies (string or list of strings)
    action: Union[str, list[str]]
    # Resource constraint (optional)
    resource: PolicyEntity | None
    # When conditions (optional)
    conditions: list[PolicyCondition]
    # Raw condition string (for advanced users)
    rawCondition: str
    # Recursive condition expression tree from Cedar JSON AST
    conditionExpression: ConditionExpression

    # Whether this rule is active - NOT embedded in Cedar (runtime state)
    enabled: bool
    # Display/evaluation order - NOT embedded in Cedar (runtime state)
    order: int


class LegacyPolicyRule(TypedDict, total=False):
    """
    Legacy policy rule format for backwards compatibility.

    Deprecated: Use PolicyRule with annotations for new code.
    """

    # Unique rule identifier
    id: str
    # Human-readable name
    name: str
    # Optional longer description
    description: str
    # Policy effect (permit or forbid)
    effect: PolicyEffect
    # Principal constraint (optional)
    principal: PolicyEntity | None
    # Action to which the rule applies (string or list of strings)
    action: Union[str, list[str]]
    # Resource constraint (optional)
    resource: PolicyEntity | None
    # When conditions (optional)
    conditions: list[PolicyCondition]
    # Raw condition string (for advanced users)
    rawCondition: str
    # Whether this rule is active
    enabled: bool
    # Display/evaluation order (0-indexed)
    order: int
    # Rule severity for display (critical, high, medium, low)
    severity: PolicySeverity
    # Optional tags for categorization and filtering
    tags: list[str]


def convert_legacy_rule(legacy: LegacyPolicyRule, index: int = 0) -> PolicyRule:
    """Convert a LegacyPolicyRule to the new annotations-based format."""
    rule_id = legacy.get("id") or generate_rule_id()
    name = legacy.get("name") or rule_id

    annotations: PolicyAnnotations = {
        "id": rule_id,
        "name": name,
    }
    if legacy.get("description"):
        annotations["description"] = legacy["description"]
    if legacy.get("severity"):
        annotations["severity"] = legacy["severity"]
    if legacy.get("tags"):
        annotations["tags"] = legacy["tags"]

    rule: PolicyRule = {
        "annotations": annotations,
        "effect": legacy.get("effect", "forbid"),
        "action": legacy.get("action", ""),
        "conditions": legacy.get("conditions", []),
        "enabled": legacy.get("enabled", True),
        "order": legacy.get("order", index),
    }

    if legacy.get("principal"):
        rule["principal"] = legacy["principal"]
    if legacy.get("resource"):
        rule["resource"] = legacy["resource"]
    if legacy.get("rawCondition"):
        rule["rawCondition"] = legacy["rawCondition"]

    return rule


def get_action_string(rule: PolicyRule) -> str:
    """Get the action as a single string (first action if array)."""
    action = rule.get("action")
    if isinstance(action, str):
        return action
    if isinstance(action, list) and len(action) > 0:
        return action[0]
    return ""


def get_action_strings(rule: PolicyRule) -> list[str]:
    """Get all actions as a list of strings."""
    action = rule.get("action")
    if isinstance(action, str):
        return [action]
    if isinstance(action, list):
        return action
    return []


def new_policy_rule(
    id: str,
    name: str,
    effect: PolicyEffect,
    action: Union[str, list[str]],
    enabled: bool = True,
    order: int = 0,
    description: str | None = None,
    principal: PolicyEntity | None = None,
    resource: PolicyEntity | None = None,
    conditions: list[PolicyCondition] | None = None,
    severity: PolicySeverity | None = None,
    tags: list[str] | None = None,
    custom_annotations: CustomAnnotations | None = None,
) -> PolicyRule:
    """Create a new PolicyRule with default values."""
    annotations: PolicyAnnotations = {
        "id": id,
        "name": name,
    }
    if description:
        annotations["description"] = description
    if severity:
        annotations["severity"] = severity
    if tags:
        annotations["tags"] = tags

    rule: PolicyRule = {
        "annotations": annotations,
        "effect": effect,
        "action": action,
        "enabled": enabled,
        "order": order,
        "conditions": conditions or [],
    }
    if principal:
        rule["principal"] = principal
    if resource:
        rule["resource"] = resource
    if custom_annotations:
        rule["customAnnotations"] = custom_annotations
    return rule


def new_policy_entity(
    type: str, id: str | None = None, operator: str | None = None
) -> PolicyEntity:
    """Create a new PolicyEntity.

    Args:
        type: Entity type (e.g., "Agent", "Tool", "FilePath", "User").
        id: Optional specific entity ID.
        operator: Scope operator - 'eq' for exact match (==), 'in' for hierarchy match (in).
            Default: 'eq'.
    """
    entity: PolicyEntity = {"type": type}
    if id:
        entity["id"] = id
    if operator:
        entity["operator"] = operator
    return entity


def new_policy_condition(
    field: str,
    operator: ConditionOperator,
    value: Union[str, int, float, bool, list[str]],
) -> PolicyCondition:
    """Create a new PolicyCondition."""
    return {"field": field, "operator": operator, "value": value}
