"""
Condition key extraction from CloudTrail events.

Resolves AWS condition key values from CloudTrail event context using
global extractors, service-specific extractors (88 manual + auto-generated
from IAM reference), and tag/unobtainable key handling.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

from ..models.report import EvaluationContext

# Resolution status for context key lookups
RESOLVED = "resolved"
NULL = "null"
UNOBTAINABLE = "unobtainable"
SERVICE_SPECIFIC = "service_specific"
UNKNOWN = "unknown"

# Global key extractors: lowercase key -> extractor function
GLOBAL_KEY_EXTRACTORS: dict[str, Callable[[EvaluationContext], Any]] = {
    "aws:sourceip": lambda ctx: ctx.source_ip,
    "aws:principalarn": lambda ctx: ctx.principal_arn,
    "aws:requestedregion": lambda ctx: ctx.region,
    "aws:principalaccount": lambda ctx: ctx.principal_account,
    "aws:principaltype": lambda ctx: ctx.principal_type,
    "aws:useragent": lambda ctx: ctx.user_agent,
    "aws:userid": lambda ctx: ctx.user_id,
    "aws:username": lambda ctx: ctx.username,
    "aws:currenttime": lambda ctx: ctx.request_time.isoformat() if ctx.request_time else None,
    "aws:epochtime": lambda ctx: ctx.epoch_time,
    "aws:tokenissuetime": lambda ctx: ctx.token_issue_time,
    "aws:multifactorauthpresent": lambda ctx: ctx.mfa_authenticated,
    "aws:securetransport": lambda ctx: ctx.secure_transport,
    "aws:viaawsservice": lambda ctx: ctx.via_aws_service,
    "aws:sourcevpc": lambda ctx: ctx.source_vpc,
    "aws:sourcevpce": lambda ctx: ctx.source_vpce,
    "aws:calledvia": lambda ctx: ctx.called_via,
    "aws:calledviaservice": lambda ctx: ctx.called_via,
    "aws:tagkeys": lambda ctx: ctx.tag_keys,
    "aws:principalorgid": lambda ctx: ctx.principal_org_id,
    "aws:principalorgpaths": lambda ctx: ctx.principal_org_paths,
}

# Keys that are valid AWS condition keys but require data we cannot obtain
# from CloudTrail events alone
KNOWN_UNOBTAINABLE_KEYS: set[str] = {
    "aws:principalorgid",
    "aws:principalorgpaths",
    "aws:principaltag/",
    "aws:multifactorauthage",
    "aws:referer",
    "aws:fatalerror",
    "aws:calledalias",
}

# Lazy-initialized merged extractors (manual + auto-generated)
_all_service_extractors: dict[str, Callable[[EvaluationContext], Any]] | None = None


def _derive_param_candidates(key_part: str) -> list[str]:
    """Derive requestParameter key candidates from a condition key name part.

    Handles PascalCase keys like "DirectInternetAccess" -> ["directInternetAccess", ...]
    and hyphenated keys like "auto-assign-public-ip" -> ["autoAssignPublicIp", ...].
    """
    candidates: list[str] = []

    if "-" in key_part:
        # Hyphenated: "auto-assign-public-ip" -> "autoAssignPublicIp", "AutoAssignPublicIp"
        parts = key_part.split("-")
        camel = parts[0].lower() + "".join(p.capitalize() for p in parts[1:])
        pascal = "".join(p.capitalize() for p in parts)
        candidates.extend([camel, pascal, key_part])
    elif key_part:
        # PascalCase: "DirectInternetAccess" -> "directInternetAccess", "DirectInternetAccess"
        camel = key_part[0].lower() + key_part[1:]
        candidates.extend([camel, key_part])

    return candidates


def _build_auto_extractors() -> dict[str, Callable[[EvaluationContext], Any]]:
    """Build condition key extractors auto-derived from IAM reference data.

    For each service condition key in the IAM reference, creates an extractor
    that tries to look up the corresponding value in requestParameters using
    derived parameter name candidates.

    Skips:
    - Keys already in SERVICE_KEY_EXTRACTORS (manual always wins)
    - Tag-based keys containing '/' or '$' (need external data)
    - Keys with nested colons (e.g., sagemaker:ModelLifeCycle:Stage)
    """
    from ..data.iam_reference import get_iam_reference

    ref = get_iam_reference()
    extractors: dict[str, Callable[[EvaluationContext], Any]] = {}

    for service in ref.services.values():
        for key_name in service.condition_keys:
            lower_key = key_name.lower()

            # Skip keys already manually defined
            if lower_key in SERVICE_KEY_EXTRACTORS:
                continue

            # Skip tag-based and template keys
            if "$" in key_name or "/" in key_name:
                continue

            # Must have service:key format
            if ":" not in key_name:
                continue
            _, key_part = key_name.split(":", 1)

            # Skip nested colon keys (e.g., sagemaker:ModelLifeCycle:Stage)
            if ":" in key_part:
                continue

            candidates = _derive_param_candidates(key_part)
            if not candidates:
                continue

            # Capture candidates in closure via default arg
            def _make_extractor(
                cands: list[str],
            ) -> Callable[[EvaluationContext], Any]:
                return lambda ctx: _get_request_param(ctx, *cands)

            extractors[lower_key] = _make_extractor(candidates)

    return extractors


def _get_all_service_extractors() -> dict[str, Callable[[EvaluationContext], Any]]:
    """Get the merged dict of manual + auto-generated service key extractors.

    Auto-generated extractors are built lazily on first call, then cached.
    Manual extractors (SERVICE_KEY_EXTRACTORS) always override auto-generated.
    """
    global _all_service_extractors
    if _all_service_extractors is None:
        auto = _build_auto_extractors()
        # Manual overrides auto
        auto.update(SERVICE_KEY_EXTRACTORS)
        _all_service_extractors = auto
    return _all_service_extractors


def _get_request_param(ctx: EvaluationContext, *keys: str) -> Any | None:
    """Helper to get a value from requestParameters with multiple key variations."""
    if not ctx.event.request_parameters:
        return None
    params = ctx.event.request_parameters
    for key in keys:
        if key in params:
            return params[key]
    return None


def _get_additional_event_data(ctx: EvaluationContext, key: str) -> Any | None:
    """Helper to get a value from additionalEventData."""
    if not ctx.event.additional_event_data:
        return None
    return ctx.event.additional_event_data.get(key)


# Service-specific key extractors: lowercase key -> extractor function
# These extract values from CloudTrail requestParameters or other event fields
SERVICE_KEY_EXTRACTORS: dict[str, Callable[[EvaluationContext], Any]] = {
    # S3 condition keys
    "s3:prefix": lambda ctx: _get_request_param(ctx, "prefix", "Prefix"),
    "s3:delimiter": lambda ctx: _get_request_param(ctx, "delimiter", "Delimiter"),
    "s3:max-keys": lambda ctx: _get_request_param(ctx, "maxKeys", "max-keys", "MaxKeys"),
    "s3:x-amz-acl": lambda ctx: _get_request_param(ctx, "x-amz-acl", "acl", "ACL"),
    "s3:x-amz-content-sha256": lambda ctx: _get_request_param(
        ctx, "x-amz-content-sha256", "contentSHA256"
    ),
    "s3:x-amz-copy-source": lambda ctx: _get_request_param(
        ctx, "x-amz-copy-source", "CopySource", "copySource"
    ),
    "s3:x-amz-server-side-encryption": lambda ctx: _get_request_param(
        ctx, "x-amz-server-side-encryption", "serverSideEncryption", "ServerSideEncryption"
    ),
    "s3:x-amz-server-side-encryption-aws-kms-key-id": lambda ctx: _get_request_param(
        ctx, "x-amz-server-side-encryption-aws-kms-key-id", "SSEKMSKeyId", "sseKmsKeyId"
    ),
    "s3:x-amz-storage-class": lambda ctx: _get_request_param(
        ctx, "x-amz-storage-class", "storageClass", "StorageClass"
    ),
    "s3:x-amz-website-redirect-location": lambda ctx: _get_request_param(
        ctx, "x-amz-website-redirect-location", "websiteRedirectLocation"
    ),
    "s3:x-amz-metadata-directive": lambda ctx: _get_request_param(
        ctx, "x-amz-metadata-directive", "metadataDirective", "MetadataDirective"
    ),
    "s3:versionid": lambda ctx: _get_request_param(ctx, "versionId", "VersionId"),
    # EC2 condition keys
    "ec2:instancetype": lambda ctx: _get_request_param(
        ctx, "instanceType", "InstanceType", "instancesSet.items.0.instanceType"
    ),
    "ec2:imageid": lambda ctx: _get_request_param(ctx, "imageId", "ImageId"),
    "ec2:region": lambda ctx: ctx.region,
    "ec2:availabilityzone": lambda ctx: _get_request_param(
        ctx, "availabilityZone", "AvailabilityZone"
    ),
    "ec2:tenancy": lambda ctx: _get_request_param(ctx, "tenancy", "Tenancy"),
    "ec2:volumetype": lambda ctx: _get_request_param(ctx, "volumeType", "VolumeType"),
    "ec2:volumesize": lambda ctx: _get_request_param(ctx, "size", "Size", "volumeSize"),
    "ec2:encrypted": lambda ctx: _get_request_param(ctx, "encrypted", "Encrypted"),
    "ec2:snapshotid": lambda ctx: _get_request_param(ctx, "snapshotId", "SnapshotId"),
    # RDS condition keys
    "rds:databaseclass": lambda ctx: _get_request_param(
        ctx, "dBInstanceClass", "DBInstanceClass", "dbInstanceClass"
    ),
    "rds:databaseengine": lambda ctx: _get_request_param(ctx, "engine", "Engine"),
    "rds:databasename": lambda ctx: _get_request_param(
        ctx, "dBName", "DBName", "databaseName", "DatabaseName"
    ),
    "rds:multi-az": lambda ctx: _get_request_param(ctx, "multiAZ", "MultiAZ"),
    "rds:storagetype": lambda ctx: _get_request_param(ctx, "storageType", "StorageType"),
    "rds:storageencrypted": lambda ctx: _get_request_param(
        ctx, "storageEncrypted", "StorageEncrypted"
    ),
    "rds:vpc": lambda ctx: _get_request_param(ctx, "dBSubnetGroupName", "DBSubnetGroupName")
    is not None,
    # Lambda condition keys
    "lambda:functionarn": lambda ctx: _get_request_param(
        ctx, "functionName", "FunctionName", "functionArn", "FunctionArn"
    ),
    "lambda:layer": lambda ctx: _get_request_param(ctx, "layerName", "LayerName"),
    "lambda:runtime": lambda ctx: _get_request_param(ctx, "runtime", "Runtime"),
    "lambda:codeuri": lambda ctx: _get_request_param(ctx, "code", "Code"),
    # KMS condition keys
    "kms:viaservice": lambda ctx: _get_additional_event_data(ctx, "viaService"),
    "kms:callerarn": lambda ctx: ctx.principal_arn,
    "kms:encryptioncontext": lambda ctx: _get_request_param(
        ctx, "encryptionContext", "EncryptionContext"
    ),
    "kms:grantisforawsresource": lambda ctx: _get_request_param(ctx, "grantIsForAWSResource"),
    "kms:retiringgrantprincipal": lambda ctx: _get_request_param(
        ctx, "retiringPrincipal", "RetiringPrincipal"
    ),
    # DynamoDB condition keys
    "dynamodb:tablename": lambda ctx: _get_request_param(ctx, "tableName", "TableName"),
    "dynamodb:returnconsumedcapacity": lambda ctx: _get_request_param(
        ctx, "returnConsumedCapacity", "ReturnConsumedCapacity"
    ),
    "dynamodb:returnvalues": lambda ctx: _get_request_param(ctx, "returnValues", "ReturnValues"),
    "dynamodb:projection": lambda ctx: _get_request_param(
        ctx, "projectionExpression", "ProjectionExpression"
    ),
    "dynamodb:select": lambda ctx: _get_request_param(ctx, "select", "Select"),
    # SNS condition keys
    "sns:endpoint": lambda ctx: _get_request_param(ctx, "endpoint", "Endpoint"),
    "sns:protocol": lambda ctx: _get_request_param(ctx, "protocol", "Protocol"),
    # SQS condition keys
    "sqs:delayseconds": lambda ctx: _get_request_param(ctx, "delaySeconds", "DelaySeconds"),
    "sqs:visibilitytimeout": lambda ctx: _get_request_param(
        ctx, "visibilityTimeout", "VisibilityTimeout"
    ),
    # Secrets Manager condition keys
    "secretsmanager:secretid": lambda ctx: _get_request_param(ctx, "secretId", "SecretId"),
    "secretsmanager:name": lambda ctx: _get_request_param(ctx, "name", "Name"),
    "secretsmanager:versionid": lambda ctx: _get_request_param(ctx, "versionId", "VersionId"),
    "secretsmanager:versionstage": lambda ctx: _get_request_param(
        ctx, "versionStage", "VersionStage"
    ),
    # CloudWatch Logs condition keys
    "logs:loggroupname": lambda ctx: _get_request_param(ctx, "logGroupName", "LogGroupName"),
    # IAM condition keys
    "iam:permissionsboundary": lambda ctx: _get_request_param(
        ctx, "permissionsBoundary", "PermissionsBoundary"
    ),
    # STS condition keys
    "sts:rolesessionname": lambda ctx: _get_request_param(
        ctx, "roleSessionName", "RoleSessionName"
    ),
    "sts:externalid": lambda ctx: _get_request_param(ctx, "externalId", "ExternalId"),
    # EKS condition keys
    "eks:cluster-name": lambda ctx: _get_request_param(ctx, "name", "clusterName"),
    "eks:nodegroup-name": lambda ctx: _get_request_param(ctx, "nodegroupName", "NodegroupName"),
    # SSM condition keys
    "ssm:resourcetag/": lambda ctx: None,  # Requires external data
    "ssm:documentname": lambda ctx: _get_request_param(ctx, "documentName", "DocumentName"),
    # ELBv2 condition keys
    "elasticloadbalancing:scheme": lambda ctx: _get_request_param(ctx, "scheme", "Scheme"),
    "elasticloadbalancing:type": lambda ctx: _get_request_param(ctx, "type", "Type"),
    "elasticloadbalancing:ipaddresstype": lambda ctx: _get_request_param(
        ctx, "ipAddressType", "IpAddressType"
    ),
    # CloudWatch condition keys
    "cloudwatch:namespace": lambda ctx: _get_request_param(ctx, "namespace", "Namespace"),
    "cloudwatch:alarmactions": lambda ctx: _get_request_param(ctx, "alarmActions", "AlarmActions"),
    # Redshift condition keys
    "redshift:dbname": lambda ctx: _get_request_param(ctx, "dbName", "DBName"),
    "redshift:dbuser": lambda ctx: _get_request_param(ctx, "dbUser", "DbUser"),
    "redshift:clustertype": lambda ctx: _get_request_param(ctx, "clusterType", "ClusterType"),
    "redshift:numberofnodes": lambda ctx: _get_request_param(ctx, "numberOfNodes", "NumberOfNodes"),
    "redshift:encrypted": lambda ctx: _get_request_param(ctx, "encrypted", "Encrypted"),
    "redshift:nodetype": lambda ctx: _get_request_param(ctx, "nodeType", "NodeType"),
    # SageMaker condition keys
    "sagemaker:instancetypes": lambda ctx: _get_request_param(ctx, "instanceType", "InstanceType"),
    "sagemaker:volumekmskeyid": lambda ctx: _get_request_param(
        ctx, "kmsKeyId", "KmsKeyId", "volumeKmsKeyId", "VolumeKmsKeyId"
    ),
    "sagemaker:rootaccess": lambda ctx: _get_request_param(ctx, "rootAccess", "RootAccess"),
    "sagemaker:directinternetaccess": lambda ctx: _get_request_param(
        ctx, "directInternetAccess", "DirectInternetAccess"
    ),
    # Organizations condition keys
    "organizations:serviceprincipal": lambda ctx: _get_request_param(
        ctx, "servicePrincipal", "ServicePrincipal"
    ),
    # Athena condition keys
    "athena:workgroup": lambda ctx: _get_request_param(ctx, "workGroup", "WorkGroup"),
    # CloudTrail condition keys
    "cloudtrail:eventcategory": lambda ctx: _get_additional_event_data(ctx, "eventCategory"),
    # Config condition keys
    "config:rule": lambda ctx: _get_request_param(ctx, "configRuleName", "ConfigRuleName"),
    # GuardDuty condition keys
    "guardduty:resource/detectorid": lambda ctx: _get_request_param(
        ctx, "detectorId", "DetectorId"
    ),
    # Access Analyzer condition keys
    "access-analyzer:analyzertype": lambda ctx: _get_request_param(
        ctx, "type", "Type", "analyzerType", "AnalyzerType"
    ),
    # WAFv2 condition keys
    "wafv2:scope": lambda ctx: _get_request_param(ctx, "scope", "Scope"),
    # ACM condition keys
    "acm:certificatetransparencylogging": lambda ctx: _get_request_param(
        ctx,
        "certificateTransparencyLoggingPreference",
        "CertificateTransparencyLoggingPreference",
    ),
    "acm:domainnames": lambda ctx: _get_request_param(ctx, "domainName", "DomainName"),
    # Cognito condition keys
    "cognito-idp:userpoolid": lambda ctx: _get_request_param(ctx, "userPoolId", "UserPoolId"),
    "cognito-idp:clientid": lambda ctx: _get_request_param(ctx, "clientId", "ClientId"),
    # AppSync condition keys
    "appsync:visibility": lambda ctx: _get_request_param(ctx, "visibility", "Visibility"),
    # Security Hub condition keys
    "securityhub:targetaccount": lambda ctx: _get_request_param(ctx, "accountId", "AccountId"),
    # Backup condition keys
    "backup:backupvaultname": lambda ctx: _get_request_param(
        ctx, "backupVaultName", "BackupVaultName"
    ),
    # Step Functions condition keys
    "states:statemachinearn": lambda ctx: _get_request_param(
        ctx, "stateMachineArn", "StateMachineArn"
    ),
}

# Backward-compatible flat mapping (used only for legacy compatibility)
CONTEXT_KEY_MAP: dict[str, str] = {
    "aws:SourceIp": "source_ip",
    "aws:sourceip": "source_ip",
    "aws:PrincipalArn": "principal_arn",
    "aws:principalarn": "principal_arn",
    "aws:RequestedRegion": "region",
    "aws:requestedregion": "region",
    "aws:CurrentTime": "request_time",
    "aws:currenttime": "request_time",
    "aws:EpochTime": "epoch_time",
    "aws:epochtime": "epoch_time",
    "aws:PrincipalAccount": "principal_account",
    "aws:principalaccount": "principal_account",
    "aws:PrincipalType": "principal_type",
    "aws:principaltype": "principal_type",
    "aws:UserAgent": "user_agent",
    "aws:useragent": "user_agent",
    "aws:userid": "user_id",
    "aws:username": "username",
    "aws:MultiFactorAuthPresent": "mfa_authenticated",
    "aws:multifactorauthpresent": "mfa_authenticated",
    "aws:SecureTransport": "secure_transport",
    "aws:securetransport": "secure_transport",
    "aws:ViaAWSService": "via_aws_service",
    "aws:viaawsservice": "via_aws_service",
    "aws:SourceVpc": "source_vpc",
    "aws:sourcevpc": "source_vpc",
    "aws:SourceVpce": "source_vpce",
    "aws:sourcevpce": "source_vpce",
    "aws:CalledVia": "called_via",
    "aws:calledvia": "called_via",
    "aws:CalledViaService": "called_via",
    "aws:calledviaservice": "called_via",
    "aws:TokenIssueTime": "token_issue_time",
    "aws:tokenissuetime": "token_issue_time",
    "aws:TagKeys": "tag_keys",
    "aws:tagkeys": "tag_keys",
}


def get_context_value_with_status(context: EvaluationContext, key: str) -> tuple[Any | None, str]:
    """
    Get a value from the evaluation context with resolution status.

    Args:
        context: The evaluation context containing CloudTrail event data
        key: The AWS condition key (e.g., "aws:SourceIp")

    Returns:
        Tuple of (value, status) where status is one of:
        - "resolved": value was found
        - "null": key is known but value is None for this event
        - "unobtainable": key is valid but requires external data
        - "service_specific": service-specific key, can't extract from CloudTrail
        - "unknown": key not recognized
    """
    lower_key = key.lower()

    # Check for tag prefix keys
    if lower_key.startswith("aws:requesttag/"):
        tag_name = key.split("/", 1)[1] if "/" in key else ""
        if context.request_tags and tag_name in context.request_tags:
            return context.request_tags[tag_name], RESOLVED
        return None, NULL

    if lower_key.startswith("aws:resourcetag/"):
        tag_name = key.split("/", 1)[1] if "/" in key else ""
        if context.resource_tags and tag_name in context.resource_tags:
            return context.resource_tags[tag_name], RESOLVED
        return None, NULL

    # Principal tags (from external context enrichment)
    if lower_key.startswith("aws:principaltag/"):
        tag_name = key.split("/", 1)[1] if "/" in key else ""
        if context.principal_tags and tag_name in context.principal_tags:
            return context.principal_tags[tag_name], RESOLVED
        return None, UNOBTAINABLE

    # Check global key extractors BEFORE unobtainable check.
    # This allows external context to resolve keys that are normally
    # unobtainable (e.g. aws:PrincipalOrgId supplied via --context file).
    extractor = GLOBAL_KEY_EXTRACTORS.get(lower_key)
    if extractor is not None:
        value = extractor(context)
        if value is not None:
            return value, RESOLVED
        # Extractor returned None -- check if this key is known unobtainable
        for prefix in KNOWN_UNOBTAINABLE_KEYS:
            if lower_key == prefix or (prefix.endswith("/") and lower_key.startswith(prefix)):
                return None, UNOBTAINABLE
        return None, NULL

    # Check remaining known unobtainable keys (those without extractors)
    for prefix in KNOWN_UNOBTAINABLE_KEYS:
        if lower_key == prefix or (prefix.endswith("/") and lower_key.startswith(prefix)):
            return None, UNOBTAINABLE

    # Check for service-specific keys (e.g., "s3:prefix", "ec2:ResourceTag/")
    if ":" in key and not lower_key.startswith("aws:"):
        # Check if we have an extractor for this service-specific key
        # (includes both manual and auto-generated extractors)
        service_extractor = _get_all_service_extractors().get(lower_key)
        if service_extractor is not None:
            value = service_extractor(context)
            if value is not None:
                return value, RESOLVED
            return None, NULL

        # Check for resource tag patterns that we can partially support
        # These require the resource to have tags in the request
        if "/resourcetag/" in lower_key or ":resourcetag/" in lower_key:
            # Try to extract from request tags if this is a create/tag operation
            tag_name = key.split("/", 1)[1] if "/" in key else ""
            if context.request_tags and tag_name in context.request_tags:
                return context.request_tags[tag_name], RESOLVED
            # Resource tags for existing resources require external data
            return None, SERVICE_SPECIFIC

        return None, SERVICE_SPECIFIC

    return None, UNKNOWN


def get_context_value(context: EvaluationContext, key: str) -> Any | None:
    """
    Get a value from the evaluation context for a given condition key.

    Backward-compatible wrapper around get_context_value_with_status.

    Args:
        context: The evaluation context containing CloudTrail event data
        key: The AWS condition key (e.g., "aws:SourceIp")

    Returns:
        The context value, or None if not available
    """
    value, _status = get_context_value_with_status(context, key)
    return value
