from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, TypedDict, final

import pulumi
import pulumi_aws

from stelvio import context
from stelvio.aws.acm import AcmValidatedDomain, AcmValidatedDomainCustomizationDict
from stelvio.component import Component
from stelvio.dns import DnsProviderNotConfiguredError

if TYPE_CHECKING:
    from stelvio.aws.s3.s3 import Bucket
    from stelvio.dns import Record


# TODO: Consider using internal names for these
# https://www.pulumi.com/registry/packages/aws/api-docs/cloudfront/distribution/#inputs
CloudfrontPriceClass = Literal["PriceClass_100", "PriceClass_200", "PriceClass_All"]


class FunctionAssociation(TypedDict):
    event_type: str
    function_arn: str


@final
@dataclass(frozen=True)
class CloudFrontDistributionResources:
    distribution: pulumi_aws.cloudfront.Distribution
    origin_access_control: pulumi_aws.cloudfront.OriginAccessControl
    cache_policy: pulumi_aws.cloudfront.CachePolicy
    acm_validated_domain: AcmValidatedDomain
    record: Record
    bucket_policy: pulumi_aws.s3.BucketPolicy
    function_associations: list[FunctionAssociation] | None


class CloudFrontDistributionCustomizationDict(TypedDict, total=False):
    distribution: pulumi_aws.cloudfront.DistributionArgs | dict[str, Any] | None
    origin_access_control: pulumi_aws.cloudfront.OriginAccessControlArgs | dict[str, Any] | None
    cache_policy: pulumi_aws.cloudfront.CachePolicyArgs | dict[str, Any] | None
    acm_validated_domain: AcmValidatedDomainCustomizationDict | dict[str, Any]
    record: dict[str, Any] | None  # No specific Pulumi Args type here, because cross cloud compat
    bucket_policy: pulumi_aws.s3.BucketPolicyArgs | dict[str, Any] | None


