"""Set Operator Validation Check.

Validates proper usage of ForAllValues and ForAnyValue set operators in IAM policies.

Based on AWS IAM best practices:
https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_condition-single-vs-multi-valued-context-keys.html
"""

import logging
from typing import ClassVar

from iam_validator.core.aws_service import AWSServiceFetcher
from iam_validator.core.aws_service.validators import find_matching_condition_key
from iam_validator.core.check_registry import CheckConfig, PolicyCheck
from iam_validator.core.condition_validators import (
    has_if_exists_suffix,
    is_multivalued_context_key,
    normalize_operator,
)
from iam_validator.core.models import Statement, ValidationIssue

logger = logging.getLogger(__name__)


class SetOperatorValidationCheck(PolicyCheck):
    """Check for proper usage of ForAllValues and ForAnyValue set operators."""

    check_id: ClassVar[str] = "set_operator_validation"
    description: ClassVar[str] = "Validates proper usage of ForAllValues and ForAnyValue set operators"
    default_severity: ClassVar[str] = "error"

    async def _is_multivalued_key(
        self,
        condition_key: str,
        fetcher: AWSServiceFetcher | None,
        actions: list[str],
    ) -> bool:
        """Check if a condition key is multivalued using hardcoded list and AWS service metadata.

        First checks the hardcoded known multivalued keys, then falls back to
        looking up the condition key's Types in AWS service data (ArrayOfString etc.).
        """
        # Fast path: check hardcoded known multivalued keys
        if is_multivalued_context_key(condition_key):
            return True

        if not fetcher:
            return False

        # Look up condition key type from AWS service data
        # Try extracting service prefix from the condition key itself (e.g. "route53:KeyName")
        service_prefixes: set[str] = set()
        if ":" in condition_key and not condition_key.startswith("aws:"):
            service_prefixes.add(condition_key.split(":")[0].lower())

        # Also extract service prefixes from the statement's actions
        for action in actions:
            if action == "*":
                continue
            try:
                service_prefix, _ = fetcher.parse_action(action)
                service_prefixes.add(service_prefix.lower())
            except Exception:
                logger.debug("Failed to parse action %s for condition key lookup", action)
                continue

        for service_prefix in service_prefixes:
            try:
                service_detail = await fetcher.fetch_service_by_name(service_prefix)
                matched_key = find_matching_condition_key(condition_key, service_detail.condition_keys)
                if matched_key:
                    condition_key_obj = service_detail.condition_keys[matched_key]
                    if condition_key_obj.types and any(t.startswith("ArrayOf") for t in condition_key_obj.types):
                        return True
            except Exception:
                logger.debug("Failed to look up condition key %s in service %s", condition_key, service_prefix)
                continue

        return False

    async def execute(
        self,
        statement: Statement,
        statement_idx: int,
        fetcher: AWSServiceFetcher,
        config: CheckConfig,
    ) -> list[ValidationIssue]:
        """
        Execute the set operator validation check.

        Validates:
        1. ForAllValues/ForAnyValue not used with single-valued context keys (anti-pattern)
        2. ForAllValues with Allow effect includes Null condition check (security)
        3. ForAnyValue with Deny effect includes Null condition check (predictability)

        Args:
            statement: The IAM statement to check
            statement_idx: Index of this statement in the policy
            fetcher: AWS service fetcher for looking up condition key metadata
            config: Check configuration

        Returns:
            List of validation issues found
        """
        issues = []

        # Only check statements with conditions
        if not statement.condition:
            return issues

        statement_sid = statement.sid
        line_number = statement.line_number
        effect = statement.effect
        actions = statement.get_actions()

        # Track which condition keys have set operators and Null checks
        set_operator_keys: dict[str, str] = {}  # key -> operator prefix
        null_checked_keys: set[str] = set()

        # First pass: Identify set operators and Null checks
        for operator, conditions in statement.condition.items():
            base_operator, _operator_type, set_prefix = normalize_operator(operator)

            # Track Null checks
            if base_operator == "Null":
                for condition_key in conditions.keys():
                    null_checked_keys.add(condition_key)

            # Track set operators
            if set_prefix in ["ForAllValues", "ForAnyValue"]:
                for condition_key in conditions.keys():
                    set_operator_keys[condition_key] = set_prefix

        # Second pass: Validate set operator usage
        for operator, conditions in statement.condition.items():
            base_operator, _operator_type, set_prefix = normalize_operator(operator)
            has_ifexists = has_if_exists_suffix(operator)

            if not set_prefix:
                continue

            # Check each condition key used with a set operator
            for condition_key, _condition_values in conditions.items():
                # Issue 1: Set operator used with single-valued context key (anti-pattern)
                if not await self._is_multivalued_key(condition_key, fetcher, actions):
                    issues.append(
                        ValidationIssue(
                            severity=self.get_severity(config),
                            message=(
                                f"Set operator `{set_prefix}` should not be used with single-valued "
                                f"condition key `{condition_key}`. This can lead to overly permissive policies. "
                                f"Set operators are designed for multivalued context keys like `aws:TagKeys`. "
                                f"See: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_condition-single-vs-multi-valued-context-keys.html"
                            ),
                            statement_sid=statement_sid,
                            statement_index=statement_idx,
                            issue_type="set_operator_on_single_valued_key",
                            condition_key=condition_key,
                            line_number=line_number,
                            field_name="condition",
                        )
                    )

                # Issue 2: ForAllValues with Allow effect without Null check (security risk)
                if set_prefix == "ForAllValues" and effect == "Allow":
                    if condition_key not in null_checked_keys:
                        if has_ifexists:
                            message = (
                                f"Compounded security risk: `{operator}` with `Allow` effect "
                                f"on `{condition_key}` is doubly permissive. `ForAllValues` "
                                f"passes when the value set is empty, and `IfExists` passes "
                                f"when the key is missing entirely. Both mechanisms "
                                f"independently allow access without matching any values. "
                                f'Add: `"Null": {{"{condition_key}": "false"}}` and consider '
                                f"removing `IfExists`."
                            )
                        else:
                            message = (
                                f"Security risk: `ForAllValues` with `Allow` effect on `{condition_key}` "
                                f"should include a `Null` condition check. Without it, requests with missing "
                                f'`{condition_key}` will be granted access. Add: `"Null": {{"{condition_key}": "false"}}`. '
                                f"See: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_condition-single-vs-multi-valued-context-keys.html"
                            )
                        issues.append(
                            ValidationIssue(
                                severity="warning",
                                message=message,
                                statement_sid=statement_sid,
                                statement_index=statement_idx,
                                issue_type="forallvalues_allow_without_null_check",
                                condition_key=condition_key,
                                line_number=line_number,
                                field_name="condition",
                            )
                        )

                # Issue 3: ForAnyValue with Deny effect without Null check (unpredictable)
                if set_prefix == "ForAnyValue" and effect == "Deny":
                    if condition_key not in null_checked_keys:
                        issues.append(
                            ValidationIssue(
                                severity="warning",
                                message=(
                                    f"Unpredictable behavior: `ForAnyValue` with `Deny` effect on `{condition_key}` "
                                    f"should include a `Null` condition check. Without it, requests with missing "
                                    f"`{condition_key}` will evaluate to `No match` instead of denying access. "
                                    f'Add: `"Null": {{"{condition_key}": "false"}}`. '
                                    f"See: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_condition-single-vs-multi-valued-context-keys.html"
                                ),
                                statement_sid=statement_sid,
                                statement_index=statement_idx,
                                issue_type="foranyvalue_deny_without_null_check",
                                field_name="condition",
                                condition_key=condition_key,
                                line_number=line_number,
                            )
                        )

        return issues
