"""
Highflame Policy Engine - Python Wrapper
Wraps cedarpy with Highflame-specific types
"""

import json
from dataclasses import dataclass, field
from typing import Any, Optional

import cedarpy

from .entities import EntityType
from .actions import ActionType


# Default validation limits - consistent across all SDK packages
@dataclass
class ValidationLimits:
    """Limits for input validation to prevent DoS attacks."""
    max_context_keys: int = 100
    max_string_length: int = 1_000_000
    max_nesting_depth: int = 10
    max_context_size_bytes: int = 10_000_000


DEFAULT_LIMITS = ValidationLimits()


class InputValidationError(Exception):
    """Raised when input validation fails."""
    pass


@dataclass
class EngineOptions:
    """Configuration options for PolicyEngine."""
    schema: Optional[str] = None
    limits: Optional[ValidationLimits] = None
    skip_validation: bool = False


@dataclass
class EntityUID:
    """Cedar entity unique identifier."""
    type: str
    id: str


@dataclass
class DeterminingPolicy:
    """A policy that contributed to the authorization decision,
    enriched with its Cedar annotations."""
    id: str
    annotations: dict[str, str]


@dataclass
class Decision:
    """Result of a policy evaluation."""
    effect: str  # "Allow" or "Deny"
    determining_policies: list[DeterminingPolicy]
    reason: Optional[str] = None

    def is_allowed(self) -> bool:
        return self.effect == "Allow"

    def is_denied(self) -> bool:
        return self.effect == "Deny"


def _extract_policy_annotations(policy_text: str) -> tuple[dict[str, str], dict[str, dict[str, str]]]:
    """Extract all annotations from Cedar policy text.

    Returns a tuple of:
    - id_map: mapping from positional IDs (policy0, policy1...) to @id annotation values
    - annotations_map: mapping from resolved policy ID to all annotations dict

    cedarpy returns positional IDs in diagnostics.reasons. The id_map
    allows translating them to meaningful @id values. The annotations_map
    stores all annotations for enriching evaluation results.

    Uses cedarpy.policies_to_json_str() for robust parsing instead of regex.
    """
    try:
        policies_json = json.loads(cedarpy.policies_to_json_str(policy_text))
    except (ValueError, json.JSONDecodeError):
        return {}, {}

    id_map: dict[str, str] = {}
    annotations_map: dict[str, dict[str, str]] = {}
    for pos_id, policy in policies_json.get("staticPolicies", {}).items():
        annotations = policy.get("annotations", {})
        str_annotations = {k: str(v) for k, v in annotations.items()}
        policy_id = annotations.get("id", pos_id)
        if "id" in annotations:
            id_map[pos_id] = annotations["id"]
        annotations_map[policy_id] = str_annotations

    return id_map, annotations_map


