import * as cdk from "aws-cdk-lib";
import {CustomResource} from "aws-cdk-lib";
import {Construct} from "constructs";
import * as cloudfront from 'aws-cdk-lib/aws-cloudfront';
import {OriginProtocolPolicy} from 'aws-cdk-lib/aws-cloudfront';
import * as eks from "aws-cdk-lib/aws-eks";
import {HelmChart} from "aws-cdk-lib/aws-eks";
import * as cloudfront_origins from 'aws-cdk-lib/aws-cloudfront-origins';
import * as elasticloadbalancingv2 from 'aws-cdk-lib/aws-elasticloadbalancingv2';
import * as cr from 'aws-cdk-lib/custom-resources';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as logs from 'aws-cdk-lib/aws-logs';

interface IdeServicesCloudfrontProps extends cdk.NestedStackProps {
    cluster: eks.Cluster;
    ideServicesChart: HelmChart;
}

export class IdeServicesCloudfront extends cdk.NestedStack {
    deploymentUrl: string;

    constructor(scope: Construct, id: string, props: IdeServicesCloudfrontProps) {
        super(scope, id, props);

        const {cluster, ideServicesChart} = props;
        const vpc = cluster.vpc;

        const findLoadBalancerFunction = this.createFindLoadBalancerFunction();

        // Provision a custom resource provider framework
        const provider = new cr.Provider(this, 'FindLoadBalancerProvider', {
            onEventHandler: findLoadBalancerFunction,
            logGroup: new logs.LogGroup(this, 'FindLoadBalancerLogs', {
                retention: logs.RetentionDays.ONE_DAY,
            }),
        });

        // Create a custom resource to look up the ALB by its DNS name
        const ingressNamePrefix = "k8s-ideservicesgroup-"
        const findLoadBalancer = new CustomResource(this , 'FindLoadBalancer', {
            serviceToken: provider.serviceToken,
            properties: {
                namePrefix: ingressNamePrefix,
            },
        });
        findLoadBalancer.node.addDependency(ideServicesChart)

        // Extract the load balancer ARN and security group ID from the Lambda function's response
        const loadBalancerArn = findLoadBalancer.getAttString('loadBalancerArn');
        const loadBalancerDnsName = findLoadBalancer.getAttString('dnsName');
        const securityGroupId = findLoadBalancer.getAttString('securityGroupId');

        // Import the ALB using its attributes
        const alb = elasticloadbalancingv2.ApplicationLoadBalancer.fromApplicationLoadBalancerAttributes(
            this,
            'ImportedAlb',
            {
                loadBalancerArn: loadBalancerArn,
                loadBalancerDnsName: loadBalancerDnsName,
                securityGroupId: securityGroupId,
                vpc: vpc,
            }
        );

        // An Application Load Balancer as a VPC origin
        const albOrigin = cloudfront_origins.VpcOrigin.withApplicationLoadBalancer(alb, {
            // Optional VPC origin configurations
            domainName: loadBalancerDnsName,
            readTimeout: cdk.Duration.seconds(30),
            keepaliveTimeout: cdk.Duration.seconds(5),
            protocolPolicy: OriginProtocolPolicy.HTTP_ONLY,
        });

        // Create CloudFront distribution with ALB as VPC origin
        const distribution = new cloudfront.Distribution(this, 'AlbCloudFrontDistribution', {
            defaultBehavior: {
                origin: albOrigin,
                originRequestPolicy: cloudfront.OriginRequestPolicy.ALL_VIEWER,
                allowedMethods: cloudfront.AllowedMethods.ALLOW_ALL,
                cachePolicy: cloudfront.CachePolicy.CACHING_DISABLED,
                viewerProtocolPolicy: cloudfront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
            },
            priceClass: cloudfront.PriceClass.PRICE_CLASS_100, // Use only North America and Europe
            enabled: true,
            comment: 'CloudFront distribution with VPC origin for ALB',
        });

        this.deploymentUrl = `https://${distribution.distributionDomainName}`
    }

    private createFindLoadBalancerFunction(): lambda.Function {
        // Create a Lambda function to find the ALB by its DNS name
        const fn = new lambda.Function(this, 'FindLoadBalancerFunction', {
            runtime: lambda.Runtime.NODEJS_22_X,
            handler: 'index.handler',
            code: lambda.Code.fromInline(`
                // Import AWS SDK v3 modules
                const { ElasticLoadBalancingV2Client, DescribeLoadBalancersCommand } = require('@aws-sdk/client-elastic-load-balancing-v2');
                const { EC2Client, DescribeSecurityGroupsCommand } = require('@aws-sdk/client-ec2');
    
                exports.handler = async (event) => {
                    console.log('Event:', JSON.stringify(event));
    
                    const namePrefix = event.ResourceProperties.namePrefix;
                    if (!namePrefix) {
                        throw new Error('ALB name prefix is required');
                    }
    
                    // Initialize AWS clients
                    const elbv2Client = new ElasticLoadBalancingV2Client();
                    const ec2Client = new EC2Client();
    
                    try {
                        // Get all load balancers
                        const describeLoadBalancersCommand = new DescribeLoadBalancersCommand({});
                        const loadBalancersResponse = await elbv2Client.send(describeLoadBalancersCommand);
                        console.log('Load balancers:', JSON.stringify(loadBalancersResponse));
    
                        // Find the load balancer with the matching DNS name
                        const loadBalancer = loadBalancersResponse.LoadBalancers.find(lb => lb.LoadBalancerName.startsWith(namePrefix));
                        if (!loadBalancer) {
                            throw new Error(\`Load balancer with name \${namePrefix} not found\`);
                        }
    
                        return {
                            PhysicalResourceId: loadBalancer.LoadBalancerArn,
                            Data: {
                                loadBalancerArn: loadBalancer.LoadBalancerArn,
                                securityGroupId: loadBalancer.SecurityGroups[0],
                                dnsName: loadBalancer.DNSName,
                            }
                        };
                    } catch (error) {
                        console.error('Error:', error);
                        throw error;
                    }
                };
            `),
            timeout: cdk.Duration.minutes(5),
        });

        // Grant the Lambda function permission to describe load balancers and security groups
        fn.addToRolePolicy(new iam.PolicyStatement({
            actions: [
                'elasticloadbalancing:DescribeLoadBalancers',
                'ec2:DescribeSecurityGroups'
            ],
            resources: ['*'],
        }));

        return fn;
    }
}