@final
class CloudFrontDistribution(
    Component[CloudFrontDistributionResources, CloudFrontDistributionCustomizationDict]
):
    def __init__(  # noqa: PLR0913
        self,
        name: str,
        bucket: Bucket,
        price_class: CloudfrontPriceClass = "PriceClass_100",
        custom_domain: str | None = None,
        function_associations: list[FunctionAssociation] | None = None,
        customize: CloudFrontDistributionCustomizationDict | None = None,
    ):
        super().__init__(name, customize=customize)
        self.bucket = bucket
        self.custom_domain = custom_domain
        self.price_class = price_class
        self.function_associations = function_associations or []
        self._resources = None

    def _create_resources(self) -> CloudFrontDistributionResources:
        # Create ACM Validated Domain if custom domain is provided
        acm_validated_domain = None
        if self.custom_domain:
            if context().dns is None:
                raise DnsProviderNotConfiguredError("DNS not configured.")
            # CloudFront requires ACM certificates in us-east-1
            acm_validated_domain = AcmValidatedDomain(
                f"{self.name}-acm-validated-domain",
                domain_name=self.custom_domain,
                customize=self._customize.get("acm_validated_domain"),
                region="us-east-1",
            )

        # Create Origin Access Control for S3
        origin_access_control = pulumi_aws.cloudfront.OriginAccessControl(
            context().prefix(f"{self.name}-oac"),
            **self._customizer(
                "origin_access_control",
                {
                    "description": f"Origin Access Control for {self.name}",
                    "origin_access_control_origin_type": "s3",
                    "signing_behavior": "always",
                    "signing_protocol": "sigv4",
                },
            ),
        )

        # Create Cache Policy
        cache_policy = pulumi_aws.cloudfront.CachePolicy(
            context().prefix(f"{self.name}-cache-policy"),
            **self._customizer(
                "cache_policy",
                {
                    "comment": f"Cache policy for {self.name}",
                    "default_ttl": 300,
                    "max_ttl": 3600,
                    "min_ttl": 0,
                    "parameters_in_cache_key_and_forwarded_to_origin": {
                        "cookies_config": {
                            "cookie_behavior": "none",
                        },
                        "headers_config": {
                            "header_behavior": "whitelist",
                            "headers": {
                                "items": ["If-Modified-Since"],
                            },
                        },
                        "query_strings_config": {
                            "query_string_behavior": "none",
                        },
                        "enable_accept_encoding_gzip": True,
                        "enable_accept_encoding_brotli": True,
                    },
                },
            ),
        )

        # Create CloudFront Distribution
        distribution = pulumi_aws.cloudfront.Distribution(
            context().prefix(self.name),
            **self._customizer(
                "distribution",
                {
                    "aliases": [self.custom_domain] if self.custom_domain else None,
                    "origins": [
                        {
                            "domain_name": self.bucket.resources.bucket.bucket_regional_domain_name,  # noqa: E501
                            "origin_id": f"{self.name}-S3-Origin",
                            "origin_access_control_id": origin_access_control.id,
                        }
                    ],
                    "enabled": True,
                    "is_ipv6_enabled": True,
                    "default_root_object": "index.html",
                    "default_cache_behavior": {
                        "allowed_methods": [
                            "GET",
                            "HEAD",
                            "OPTIONS",
                        ],  # Reduced to read-only methods
                        "cached_methods": ["GET", "HEAD"],
                        "target_origin_id": f"{self.name}-S3-Origin",
                        "compress": True,
                        "viewer_protocol_policy": "redirect-to-https",
                        "cache_policy_id": cache_policy.id,
                        "function_associations": self.function_associations,
                    },
                    "price_class": self.price_class,
                    "restrictions": {
                        "geo_restriction": {
                            "restriction_type": "none",
                        }
                    },
                    "viewer_certificate": {
                        "acm_certificate_arn": acm_validated_domain.resources.certificate.arn,
                        "ssl_support_method": "sni-only",
                        "minimum_protocol_version": "TLSv1.2_2021",
                    }
                    if self.custom_domain
                    else {
                        "cloudfront_default_certificate": True,
                    },
                    "custom_error_responses": [
                        {
                            "error_code": 403,
                            "response_code": 404,
                            "response_page_path": "/error.html",
                            "error_caching_min_ttl": 0,  # Don't cache 403 errors
                        },
                        {
                            "error_code": 404,
                            "response_code": 404,
                            "response_page_path": "/error.html",
                            "error_caching_min_ttl": 300,  # Cache 404s for only 5 minutes
                        },
                    ],
                },
            ),
        )

        # Update S3 bucket policy to allow CloudFront access
        bucket_policy = pulumi_aws.s3.BucketPolicy(
            context().prefix(f"{self.name}-bucket-policy"),
            **self._customizer(
                "bucket_policy",
                {
                    "bucket": self.bucket.resources.bucket.id,
                    "policy": pulumi.Output.all(
                        distribution_arn=distribution.arn,
                        bucket_arn=self.bucket.arn,
                    ).apply(
                        lambda args: pulumi.Output.json_dumps(
                            {
                                "Version": "2012-10-17",
                                "Statement": [
                                    {
                                        "Sid": "AllowCloudFrontServicePrincipal",
                                        "Effect": "Allow",
                                        "Principal": {"Service": "cloudfront.amazonaws.com"},
                                        "Action": "s3:GetObject",
                                        "Resource": f"{args['bucket_arn']}/*",
                                        "Condition": {
                                            "StringEquals": {
                                                "AWS:SourceArn": args["distribution_arn"]
                                            }
                                        },
                                    }
                                ],
                            }
                        )
                    ),
                },
            ),
            opts=pulumi.ResourceOptions(
                depends_on=[distribution]
            ),  # Ensure policy is applied after distribution
        )

        record = None
        if self.custom_domain:
            record = context().dns.create_record(
                resource_name=context().prefix(f"{self.name}-cloudfront-record"),
                **self._customizer(
                    "record",
                    {
                        "name": self.custom_domain,
                        "record_type": "CNAME",
                        "value": distribution.domain_name,
                        "ttl": 1,
                    },
                ),
            )

        pulumi.export(f"cloudfront_{self.name}_domain_name", distribution.domain_name)
        pulumi.export(f"cloudfront_{self.name}_distribution_id", distribution.id)
        pulumi.export(f"cloudfront_{self.name}_arn", distribution.arn)

        if record:
            pulumi.export(f"cloudfront_{self.name}_record_name", record.pulumi_resource.name)

        pulumi.export(f"cloudfront_{self.name}_bucket_policy", bucket_policy.id)

        return CloudFrontDistributionResources(
            distribution,
            origin_access_control,
            cache_policy,
            acm_validated_domain,
            record,
            bucket_policy,
            self.function_associations,
        )
