import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as elbv2 from 'aws-cdk-lib/aws-elasticloadbalancingv2';
import * as autoscaling from 'aws-cdk-lib/aws-autoscaling';
import { config } from '../../config';

export interface DataLayerProxyStackProps extends cdk.StackProps {
  accountStage: string;
  allowedClientCidrs?: string[];
  tcpPort?: number; // default 6379
  // Backends (all required)
  defaultValkeyEndpoint: string;
  recValkeyEndpoint: string;
  orpheusValkeyEndpoint: string;
  hookListenHistoryValkeyEndpoint: string;
  listenHistoryValkeyEndpoint: string;
  desiredCapacity?: number;
}

export class DataLayerProxyStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props: DataLayerProxyStackProps) {
    super(scope, id, props);

    const accountStage = props.accountStage;
    const port = props.tcpPort ?? 6379;
    const desiredCapacity = props.desiredCapacity ?? 2;

    // VPC and subnets
    const vpc = ec2.Vpc.fromVpcAttributes(this, 'Vpc', {
        vpcId: config[props.accountStage].vpc,
        availabilityZones: ['us-east-2a', 'us-east-2b', 'us-east-2c'],
        publicSubnetIds: config[props.accountStage].publicSubnets,
        privateSubnetIds: config[props.accountStage].dbSubnets,
        vpcCidrBlock: accountStage === 'prod' ? '10.1.0.0/16' : '10.0.0.0/16',
    });
    const privateDbSubnets = config[accountStage].dbSubnets.map((subnetId: string) =>
      ec2.Subnet.fromSubnetId(this, `DbSubnet-${subnetId}`, subnetId)
    );

    // Shared instance Security Group
    const proxySg = new ec2.SecurityGroup(this, 'DataLayerProxyInstanceSG', {
      vpc,
      description: 'Security group for Data Layer TCP proxy instances',
      allowAllOutbound: true,
    });

    // Allow VPC-internal traffic (incl. NLB health checks) on the TCP port
    proxySg.addIngressRule(ec2.Peer.ipv4(vpc.vpcCidrBlock), ec2.Port.tcp(port), 'VPC internal access');

    // Optional allowlist for external clients
    for (const cidr of props.allowedClientCidrs ?? []) {
      proxySg.addIngressRule(ec2.Peer.ipv4(cidr), ec2.Port.tcp(port), `Allowed client ${cidr}`);
    }

    // Ubuntu 22.04 ARM64 for c8g
    const ubuntuArmAmi = ec2.MachineImage.lookup({
      name: 'ubuntu/images/hvm-ssd/ubuntu-jammy-22.04-arm64-server-*',
      owners: ['099720109477'],
    });

    // Helper to create one proxy unit (ASG + NLB + TG + listener + userData)
    const createProxy = (
      label: 'Default' | 'Rec' | 'Orpheus' | 'HookListenHistory' | 'ListenHistory',
      endpoint: string
    ) => {
      const asg = new autoscaling.AutoScalingGroup(this, `DataLayerProxyAsg${label}`, {
        vpc,
        vpcSubnets: { subnets: privateDbSubnets },
        minCapacity: 4,
        maxCapacity: Math.max(desiredCapacity, 8),
        instanceType: ec2.InstanceType.of(ec2.InstanceClass.C8G, accountStage === 'prod' ? ec2.InstanceSize.LARGE : ec2.InstanceSize.MEDIUM),
        machineImage: ubuntuArmAmi,
        securityGroup: proxySg,
        associatePublicIpAddress: false,
        keyPair: ec2.KeyPair.fromKeyPairName(
          this,
          `DataLayerProxyKeyPair${label}`,
          accountStage === 'staging' ? 'rider-backfill-staging' : 'rider-backfill'
        ),
      });

      const haproxyCfg = `
global
  maxconn 200000
  nbproc 1

defaults
  mode tcp
  timeout connect 5s
  timeout client  300s
  timeout server  300s

frontend fe_valkey_${label.toLowerCase()}
  bind 0.0.0.0:${port}
  default_backend be_valkey_${label.toLowerCase()}

backend be_valkey_${label.toLowerCase()}
  balance roundrobin
  server valkey_${label.toLowerCase()} ${endpoint}:6379 check fall 3 rise 3 maxconn 8000
`;

      asg.addUserData(
        'export DEBIAN_FRONTEND=noninteractive',
        'apt-get update -y || true',
        'apt-get install -y jq curl haproxy || true',
        'if [ -f /etc/default/haproxy ]; then sed -i "s/^ENABLED=.*/ENABLED=1/" /etc/default/haproxy; else echo "ENABLED=1" >/etc/default/haproxy; fi',
        `cat >/etc/haproxy/haproxy.cfg <<'EOF'
${haproxyCfg}
EOF
`,
        'systemctl daemon-reload',
        'systemctl enable haproxy',
        'systemctl restart haproxy'
      );

      const nlb = new elbv2.NetworkLoadBalancer(this, `DataLayerNlb${label}`, {
        vpc,
        loadBalancerName: `data-layer-nlb-${label.toLowerCase()}`,
        internetFacing: true,
        crossZoneEnabled: true,
      });

      const tg = new elbv2.NetworkTargetGroup(this, `DataLayerTg${label}`, {
        vpc,
        port,
        protocol: elbv2.Protocol.TCP,
        targetType: elbv2.TargetType.INSTANCE,
        healthCheck: {
          enabled: true,
          protocol: elbv2.Protocol.TCP,
          port: String(port),
          healthyThresholdCount: 3,
          unhealthyThresholdCount: 3,
        },
        deregistrationDelay: cdk.Duration.seconds(10),
      });
      tg.addTarget(asg);

      nlb.addListener(`TcpListener${label}`, {
        port,
        protocol: elbv2.Protocol.TCP,
        defaultTargetGroups: [tg],
      });

      new cdk.CfnOutput(this, `DataLayerNlb${label}DnsName`, {
        value: nlb.loadBalancerDnsName,
        description: `Public DNS of the ${label} Data Layer NLB`,
      });
    };

    // Create four independent proxy stacks (same port, separate NLBs/ASGs)
    createProxy('Default', props.defaultValkeyEndpoint);
    createProxy('Rec', props.recValkeyEndpoint);
    createProxy('Orpheus', props.orpheusValkeyEndpoint);
    createProxy('HookListenHistory', props.hookListenHistoryValkeyEndpoint);
    createProxy('ListenHistory', props.listenHistoryValkeyEndpoint);

    new cdk.CfnOutput(this, 'DataLayerPort', {
      value: String(port),
      description: 'TCP port exposed on all NLBs',
    });
  }
}