import * as cdk from 'aws-cdk-lib';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as elasticache from 'aws-cdk-lib/aws-elasticache';
import * as memorydb from 'aws-cdk-lib/aws-memorydb';
import * as rds from 'aws-cdk-lib/aws-rds';
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
import { Construct } from 'constructs';
import { config } from '../../config'; // Import the configuration

export interface RedisCacheStackProps extends cdk.StackProps {
  accountStage: string;
}

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

    const vpc = ec2.Vpc.fromVpcAttributes(this, 'SunoMainVpc', {
      vpcId: config[props.accountStage].vpc,
      availabilityZones: ['us-east-2a', 'us-east-2b', 'us-east-2c'],
      privateSubnetIds: config[props.accountStage].dbSubnets,
    });
    // Create a security group for the Redis cluster
    const securityGroup = new ec2.SecurityGroup(this, 'RedisSecurityGroup', {
      vpc: ec2.Vpc.fromVpcAttributes(this, 'Vpc', {
        vpcId: config[props.accountStage].vpc,
        availabilityZones: ['us-east-2a', 'us-east-2b', 'us-east-2c'], // Replace with your availability zones
        privateSubnetIds: config[props.accountStage].dbSubnets,
      }),
      description: 'Allow redis outbound access',
      allowAllOutbound: true,
      securityGroupName: 'RedisSecurityGroup',
    });

    securityGroup.addIngressRule(ec2.Peer.anyIpv4(), ec2.Port.tcp(6379), 'Allow Redis traffic at 6379');

    // ========================= MemoryDB Cluster for Semantic Caching =========================
    // MemoryDB engine configuration for vector search
    // Vector search requirements:
    // - Engine version 7.1+ (required)
    // - Single shard configuration (numShards: 1)
    // - Parameter group with search-enabled: 'yes'
    // - Supported node types: R6g, R7g, T4g
    const memorydbEngine = {
      engine: 'redis',
      engineVersion: '7.1', // Required for vector search
      parameterGroupFamily: 'memorydb_redis7',
    };

    // Create a security group specifically for MemoryDB with external access
    const memorydbSecurityGroup = new ec2.SecurityGroup(this, 'MemoryDBSecurityGroup', {
      vpc: vpc,
      description: 'Allow memoryDB outbound access',
      allowAllOutbound: true,
      securityGroupName: 'MemoryDBSecurityGroup',
    });
    memorydbSecurityGroup.addIngressRule(ec2.Peer.anyIpv4(), ec2.Port.tcp(6379), 'Allow MemoryDB traffic at 6379');

    // Create MemoryDB Subnet Group
    const memorydbSubnetGroup = new memorydb.CfnSubnetGroup(this, 'MemoryDBSubnetGroup', {
      subnetGroupName: `semantic-cache-${props.accountStage}`,
      description: 'Subnet group for MemoryDB semantic cache cluster',
      subnetIds: config[props.accountStage].dbSubnets,
    });

    // Note: 'open-access' is a built-in ACL in MemoryDB, so we reference it directly
    // If you need a custom ACL, create it here. For now, we'll use the built-in 'open-access'

    // Get MemoryDB node type for this stage
    const memoryDbNodeType = config[props.accountStage].memorydbCacheNodes.default;

    // Create MemoryDB Parameter Group with vector search enabled
    // Due AWS limitations, we need to use the default parameter group, but in UI, we need to update to the custom parameter group
    const memorydbParameterGroup = new memorydb.CfnParameterGroup(this, 'MemoryDBParameterGroup', {
      parameterGroupName: `memorydb-redis7-search-${props.accountStage}`,
      description: 'Parameter group for memorydb_redis7 with vector search enabled (search-enabled: yes)',
      family: memorydbEngine.parameterGroupFamily,
      parameters: {
        'search-enabled': 'yes', // Need to changed manually in UI to 'yes'
        'maxmemory-policy': 'allkeys-lru', // Need to changed manually in UI to 'allkeys-lru'
      },
    });

    // Create MemoryDB Cluster
    const memorydbCluster = new memorydb.CfnCluster(this, 'MemoryDBCluster', {
      clusterName: `semantic-cache-${props.accountStage}`,
      nodeType: memoryDbNodeType,
      engine: memorydbEngine.engine,
      engineVersion: memorydbEngine.engineVersion,
      tlsEnabled: true,
      subnetGroupName: memorydbSubnetGroup.ref,
      securityGroupIds: [memorydbSecurityGroup.securityGroupId],
      aclName: 'open-access', // Built-in ACL, no need to create
      parameterGroupName: 'default.memorydb-redis7.search',  // Has to use the default parameter group, but in UI, we need to update to the custom parameter group
      numShards: 1, // Required for vector search (vector search only supports single shard)
      snapshotRetentionLimit: 7,
      snapshotWindow: '03:00-05:00',
    });

    // Retrieve the cacheNodes object from the configuration
    const cacheNodes = config[props.accountStage].cacheNodes;

    // Loop through each key-value pair in the cacheNodes object
    for (const [key, value] of Object.entries(cacheNodes)) {
      const redisCluster = new elasticache.CfnCacheCluster(this, `RedisCluster${key}`, {
        cacheNodeType: value,
        engine: 'redis',
        numCacheNodes: 1,
        vpcSecurityGroupIds: [securityGroup.securityGroupId],
        cacheSubnetGroupName: new elasticache.CfnSubnetGroup(this, `RedisSubnetGroup${key}`, {
          description: `Subnet group for Redis cluster ${key}`,
          subnetIds: config[props.accountStage].dbSubnets,
        }).ref,
        clusterName: `redis-cluster-${key}`,
      });

      // Output the Redis endpoint for each cluster
      new cdk.CfnOutput(this, `RedisEndpoint${key}`, {
        value: redisCluster.attrRedisEndpointAddress,
        exportName: `AWSRedisEndpoint${key}`,
      });
    }

    // ========================= Aurora PostgreSQL (RDS) + Proxy =========================
    // Security Group for PostgreSQL
    const dbSecurityGroup = new ec2.SecurityGroup(this, 'BackendSecondDbSecurityGroup', {
      vpc: vpc,
      description: 'Allow PostgreSQL access',
      allowAllOutbound: true,
      securityGroupName: 'BackendSecondDbSecurityGroup',
    });
    dbSecurityGroup.addIngressRule(ec2.Peer.anyIpv4(), ec2.Port.tcp(5432), 'Allow PostgreSQL traffic');

    // Resolve DB subnets as constructs
    const dbSubnetConstructs = config[props.accountStage].dbSubnets.map((subnetId) =>
      ec2.Subnet.fromSubnetId(this, `BackendDbSubnet-${subnetId}`, subnetId)
    );

    // Aurora PostgreSQL engine (use 16.x for now; upgrade to 17 when region/CDK support it)
    const backendSecondPgEngine = rds.DatabaseClusterEngine.auroraPostgres({
      version: rds.AuroraPostgresEngineVersion.VER_17_5,
    });

    // Custom cluster parameter group (add params as needed)
    const auroraPgParamGroup = new rds.ParameterGroup(this, 'BackendSecondPgParamGroup', {
      engine: backendSecondPgEngine,
      name: 'backend-second-postgresql-17-5',
      description: 'Backend Second PostgreSQL parameter group',
      parameters: {
        // e.g. 'max_connections': '500'
      },
    });

    // Aurora cluster with r8g.4xlarge writer
    const backendSecondCluster = new rds.DatabaseCluster(this, 'BackendSecondCluster', {
      clusterIdentifier: 'backend-second-cluster',
      engine: backendSecondPgEngine,
      
      writer: rds.ClusterInstance.provisioned('WriterInstance', {
        instanceIdentifier: 'backend-second-writer',
        instanceType: ec2.InstanceType.of(ec2.InstanceClass.R8G, props.accountStage === 'prod' ? ec2.InstanceSize.XLARGE4 : ec2.InstanceSize.LARGE),
        publiclyAccessible: false,
      }),
      vpc: vpc,
      vpcSubnets: { subnets: dbSubnetConstructs },
      securityGroups: [dbSecurityGroup],
      parameterGroup: auroraPgParamGroup,
      storageType: rds.DBClusterStorageType.AURORA_IOPT1,
      defaultDatabaseName: 'suno',
      removalPolicy: cdk.RemovalPolicy.RETAIN,
      backup: { retention: cdk.Duration.days(7) },
      cloudwatchLogsExports: ['postgresql'],
    });

    // RDS Proxy in front of the cluster
    const backendSecondProxy = new rds.DatabaseProxy(this, 'BackendSecondProxy', {
      dbProxyName: 'backend-second-proxy',
      proxyTarget: rds.ProxyTarget.fromCluster(backendSecondCluster),
      vpc: vpc,
      securityGroups: [dbSecurityGroup],
      secrets: [backendSecondCluster.secret!],
      requireTLS: false,
      iamAuth: false,
    });

    // Aurora cluster with r8g.4xlarge writer
    const sunoContentCluster = new rds.DatabaseCluster(this, 'SunoContentCluster', {
      clusterIdentifier: 'suno-content-cluster',
      engine: backendSecondPgEngine,
      
      writer: rds.ClusterInstance.provisioned('SunoContentWriterInstance', {
        instanceIdentifier: 'suno-content-writer',
        instanceType: ec2.InstanceType.of(ec2.InstanceClass.R8G, props.accountStage === 'prod' ? ec2.InstanceSize.XLARGE2 : ec2.InstanceSize.XLARGE),
        publiclyAccessible: false,
      }),
      vpc: vpc,
      vpcSubnets: { subnets: dbSubnetConstructs },
      securityGroups: [dbSecurityGroup],
      parameterGroup: auroraPgParamGroup,
      storageType: rds.DBClusterStorageType.AURORA_IOPT1,
      defaultDatabaseName: 'suno_content',
      removalPolicy: cdk.RemovalPolicy.RETAIN,
      backup: { retention: cdk.Duration.days(7) },
      cloudwatchLogsExports: ['postgresql'],
    });

    // RDS Proxy in front of the cluster
    const sunoContentProxy = new rds.DatabaseProxy(this, 'SunoContentProxy', {
      dbProxyName: 'suno-content-proxy',
      proxyTarget: rds.ProxyTarget.fromCluster(sunoContentCluster),
      vpc: vpc,
      securityGroups: [dbSecurityGroup],
      secrets: [sunoContentCluster.secret!],
      requireTLS: false,
      iamAuth: false,
    });

    // Outputs for convenience
    new cdk.CfnOutput(this, 'BackendSecondClusterEndpoint', {
      value: backendSecondCluster.clusterEndpoint.hostname,
      exportName: 'BackendSecondClusterEndpoint',
    });
    new cdk.CfnOutput(this, 'BackendSecondReaderEndpoint', {
      value: backendSecondCluster.clusterReadEndpoint.hostname,
      exportName: 'BackendSecondReaderEndpoint',
    });
    new cdk.CfnOutput(this, 'BackendSecondProxyEndpoint', {
      value: backendSecondProxy.endpoint,
      exportName: 'BackendSecondProxyEndpoint',
    });
    new cdk.CfnOutput(this, 'BackendSecondAdminSecretArn', {
      value: backendSecondCluster.secret!.secretArn,
      exportName: 'BackendSecondAdminSecretArn',
    });

    // Retrieve the cacheNodes object from the configuration
    const valkeyCacheNodes = config[props.accountStage].valkeyCacheNodes;

    // Loop through each key-value pair in the cacheNodes object
    for (const [key, value] of Object.entries(valkeyCacheNodes)) {
      const valkeyCluster = new elasticache.CfnReplicationGroup(this, `ValkeyCluster${key}`, {
        cacheNodeType: value,
        engine: 'valkey',
        engineVersion: '8.0',
        numNodeGroups: 1,
        securityGroupIds: [securityGroup.securityGroupId],
        automaticFailoverEnabled: false,
        replicasPerNodeGroup: 0,
        cacheSubnetGroupName: new elasticache.CfnSubnetGroup(this, `ValkeySubnetGroup${key}`, {
          description: `Subnet group for Valkey cluster ${key}`,
          subnetIds: config[props.accountStage].dbSubnets,
        }).ref,
        replicationGroupId: `valkey-cluster-${key}`,
        replicationGroupDescription: `Non-clustered Valkey instance ${key}`,
        transitEncryptionEnabled: false,
      });
    }
    // Create the DynamoDB table
    const lyricsTable = new dynamodb.Table(this, 'LyricsTable', {
      tableName: 'clip-meta-heavy',
      partitionKey: {
        name: 'clipId',
        type: dynamodb.AttributeType.STRING,
      },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, // On-demand capacity
      removalPolicy: cdk.RemovalPolicy.RETAIN, // Prevents accidental deletion
      pointInTimeRecovery: true, // Enable point-in-time recovery
    });

    // Create the DynamoDB table
    const clipConfigTable = new dynamodb.Table(this, 'ClipConfigTable', {
      tableName: 'clip-gen-config',
      partitionKey: {
        name: 'clipId',
        type: dynamodb.AttributeType.STRING,
      },
      sortKey: {
        name: 'type',
        type: dynamodb.AttributeType.STRING,
      },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, // On-demand capacity
      removalPolicy: cdk.RemovalPolicy.RETAIN, // Prevents accidental deletion
      pointInTimeRecovery: true, // Enable point-in-time recovery
    });

    // Create the DynamoDB table
    const itemInfoTable = new dynamodb.Table(this, 'ItemInfoTable', {
      tableName: 'item-info',
      partitionKey: {
        name: 'itemId',
        type: dynamodb.AttributeType.STRING,
      },
      sortKey: {
        name: 'type',
        type: dynamodb.AttributeType.STRING,
      },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, // On-demand capacity
      removalPolicy: cdk.RemovalPolicy.RETAIN, // Prevents accidental deletion
      pointInTimeRecovery: true, // Enable point-in-time recovery
    });

    // Create the DynamoDB table
    const UserPersonalizationTable = new dynamodb.Table(this, 'UserPersonalizationTable', {
      tableName: 'user-personalization',
      partitionKey: {
        name: 'userId',
        type: dynamodb.AttributeType.STRING,
      },
      sortKey: {
        name: 'type',
        type: dynamodb.AttributeType.STRING,
      },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, // On-demand capacity
      removalPolicy: cdk.RemovalPolicy.RETAIN, // Prevents accidental deletion
      pointInTimeRecovery: true, // Enable point-in-time recovery
    });

    // Create the DynamoDB table
    const queryEmbeddingCacheTable = new dynamodb.Table(this, 'QueryEmbeddingCacheTable', {
      tableName: 'query-embedding-cache',
      partitionKey: {
        name: 'query_hash',
        type: dynamodb.AttributeType.STRING,
      },
      sortKey: {
        name: 'embedding_type',
        type: dynamodb.AttributeType.STRING,
      },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, // On-demand capacity
      removalPolicy: cdk.RemovalPolicy.RETAIN, // Prevents accidental deletion
      pointInTimeRecovery: true, // Enable point-in-time recovery
    });

    // Output the table name
    new cdk.CfnOutput(this, 'LyricsTableName', {
      value: lyricsTable.tableName,
      description: 'The name of the lyrics table',
    });

    // Output the table ARN
    new cdk.CfnOutput(this, 'LyricsTableArn', {
      value: lyricsTable.tableArn,
      description: 'The ARN of the lyrics table',
    });

    // Output the table name
    new cdk.CfnOutput(this, 'ClipGenerationConfigTableName', {
      value: clipConfigTable.tableName,
      description: 'The name of the lyrics table',
    });

    // Output the table ARN
    new cdk.CfnOutput(this, 'ClipGenerationConfigTableArn', {
      value: clipConfigTable.tableArn,
      description: 'The ARN of the lyrics table',
    });
  }
}