"""
Cedar Policy Parser

Converts Cedar policy text to structured PolicyRule format using the
official Cedar engine (cedarpy) for parsing.

Architecture:
1. Cedar text → Cedar JSON (via cedarpy.policies_to_json_str)
2. Cedar JSON → PolicyRule (simple JSON mapping)
"""

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

import cedarpy

from .rule import (
    PolicyRule,
    PolicyCondition,
    PolicyEntity,
    PolicyEffect,
    ConditionOperator,
    ConditionExpression,
    new_policy_rule,
    new_policy_entity,
    new_policy_condition,
)
from .errors import ParserError, ErrorCodes


@dataclass
class ParseResult:
    """Result of parsing Cedar policies."""

    # Policies successfully converted to PolicyRule format
    rules: list[PolicyRule] = field(default_factory=list)
    # Policies that couldn't be fully represented (raw Cedar text)
    unstructured: list[str] = field(default_factory=list)
    # Any parsing errors encountered
    errors: list[str] = field(default_factory=list)


def parse_cedar_to_rules(cedar_text: str) -> ParseResult:
    """
    Parse Cedar policy text and convert to PolicyRule format.

    Uses the official Cedar engine (cedarpy) for parsing.
    Policies with features that can't be represented as PolicyRule (e.g.,
    unless clauses, complex expressions) are returned in the unstructured list.

    Args:
        cedar_text: Cedar policy text to parse

    Returns:
        ParseResult with structured rules, unstructured policies, and errors
    """
    result = ParseResult()

    try:
        # Use cedarpy to convert Cedar text to JSON
        json_str = cedarpy.policies_to_json_str(cedar_text)
        policy_set = json.loads(json_str)

        # Process each static policy
        static_policies = policy_set.get("staticPolicies", {})
        for index, (policy_id, policy) in enumerate(static_policies.items()):
            # Get engine-serialized Cedar text for rawCondition extraction.
            # Using cedarpy's policies_from_json_str ensures we get the official
            # engine's representation rather than relying on our own text extraction.
            single_policy_json = json.dumps({
                "staticPolicies": {policy_id: policy},
                "templates": {},
                "templateLinks": [],
            })
            try:
                engine_text = cedarpy.policies_from_json_str(single_policy_json)
            except Exception:
                engine_text = ""

            rule, raw, error = _cedar_json_to_rule(policy, policy_id, index, engine_text)

            if error:
                result.errors.append(f"Policy {policy_id}: {error}")

            if rule:
                result.rules.append(rule)
            elif raw:
                result.unstructured.append(raw)

        # Templates can't be represented as PolicyRule
        # Preserve full template body for downstream systems
        templates = policy_set.get("templates", {})
        for template_id, template_body in templates.items():
            # Store template as JSON for round-trip fidelity
            result.unstructured.append(json.dumps({
                "type": "template",
                "id": template_id,
                "body": template_body
            }))

    except Exception as e:
        result.errors.append(f"Parse error: {str(e)}")

    # Check for duplicate policy IDs and add warnings
    id_occurrences: dict[str, list[int]] = {}
    for idx, rule in enumerate(result.rules):
        rule_id = rule.get("id", f"policy{idx}")
        if rule_id not in id_occurrences:
            id_occurrences[rule_id] = []
        id_occurrences[rule_id].append(idx)

    for policy_id, indices in id_occurrences.items():
        if len(indices) > 1:
            result.errors.append(
                f"[{ErrorCodes.PARSE_DUPLICATE_ID}] Duplicate policy ID '{policy_id}' found at indices {indices}"
            )

    return result


