"""
Cedar Policy Builder

Provides functions to convert PolicyRule to Cedar policy text with proper annotations.
"""

import re
from typing import Union

from .annotations import generate_annotation_lines
from .rule import PolicyCondition, PolicyEffect, PolicyEntity, PolicyRule

# ============================================================================
# Security: Input Validation and Escaping
# ============================================================================

# Valid identifier pattern for Cedar (alphanumeric, underscore, with optional namespace)
_VALID_IDENTIFIER_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(::[A-Za-z_][A-Za-z0-9_]*)*$")

# Pattern that matches potentially dangerous content in raw conditions.
# Note: '}' is NOT blocked because raw conditions may contain valid Cedar
# record literals or nested expressions. The ';' blocks when-clause escapes.
_DANGEROUS_PATTERN_REGEX = re.compile(r";|//|/\*|\*/|permit\s*\(|forbid\s*\(")


def _escape_cedar_string(value: str) -> str:
    """
    Escape a string value for use in Cedar string literals.
    This prevents injection attacks by escaping backslashes and double quotes.
    """
    return value.replace("\\", "\\\\").replace('"', '\\"')


def _is_valid_identifier(s: str) -> bool:
    """Check if a string is a valid Cedar identifier."""
    return bool(_VALID_IDENTIFIER_REGEX.match(s))


def _sanitize_identifier(s: str, context: str) -> str:
    """
    Sanitize an identifier, replacing invalid characters with underscores.
    """
    if _is_valid_identifier(s):
        return s
    # Replace invalid characters with underscores
    sanitized = re.sub(r"[^A-Za-z0-9_:]", "_", s)
    if not sanitized or not _is_valid_identifier(sanitized):
        return f"invalid_{context}"
    return sanitized


def _is_valid_raw_condition(condition: str) -> bool:
    """
    Validate a raw condition string for potentially dangerous patterns.
    Returns True if the condition is safe to use.
    """
    return not bool(_DANGEROUS_PATTERN_REGEX.search(condition))


def rule_to_cedar(
    rule: PolicyRule,
    optional_fields: set[str] | None = None,
) -> str:
    """
    Convert a PolicyRule to Cedar policy text with proper annotations.

    Args:
        rule: The PolicyRule to convert.
        optional_fields: Set of context field names that are optional and need
            ``context has`` guards. Pass None to skip guard injection.

    Example output:
        @id("rule-001")
        @name("Block threats")
        @severity("high")
        forbid (
            principal,
            action == Action::"call_tool",
            resource
        )
        when { context has threat_count && context.threat_count > 0 };
    """
    lines: list[str] = []

    # Generate Cedar annotations
    annotations = rule.get("annotations", {})
    custom_annotations = rule.get("customAnnotations")
    lines.extend(generate_annotation_lines(annotations, custom_annotations))

    # Generate policy body
    lines.append(
        _generate_policy_body(
            rule.get("effect", "forbid"),
            rule.get("principal"),
            rule.get("action", ""),
            rule.get("resource"),
            rule.get("conditions", []),
            rule.get("rawCondition"),
            optional_fields,
        )
    )

    return "\n".join(lines)


def rules_to_cedar(
    rules: list[PolicyRule],
    include_disabled: bool = False,
    optional_fields: set[str] | None = None,
) -> str:
    """
    Convert multiple PolicyRules to Cedar policy text.
    Only enabled rules are included, sorted by order.

    Args:
        rules: List of PolicyRules to convert.
        include_disabled: If True, disabled rules are included as comments.
        optional_fields: Set of context field names that are optional and need
            ``context has`` guards. Pass None to skip guard injection.
    """
    # Sort by order
    sorted_rules = sorted(rules, key=lambda r: r.get("order", 0))

    cedar_policies: list[str] = []
    for rule in sorted_rules:
        if rule.get("enabled", True):
            cedar_policies.append(rule_to_cedar(rule, optional_fields))
        elif include_disabled:
            # Include disabled rules as comments
            cedar_lines = rule_to_cedar(rule, optional_fields).split("\n")
            commented_lines = [f"// [DISABLED] {line}" for line in cedar_lines]
            cedar_policies.append("\n".join(commented_lines))

    return "\n\n".join(cedar_policies)


