"""
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


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


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."""

    # Entity type (e.g., "Agent", "Tool", "FilePath", "User")
    type: str
    # Optional specific entity ID. If omitted, matches any entity of this type.
    id: 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 UI/storage metadata.
    Extends PolicyJSON with fields needed for UI editing and database storage.

    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.
    """

    # Unique rule identifier
    id: str
    # Human-readable name (becomes @id annotation in Cedar)
    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 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,
) -> PolicyRule:
    """Create a new PolicyRule with default values."""
    rule: PolicyRule = {
        "id": id,
        "name": name,
        "effect": effect,
        "action": action,
        "enabled": enabled,
        "order": order,
        "conditions": conditions or [],
    }
    if description:
        rule["description"] = description
    if principal:
        rule["principal"] = principal
    if resource:
        rule["resource"] = resource
    if severity:
        rule["severity"] = severity
    if tags:
        rule["tags"] = tags
    return rule


def new_policy_entity(type: str, id: str | None = None) -> PolicyEntity:
    """Create a new PolicyEntity."""
    entity: PolicyEntity = {"type": type}
    if id:
        entity["id"] = id
    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}