def _cedar_json_to_rule(
    policy: dict[str, Any],
    policy_id: str,
    index: int,
    original_cedar: str,
) -> tuple[Optional[PolicyRule], Optional[str], Optional[str]]:
    """
    Convert Cedar JSON policy to PolicyRule.
    This is pure JSON mapping - no parsing logic.

    Returns:
        Tuple of (rule, raw_cedar, error)
    """
    try:
        # Check if this policy can be represented as PolicyRule
        if not _can_represent_as_rule(policy):
            # Return raw Cedar for this policy (best effort extraction)
            raw = _extract_policy_cedar(policy_id, original_cedar)
            return None, raw, None

        annotations = policy.get("annotations", {})

        rule = new_policy_rule(
            id=annotations.get("id", policy_id),
            name=annotations.get("name", annotations.get("id", policy_id)),
            effect=policy["effect"],
            action=_map_action_scope(policy.get("action", {})),
            enabled=True,
            order=index,
            description=annotations.get("description"),
            principal=_map_scope_to_entity(policy.get("principal", {}), "principal"),
            resource=_map_scope_to_entity(policy.get("resource", {}), "resource"),
        )

        # Map conditions - engine_text is the single-policy Cedar text from
        # cedarpy's serialization, so no policy_id lookup is needed
        conditions, raw_condition, condition_expression = _map_conditions(
            policy.get("conditions", []), original_cedar
        )
        rule["conditions"] = conditions
        if raw_condition:
            rule["rawCondition"] = raw_condition
        if condition_expression:
            rule["conditionExpression"] = condition_expression

        return rule, None, None

    except ParserError as e:
        # Malformed constraint - return raw Cedar with error
        raw = _extract_policy_cedar(policy_id, original_cedar)
        return None, raw, str(e)
    except Exception as e:
        return None, None, str(e)


def _can_represent_as_rule(policy: dict[str, Any]) -> bool:
    """Check if a Cedar policy can be represented as PolicyRule."""
    # Unless clauses can't be represented
    for cond in policy.get("conditions", []):
        if cond.get("kind") == "unless":
            return False

    # Template slots can't be represented
    principal = policy.get("principal", {})
    resource = policy.get("resource", {})
    if principal.get("slot") or resource.get("slot"):
        return False

    # Multiple entity "in" constraints are complex
    if principal.get("op") == "in":
        entities = principal.get("entities", [])
        if len(entities) > 1:
            return False

    if resource.get("op") == "in":
        entities = resource.get("entities", [])
        if len(entities) > 1:
            return False

    return True


def _extract_policy_cedar(policy_id: str, original_cedar: str) -> str:
    """Extract raw Cedar for a policy that can't be represented as PolicyRule."""
    # Simple fallback - return a comment with the policy ID
    # In production, you'd want to use cedarpy_conversor.convert_json_to_cedar_policies
    # Sanitize policy_id to prevent newline injection attacks
    safe_id = policy_id.replace("\n", " ").replace("\r", " ")
    return f"// Complex policy: {safe_id}"


def _map_scope_to_entity(
    scope: dict[str, Any], field: str
) -> Optional[PolicyEntity]:
    """
    Map Cedar scope constraint to PolicyEntity.

    Returns None for unconstrained ("All") scopes.
    Raises ParserError for malformed constraints.

    Args:
        scope: The Cedar scope constraint
        field: The field name ("principal" or "resource") for error context
    """
    op = scope.get("op", "All")

    if op == "All":
        return None
    elif op == "==":
        entity = scope.get("entity", {})
        if entity:
            return new_policy_entity(entity.get("type", ""), entity.get("id"))
        if scope.get("slot"):
            raise ParserError.scope_slot_not_supported("==", field)
        raise ParserError.scope_missing_entity("==", field)
    elif op == "is":
        entity_type = scope.get("entity_type", "")
        if entity_type:
            return new_policy_entity(entity_type)
        raise ParserError.scope_missing_entity_type(field)
    elif op == "in":
        # Single entity "in" - treat as specific entity with "in" operator
        entities = scope.get("entities", [])
        if len(entities) == 1:
            return new_policy_entity(
                entities[0].get("type", ""), entities[0].get("id"), operator="in"
            )
        # Single entity reference
        entity = scope.get("entity", {})
        if entity:
            return new_policy_entity(
                entity.get("type", ""), entity.get("id"), operator="in"
            )
        if scope.get("slot"):
            raise ParserError.scope_slot_not_supported("in", field)
        raise ParserError.scope_missing_entity_list(field)
    else:
        raise ParserError.scope_unsupported_op(op, field)


