"""
Cedar Policy Annotations

Provides types and utilities for working with Cedar policy annotations.
Annotations are key-value pairs attached to Cedar policies that provide
metadata without affecting policy evaluation.

Cedar annotation syntax:
    @id("rule-001")
    @name("Block critical threats")
    @severity("high")
    permit(...) when {...};

See: https://docs.cedarpolicy.com/policies/syntax-policy.html
"""

import re
import uuid
from typing import Literal, TypedDict

# Type aliases
PolicySeverity = Literal["critical", "high", "medium", "low"]

# Predefined annotation keys with known semantics
PREDEFINED_ANNOTATION_KEYS = frozenset(["id", "name", "description", "severity", "tags"])

# Regex for valid Cedar identifier
_VALID_ANNOTATION_KEY_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")


class PolicyAnnotations(TypedDict, total=False):
    """
    Predefined Cedar annotations with known semantics.
    These are embedded in Cedar policy text using @annotation("value") syntax.

    id and name are required for all rules created via the UI.
    """

    # Unique identifier for this rule (auto-generated UUID if not provided)
    id: str
    # Human-readable rule name (required)
    name: str
    # Longer explanation of what this rule does
    description: str
    # Severity/priority level for display and filtering
    severity: PolicySeverity
    # Categorization tags for grouping and filtering
    tags: list[str]


# Custom annotations are just a string dict
CustomAnnotations = dict[str, str]


def is_predefined_annotation_key(key: str) -> bool:
    """Check if a key is a predefined annotation key."""
    return key in PREDEFINED_ANNOTATION_KEYS


def is_valid_annotation_key(key: str) -> bool:
    """
    Validate a custom annotation key.
    Cedar annotation keys must be valid identifiers: start with letter or underscore,
    followed by letters, numbers, or underscores.
    Returns False for predefined keys (those are handled separately).
    """
    if is_predefined_annotation_key(key):
        return False
    return bool(_VALID_ANNOTATION_KEY_PATTERN.match(key))


def escape_annotation_value(value: str) -> str:
    """
    Escape a string value for use in Cedar annotation.
    Escapes backslashes and double quotes.
    """
    return value.replace("\\", "\\\\").replace('"', '\\"')


def unescape_annotation_value(value: str) -> str:
    """
    Unescape a Cedar annotation value.
    Reverses the escaping done by escape_annotation_value.
    """
    return value.replace('\\"', '"').replace("\\\\", "\\")


def format_annotation(key: str, value: str) -> str:
    """
    Generate Cedar annotation syntax for a single annotation.
    Returns a string like '@severity("high")'.
    """
    return f'@{key}("{escape_annotation_value(value)}")'


def generate_annotation_lines(
    annotations: PolicyAnnotations,
    custom_annotations: CustomAnnotations | None = None,
) -> list[str]:
    """
    Generate all Cedar annotations from PolicyAnnotations and optional custom annotations.
    Returns array of annotation lines to prepend to policy.
    """
    lines: list[str] = []

    # Predefined annotations in consistent order
    if "id" in annotations:
        lines.append(format_annotation("id", annotations["id"]))
    if "name" in annotations:
        lines.append(format_annotation("name", annotations["name"]))
    if "description" in annotations and annotations["description"]:
        lines.append(format_annotation("description", annotations["description"]))
    if "severity" in annotations and annotations["severity"]:
        lines.append(format_annotation("severity", annotations["severity"]))
    if "tags" in annotations and annotations["tags"]:
        # Cedar annotations are single string values, so join tags with comma
        lines.append(format_annotation("tags", ",".join(annotations["tags"])))

    # Custom annotations (alphabetical order for consistency)
    if custom_annotations:
        for key in sorted(custom_annotations.keys()):
            if is_valid_annotation_key(key):
                lines.append(format_annotation(key, custom_annotations[key]))

    return lines


class ParseAnnotationsResult(TypedDict):
    """Result of parsing Cedar annotations."""

    annotations: PolicyAnnotations
    custom_annotations: CustomAnnotations | None


def parse_annotations(
    raw_annotations: dict[str, str] | None,
) -> ParseAnnotationsResult:
    """
    Parse Cedar annotations from a dict (as returned by cedarpy).
    Separates predefined annotations from custom annotations.
    """
    annotations: PolicyAnnotations = {}
    custom_annotations: CustomAnnotations = {}

    if not raw_annotations:
        return {"annotations": annotations, "custom_annotations": None}

    for key, value in raw_annotations.items():
        unescaped_value = unescape_annotation_value(value)

        if key == "id":
            annotations["id"] = unescaped_value
        elif key == "name":
            annotations["name"] = unescaped_value
        elif key == "description":
            annotations["description"] = unescaped_value
        elif key == "severity":
            if unescaped_value in ("critical", "high", "medium", "low"):
                annotations["severity"] = unescaped_value  # type: ignore
        elif key == "tags":
            tags = [t.strip() for t in unescaped_value.split(",") if t.strip()]
            annotations["tags"] = tags
        else:
            # Custom annotation
            custom_annotations[key] = unescaped_value

    # Use id as name if name not provided
    if not annotations.get("name") and annotations.get("id"):
        annotations["name"] = annotations["id"]

    return {
        "annotations": annotations,
        "custom_annotations": custom_annotations if custom_annotations else None,
    }


def generate_rule_id() -> str:
    """Generate a UUID v4 for rule IDs."""
    return str(uuid.uuid4())