class PolicyEngine:
    """
    PolicyEngine wraps cedarpy with Highflame schema types.

    Example:
        engine = PolicyEngine()
        engine.load_policies_from_file("policies/palisade/policy.cedar")

        decision = engine.evaluate(
            principal_type=EntityType.SCANNER,
            principal_id="palisade",
            action=ActionType.SCAN_ARTIFACT,
            resource_type=EntityType.ARTIFACT,
            resource_id="/path/to/model.safetensors",
            context={
                "artifact_format": "safetensors",
                "severity": "HIGH",
                "environment": "production",
            }
        )

        if decision.is_denied():
            print(f"Blocked by: {decision.determining_policies}")
    """

    def __init__(self, schema: Optional[str] = None, options: Optional[EngineOptions] = None):
        self._policies: str = ""
        self._policy_id_map: dict[str, str] = {}  # policy0 -> @id value
        self._policy_annotations: dict[str, dict[str, str]] = {}  # policyID -> annotations
        self._schema: Optional[str] = schema or (options.schema if options else None)
        self._limits = options.limits if options and options.limits else DEFAULT_LIMITS
        self._skip_validation = options.skip_validation if options else False

    def _load_policy_text(self, policy_text: str) -> None:
        """Internal helper to load policy text and extract annotations."""
        self._policies = policy_text
        self._policy_id_map, self._policy_annotations = _extract_policy_annotations(policy_text)

    def load_policies_from_file(self, path: str) -> None:
        """Load Cedar policies from a file."""
        with open(path, "r") as f:
            policy_text = f.read()
        self._load_policy_text(policy_text)

    def load_policy(self, policy: str) -> None:
        """Load a single Cedar policy text string.
        Uses @id annotations as policy IDs when available.
        Stores all annotations per policy for enriching evaluation results."""
        self._load_policy_text(policy)

    def load_policies(self, policies: list[str]) -> None:
        """Load multiple Cedar policy texts (concatenated with newlines).
        Uses @id annotations as policy IDs when available.
        Stores all annotations per policy for enriching evaluation results."""
        self._load_policy_text("\n".join(policies))

    def load_schema_from_file(self, path: str) -> None:
        """Load Cedar schema from a file."""
        with open(path, "r") as f:
            self._schema = f.read()

    def load_schema(self, schema: str) -> None:
        """Load Cedar schema from a string."""
        self._schema = schema

    def get_policy_annotations(self, policy_id: str) -> dict[str, str] | None:
        """Returns stored annotations for a given policy ID.
        Returns None if the policy ID is not found."""
        return self._policy_annotations.get(policy_id)

    def evaluate(
        self,
        principal_type: str,
        principal_id: str,
        action: str,
        resource_type: str,
        resource_id: str,
        context: Optional[dict[str, Any]] = None,
        entities: Optional[list[dict[str, Any]]] = None,
    ) -> Decision:
        """
        Evaluate a policy request and return a decision.

        Args:
            principal_type: Entity type of the principal (e.g., EntityType.USER)
            principal_id: ID of the principal
            action: Action being performed (e.g., ActionType.SCAN_ARTIFACT)
            resource_type: Entity type of the resource (e.g., EntityType.ARTIFACT)
            resource_id: ID of the resource
            context: Optional context attributes for the request

        Returns:
            Decision with effect, determining policies, and optional reason

        Raises:
            InputValidationError: If context validation fails
        """
        # Validate context if validation is enabled
        if not self._skip_validation and context:
            self._validate_context(context)

        # Format entity UIDs for Cedar
        # Handle both namespaced (Overwatch::User) and non-namespaced (User) types
        principal = f'{principal_type}::"{principal_id}"'

        # Handle both pre-formatted actions (Namespace::Action::"name") and simple action names
        # Check if action is already fully qualified by looking for Action::" pattern
        if 'Action::"' in action:
            # Action is already fully qualified (e.g., "Overwatch::Action::\"call_tool\"")
            action_uid = action
        else:
            # Simple action name, add Action:: prefix
            action_uid = f'Action::"{action}"'

        resource = f'{resource_type}::"{resource_id}"'

        # Call cedarpy (v4.8.0+ uses request dict instead of separate args)
        request = {
            "principal": principal,
            "action": action_uid,
            "resource": resource,
            "context": context or {},
        }
        result = cedarpy.is_authorized(
            request=request,
            policies=self._policies,
            entities=entities or [],
        )

        # Map positional IDs (policy0, policy1...) to @id annotation values
        # and build enriched DeterminingPolicy objects with annotations
        raw_reasons = list(result.diagnostics.reasons) if result.diagnostics else []
        determining: list[DeterminingPolicy] = []
        for pid in raw_reasons:
            resolved_id = self._policy_id_map.get(pid, pid)
            annots = self._policy_annotations.get(resolved_id, {})
            determining.append(DeterminingPolicy(id=resolved_id, annotations=annots))

        return Decision(
            effect="Allow" if result.allowed else "Deny",
            determining_policies=determining,
            reason="; ".join(result.diagnostics.errors) if result.diagnostics and result.diagnostics.errors else None,
        )

    def validate_policies(self, policies: str) -> list[str]:
        """
        Validate policies against the loaded schema.

        Returns:
            List of validation error messages, or empty list if valid.
        """
        if not self._schema:
            return ["No schema loaded for validation"]

        result = cedarpy.validate(
            policies=policies,
            schema=self._schema,
        )

        if result.validation_errors:
            return [str(e) for e in result.validation_errors]
        return []

    def _validate_context(self, context: dict[str, Any]) -> None:
        """
        Validate context against configured limits.

        Raises:
            InputValidationError: If validation fails
        """
        # Check total number of keys
        if len(context) > self._limits.max_context_keys:
            raise InputValidationError(
                f"context exceeds maximum key count: {len(context)} > {self._limits.max_context_keys}"
            )

        # Check total size (approximate via JSON)
        try:
            context_json = json.dumps(context)
            if len(context_json.encode("utf-8")) > self._limits.max_context_size_bytes:
                raise InputValidationError(
                    f"context exceeds maximum size: {len(context_json)} bytes > {self._limits.max_context_size_bytes} bytes"
                )
        except (TypeError, ValueError) as e:
            raise InputValidationError(
                f"context is invalid or too complex: {str(e)}"
            )

        # Validate each value recursively
        for key, value in context.items():
            self._validate_value(value, depth=0, key=key)

    def _validate_value(self, value: Any, depth: int, key: str = "") -> None:
        """
        Recursively validate a value against configured limits.

        Raises:
            InputValidationError: If validation fails
        """
        # Check nesting depth
        if depth > self._limits.max_nesting_depth:
            raise InputValidationError(
                f"value exceeds maximum nesting depth: {depth} > {self._limits.max_nesting_depth}"
            )

        if isinstance(value, str):
            if len(value) > self._limits.max_string_length:
                raise InputValidationError(
                    f"string exceeds maximum length: {len(value)} > {self._limits.max_string_length}"
                )
        elif isinstance(value, list):
            for item in value:
                self._validate_value(item, depth + 1, key)
        elif isinstance(value, dict):
            for k, v in value.items():
                self._validate_value(v, depth + 1, k)