def _format_action_entity(entity: dict[str, Any]) -> str:
    """Format action entity reference, preserving namespace if present.

    If entity type is just "Action", returns just the id (e.g., "process_prompt").
    If entity type has a namespace (e.g., "Overwatch::Action"),
    returns the fully qualified action string (e.g., 'Overwatch::Action::"process_prompt"').
    """
    entity_type = entity.get("type", "Action")
    action_id = entity.get("id", "")
    if entity_type != "Action":
        return f'{entity_type}::"{action_id}"'
    return action_id


def _map_action_scope(scope: dict[str, Any]) -> Union[str, list[str]]:
    """
    Map action scope to action string(s).

    Raises ParserError for malformed constraints.
    """
    op = scope.get("op", "All")

    if op == "All":
        return "*"
    elif op == "==":
        entity = scope.get("entity", {})
        if entity:
            return _format_action_entity(entity)
        raise ParserError.action_missing_entity("==")
    elif op == "in":
        entities = scope.get("entities", [])
        if entities:
            actions = [_format_action_entity(e) for e in entities]
            return actions[0] if len(actions) == 1 else actions
        entity = scope.get("entity", {})
        if entity:
            return _format_action_entity(entity)
        raise ParserError.action_missing_entities()
    else:
        raise ParserError.action_unsupported_op(op)


def _map_conditions(
    conditions: list[dict[str, Any]],
    engine_cedar: str = "",
) -> tuple[list[PolicyCondition], Optional[str], Optional[ConditionExpression]]:
    """Map Cedar conditions to PolicyCondition list + recursive ConditionExpression.

    When conditions can't be mapped to structured format, extract the raw Cedar
    condition text from the engine-serialized policy text (not JSON AST).
    """
    result: list[PolicyCondition] = []
    has_unmapped = False

    # Collect condition expressions from all when clauses
    expressions: list[ConditionExpression] = []

    for cond in conditions:
        if cond.get("kind") != "when":
            continue

        # Flat mapping (backward compat)
        parsed, raw = _map_condition_body(cond.get("body", {}))
        if parsed:
            result.append(parsed)
        elif raw:
            has_unmapped = True

        # Recursive mapping (new)
        expressions.append(
            _map_condition_body_to_expression(cond.get("body", {}), engine_cedar)
        )

    # Extract readable Cedar condition text instead of storing JSON AST
    raw_condition = None
    if has_unmapped and engine_cedar:
        raw_condition = _extract_when_clause(engine_cedar)

    # Build the final condition expression
    condition_expression: Optional[ConditionExpression] = None
    if len(expressions) == 1:
        condition_expression = expressions[0]
    elif len(expressions) > 1:
        condition_expression = {"kind": "and", "children": expressions}

    return result, raw_condition, condition_expression


def _extract_when_clause(cedar_text: str) -> str:
    """Extract the readable condition text from a Cedar policy's when clause.

    Expects single-policy Cedar text (from engine serialization).

    Given: `forbid (...)\\nwhen { context.path like "/etc/*" };`
    Returns: `context.path like "/etc/*"`
    """
    when_prefix = "when {"
    idx = cedar_text.find(when_prefix)
    if idx < 0:
        return ""

    body = cedar_text[idx + len(when_prefix):]
    depth = 1
    for i, ch in enumerate(body):
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return body[:i].strip()

    return body.strip()