def _generate_policy_body(
    effect: PolicyEffect,
    principal: PolicyEntity | None,
    action: Union[str, list[str]],
    resource: PolicyEntity | None,
    conditions: list[PolicyCondition],
    raw_condition: str | None,
    optional_fields: set[str] | None = None,
) -> str:
    """
    Generate the Cedar policy body (permit/forbid statement).
    All inputs are sanitized/escaped to prevent injection attacks.
    """
    parts: list[str] = [f"{effect} ("]

    # Principal
    if principal:
        entity_type = _sanitize_identifier(principal["type"], "principal_type")
        if principal.get("id"):
            op = "in" if principal.get("operator") == "in" else "=="
            escaped_id = _escape_cedar_string(principal["id"])
            parts.append(f'\n    principal {op} {entity_type}::"{escaped_id}"')
        else:
            parts.append(f"\n    principal is {entity_type}")
    else:
        parts.append("\n    principal")

    # Action
    if isinstance(action, list):
        filtered = [a for a in action if a != "*"]
        if len(filtered) == 0:
            parts.append(",\n    action")
        elif len(filtered) == 1:
            parts.append(f",\n    action == {_format_action(filtered[0])}")
        else:
            actions = ", ".join(_format_action(a) for a in filtered)
            parts.append(f",\n    action in [{actions}]")
    elif not action or action == "*":
        parts.append(",\n    action")
    else:
        parts.append(f",\n    action == {_format_action(action)}")

    # Resource
    if resource:
        entity_type = _sanitize_identifier(resource["type"], "resource_type")
        if resource.get("id"):
            op = "in" if resource.get("operator") == "in" else "=="
            escaped_id = _escape_cedar_string(resource["id"])
            parts.append(f',\n    resource {op} {entity_type}::"{escaped_id}"')
        else:
            parts.append(f",\n    resource is {entity_type}")
    else:
        parts.append(",\n    resource")

    parts.append("\n)")

    # When clause
    # SECURITY: raw_condition is validated to prevent injection attacks.
    # If validation fails, fall back to structured conditions.
    if raw_condition:
        if _is_valid_raw_condition(raw_condition):
            parts.append(f"\nwhen {{ {raw_condition} }};")
        elif conditions:
            # Fallback to structured conditions if raw_condition is rejected
            condition_strs = [_condition_to_cedar(c, optional_fields) for c in conditions]
            parts.append(f"\nwhen {{ {' && '.join(condition_strs)} }};")
        else:
            parts.append(";")
    elif conditions:
        condition_strs = [_condition_to_cedar(c, optional_fields) for c in conditions]
        parts.append(f"\nwhen {{ {' && '.join(condition_strs)} }};")
    else:
        parts.append(";")

    return "".join(parts)


def _format_action(action: str) -> str:
    """
    Format an action string for Cedar policy text.
    Escapes the action name to prevent injection attacks.
    """
    if 'Action::"' in action:
        # Already namespaced - extract and escape the action name
        parts = action.split('Action::"', 1)
        if len(parts) == 2:
            action_name = parts[1].rstrip('"')
            return f'{parts[0]}Action::"{_escape_cedar_string(action_name)}"'
        return action
    return f'Action::"{_escape_cedar_string(action)}"'


def _condition_to_cedar(
    condition: PolicyCondition,
    optional_fields: set[str] | None = None,
) -> str:
    """
    Convert a condition to Cedar syntax.
    Field names are sanitized to prevent injection attacks.

    When optional_fields is provided and the condition's field is in the set,
    the output is wrapped with a ``context has`` guard::

        context has field && context.field > value
    """
    field = _sanitize_identifier(condition["field"], "field")
    operator = condition["operator"]
    value = condition["value"]
    value_str = _value_to_string(value)

    operator_map = {
        "eq": "==",
        "neq": "!=",
        "lt": "<",
        "lte": "<=",
        "gt": ">",
        "gte": ">=",
    }

    if operator in operator_map:
        expr = f"context.{field} {operator_map[operator]} {value_str}"
    elif operator == "contains":
        expr = f"context.{field}.contains({value_str})"
    elif operator == "in":
        expr = f"context.{field} in {value_str}"
    elif operator == "like":
        expr = f"context.{field} like {value_str}"
    else:
        expr = f"context.{field} == {value_str}"

    # Auto-inject `context has` guard for optional fields
    if optional_fields and field in optional_fields:
        return f"context has {field} && {expr}"
    return expr


def _value_to_string(value: Union[str, int, float, bool, list[str]]) -> str:
    """
    Convert a value to Cedar string representation.
    String values are escaped to prevent injection attacks.
    """
    if isinstance(value, str):
        return f'"{_escape_cedar_string(value)}"'
    elif isinstance(value, bool):
        return "true" if value else "false"
    elif isinstance(value, (int, float)):
        if isinstance(value, float) and value == int(value):
            return str(int(value))
        return str(value)
    elif isinstance(value, list):
        quoted = [f'"{_escape_cedar_string(v)}"' for v in value]
        return f"[{', '.join(quoted)}]"
    else:
        return str(value)