def _map_condition_body(
    body: dict[str, Any],
) -> tuple[Optional[PolicyCondition], Optional[str]]:
    """
    Map a Cedar expression body to PolicyCondition.

    Cedar JSON expressions use nested objects with operator keys.
    Format: {"==": {"left": {".": {"left": {"Var": "context"}, "attr": "field"}}, "right": {"Value": "x"}}}
    """
    for op, args in body.items():
        if op in ("==", "!=", "<", "<=", ">", ">="):
            condition = _map_comparison(op, args)
            if condition:
                return condition, None
        elif op == "contains":
            condition = _map_contains(args)
            if condition:
                return condition, None
        elif op == "like":
            condition = _map_like(args)
            if condition:
                return condition, None

    # Can't map - return as raw JSON
    return None, json.dumps(body)


def _map_comparison(op: str, args: Any) -> Optional[PolicyCondition]:
    """Map a comparison expression to PolicyCondition."""
    if not isinstance(args, dict):
        return None

    left = args.get("left")
    right = args.get("right")
    if not left or not right:
        return None

    field = _extract_context_field(left)
    if not field:
        return None

    value = _extract_literal_value(right)
    if value is None:
        return None

    operator = _map_operator(op)
    if not operator:
        return None

    return new_policy_condition(field, operator, value)


def _map_contains(args: Any) -> Optional[PolicyCondition]:
    """Map a contains expression to PolicyCondition."""
    if not isinstance(args, dict):
        return None

    left = args.get("left")
    right = args.get("right")
    if not left or not right:
        return None

    field = _extract_context_field(left)
    if not field:
        return None

    value = _extract_literal_value(right)
    if value is None:
        return None

    return new_policy_condition(field, "contains", value)


def _map_like(args: Any) -> Optional[PolicyCondition]:
    """Map a like expression to PolicyCondition."""
    if not isinstance(args, dict):
        return None

    left = args.get("left")
    pattern = args.get("pattern")
    if not left or pattern is None:
        return None

    field = _extract_context_field(left)
    if not field:
        return None

    # Convert pattern list to string (e.g., ["Wildcard", {"Literal": "foo"}] -> "*foo")
    # Escape literal * characters to distinguish from wildcards on round-trip
    pattern_str = ""
    for p in pattern:
        if p == "Wildcard":
            pattern_str += "*"
        elif isinstance(p, dict) and "Literal" in p:
            # Escape literal * to distinguish from wildcards
            pattern_str += p["Literal"].replace("*", "\\*")

    return new_policy_condition(field, "like", pattern_str)


def _extract_context_field(expr: Any) -> Optional[str]:
    """
    Extract field name from context.field access pattern.
    Pattern: {".": {"left": {"Var": "context"}, "attr": "field_name"}}
    """
    if not isinstance(expr, dict):
        return None

    dot_access = expr.get(".")
    if not isinstance(dot_access, dict):
        return None

    left = dot_access.get("left")
    if not isinstance(left, dict) or left.get("Var") != "context":
        return None

    attr = dot_access.get("attr")
    return attr if isinstance(attr, str) else None


def _extract_literal_value(
    expr: Any,
) -> Optional[Union[str, int, float, bool, list[str]]]:
    """
    Extract literal value from Cedar JSON.
    Pattern: {"Value": <literal>}
    """
    if not isinstance(expr, dict):
        return None

    value = expr.get("Value")
    if value is None:
        return None

    if isinstance(value, (str, int, float, bool)):
        return value

    if isinstance(value, list) and all(isinstance(v, str) for v in value):
        return value

    return None


def _map_operator(cedar_op: str) -> Optional[ConditionOperator]:
    """Map Cedar operator to ConditionOperator."""
    mapping: dict[str, ConditionOperator] = {
        "==": "eq",
        "!=": "neq",
        "<": "lt",
        "<=": "lte",
        ">": "gt",
        ">=": "gte",
    }
    return mapping.get(cedar_op)


# ---------------------------------------------------------------------------
# Recursive condition expression walker
# ---------------------------------------------------------------------------


def _map_condition_body_to_expression(
    body: dict[str, Any], original_text: str = ""
) -> ConditionExpression:
    """Recursively walk the Cedar JSON AST to build a ConditionExpression tree.

    Flattens binary && / || chains into n-ary and/or nodes.
    """
    # Logical AND -- flatten binary chain
    if "&&" in body:
        children = _flatten_binary_chain("&&", body, original_text)
        return {"kind": "and", "children": children}

    # Logical OR -- flatten binary chain
    if "||" in body:
        children = _flatten_binary_chain("||", body, original_text)
        return {"kind": "or", "children": children}

    # Negation
    if "!" in body:
        child = _map_condition_body_to_expression(body["!"]["arg"], original_text)
        return {"kind": "not", "child": child}

    # Has (existence check): {"has": {"left": {"Var": "context"}, "attr": "field"}}
    if "has" in body:
        has_expr = body["has"]
        field = _extract_has_field(has_expr)
        if field:
            return {"kind": "has", "field": field}

    # Leaf: comparison operators
    for op in ("==", "!=", "<", "<=", ">", ">="):
        if op in body:
            mapped = _map_comparison_to_expression(op, body[op])
            if mapped:
                return mapped

    # Leaf: contains
    if "contains" in body:
        mapped = _map_contains_to_expression(body["contains"])
        if mapped:
            return mapped

    # Leaf: like
    if "like" in body:
        mapped = _map_like_to_expression(body["like"])
        if mapped:
            return mapped

    # Fallback -- extract readable text if possible
    text = _extract_when_clause(original_text) if original_text else json.dumps(body)
    return {"kind": "raw", "text": text}


def _flatten_binary_chain(
    op: str, body: dict[str, Any], original_text: str = ""
) -> list[ConditionExpression]:
    """Flatten a binary chain of && or || into a flat array of expressions.

    (A && B) && C becomes [A, B, C] instead of nested binary pairs.
    """
    node = body.get(op)
    if not node:
        return [_map_condition_body_to_expression(body, original_text)]
    return [
        *_flatten_binary_chain(op, node["left"], original_text),
        *_flatten_binary_chain(op, node["right"], original_text),
    ]


def _extract_has_field(has_expr: dict[str, Any]) -> Optional[str]:
    """Extract field name from a 'has' expression.

    Pattern: {"left": {"Var": "context"}, "attr": "field_name"}
    """
    left = has_expr.get("left", {})
    if isinstance(left, dict) and left.get("Var") == "context":
        attr = has_expr.get("attr")
        return attr if isinstance(attr, str) else None
    return None


def _map_comparison_to_expression(
    op: str, args: dict[str, Any]
) -> Optional[ConditionExpression]:
    """Map a comparison expression to a ConditionExpression comparison node."""
    if not isinstance(args, dict):
        return None
    field = _extract_context_field(args.get("left"))
    if not field:
        return None
    value = _extract_literal_value(args.get("right"))
    if value is None:
        return None
    operator = _map_operator(op)
    if not operator:
        return None
    return {"kind": "comparison", "field": field, "operator": operator, "value": value}


def _map_contains_to_expression(
    args: dict[str, Any],
) -> Optional[ConditionExpression]:
    """Map a contains expression to a ConditionExpression contains node."""
    if not isinstance(args, dict):
        return None
    field = _extract_context_field(args.get("left"))
    if not field:
        return None
    value = _extract_literal_value(args.get("right"))
    if value is None or isinstance(value, list):
        return None
    return {"kind": "contains", "field": field, "value": value}


def _map_like_to_expression(
    args: dict[str, Any],
) -> Optional[ConditionExpression]:
    """Map a like expression to a ConditionExpression like node."""
    if not isinstance(args, dict):
        return None
    field = _extract_context_field(args.get("left"))
    if not field:
        return None
    pattern = args.get("pattern")
    if pattern is None:
        return None
    # Convert pattern list to string
    pattern_str = ""
    for p in pattern:
        if p == "Wildcard":
            pattern_str += "*"
        elif isinstance(p, dict) and "Literal" in p:
            pattern_str += p["Literal"].replace("*", "\\*")
    return {"kind": "like", "field": field, "pattern": pattern_str}
