import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import * as s3 from "aws-cdk-lib/aws-s3";
import * as glue from "aws-cdk-lib/aws-glue";
import * as iam from "aws-cdk-lib/aws-iam";
import * as athena from "aws-cdk-lib/aws-athena";
import * as firehose from "aws-cdk-lib/aws-kinesisfirehose";
import * as lambda from "aws-cdk-lib/aws-lambda";
import * as logs from "aws-cdk-lib/aws-logs";
import * as dynamodb from "aws-cdk-lib/aws-dynamodb";

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

export class OrpheusInfraStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: OrpheusInfraStackProps) {
    super(scope, id, props);


    // Create the DynamoDB table
    const orpheusTable = new dynamodb.Table(this, 'OrpheusTable', {
      tableName: 'orpheus',
      partitionKey: {
        name: 'item_id',
        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
    });


    // Query per-user by type (user's chats/sessions)
    orpheusTable.addGlobalSecondaryIndex({
      indexName: 'UserTypeIndex',
      partitionKey: { name: 'user_id', type: dynamodb.AttributeType.STRING },
      sortKey: { name: 'type', type: dynamodb.AttributeType.STRING },
      projectionType: dynamodb.ProjectionType.ALL,
    });

    // 1. S3 Bucket for chat logs
    const chatBucket = new s3.Bucket(this, "ChatMessagesBucket", {
      bucketName: "orpheus-chat-messages-bucket" + (props?.accountStage === 'prod' ? "-prod" : ""),
      versioned: true,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    });

    // 2. Glue Database
    const db = new glue.CfnDatabase(this, "ChatMessagesDatabase", {
      catalogId: this.account,
      databaseInput: { name: "orpheus_chat_messages_db" },
    });

    // Simplified: Just hex characters for UUIDs
    const generateHexPrefixes = (): string[] => {
      const hexChars = '0123456789abcdef';
      const combinations: string[] = [];
      
      for (let i = 0; i < hexChars.length; i++) {
        for (let j = 0; j < hexChars.length; j++) {
          combinations.push(hexChars[i] + hexChars[j]);
        }
      }
      return combinations; // 16² = 256 combinations - perfect!
    };

    // 3. Glue Table with 2-Character Partition Projection
    new glue.CfnTable(this, "ChatMessagesTable", {
      catalogId: this.account,
      databaseName: 'orpheus_chat_messages_db',
      tableInput: {
        name: "orpheus_chat_messages",
        tableType: "EXTERNAL_TABLE",
        parameters: { 
          classification: "json",
          "projection.enabled": "true",
          "projection.prefix.type": "enum",
          "projection.prefix.values": generateHexPrefixes().join(","),
          "storage.location.template": `s3://${chatBucket.bucketName}/chat_messages/prefix=\${prefix}/`
        },
        storageDescriptor: {
          columns: [
            { name: "session_id", type: "string" },
            { name: "role", type: "string" },
            { name: "content", type: "string" },
            { name: "finish_reason", type: "string" },
            { name: "tool_calls", type: "string" },
            { name: "message_id", type: "string" },
            { name: "timestamp", type: "timestamp" },
            { name: "user_id", type: "string" },
            { name: "tool_call_id", type: "string" },
            { name: "referenced_clip_id", type: "string" },
          ],
          location: `s3://${chatBucket.bucketName}/chat_messages/`,
          inputFormat: "org.apache.hadoop.mapred.TextInputFormat",
          outputFormat: "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
          serdeInfo: {
            serializationLibrary: "org.openx.data.jsonserde.JsonSerDe",
          },
        },
        partitionKeys: [
          { name: "prefix", type: "string" },
        ],
      },
    });

    // 4. Glue Crawler (auto schema evolution)
    new glue.CfnCrawler(this, "ChatMessagesCrawler", {
      role: new iam.Role(this, "CrawlerRole", {
        assumedBy: new iam.ServicePrincipal("glue.amazonaws.com"),
        managedPolicies: [
          iam.ManagedPolicy.fromAwsManagedPolicyName("service-role/AWSGlueServiceRole"),
        ],
      }).roleArn,
      databaseName: 'orpheus_chat_messages_db',
      targets: { s3Targets: [{ path: `s3://${chatBucket.bucketName}/chat_messages/` }] },
      schedule: { scheduleExpression: "cron(0 0 * * ? *)" }, // daily
    });

    // 5. Athena WorkGroup
    const athenaResultsBucket = new s3.Bucket(this, "AthenaResultsBucket", {
      bucketName: "orpheus-chat-athena-results" + (props?.accountStage === 'prod' ? "-prod" : ""),
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    });

    new athena.CfnWorkGroup(this, "ChatAthenaWorkGroup", {
      name: "orpheus_chat_messages_wg",
      workGroupConfiguration: {
        resultConfiguration: {
          outputLocation: `s3://${athenaResultsBucket.bucketName}/results/`,
        },
      },
    });

    // 6. Lambda function for chat message processing and partitioning
    const chatProcessorLambda = new lambda.Function(this, "ChatProcessorLambda", {
      functionName: "orpheus-chat-processor",
      runtime: lambda.Runtime.PYTHON_3_9,
      handler: "index.lambda_handler",
      timeout: cdk.Duration.seconds(120),
      code: lambda.Code.fromInline(`
import base64
import json
import logging
import re

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def lambda_handler(event, context):
    logger.info("=== ChatProcessor Lambda Started ===")
    logger.info(f"Request ID: {context.aws_request_id}")
    logger.info(f"Function Name: {context.function_name}")
    logger.info(f"Remaining Time: {context.get_remaining_time_in_millis()}ms")
    
    # Log the incoming event structure
    logger.info(f"Event keys: {list(event.keys())}")
    if 'records' in event:
        logger.info(f"Number of records to process: {len(event['records'])}")
    else:
        logger.error("No 'records' key in event!")
        return {'records': []}
    
    output = {'records': []}
    
    for i, record in enumerate(event['records']):
        logger.info(f"--- Processing Record {i+1}/{len(event['records'])} ---")
        logger.info(f"Record ID: {record.get('recordId', 'unknown')}")
        logger.info(f"Record keys: {list(record.keys())}")
        
        try:
            # Log raw record data info
            raw_data = record.get('data', '')
            logger.info(f"Raw data length: {len(raw_data)} characters")
            logger.info(f"Raw data preview: {raw_data[:100]}...")
            
            # Decode the record data
            logger.info("Attempting to decode base64 data...")
            payload = base64.b64decode(raw_data)
            logger.info(f"Decoded payload length: {len(payload)} bytes")
            logger.info(f"Decoded payload preview: {payload[:200]}")
            
            # Parse JSON
            logger.info("Attempting to parse JSON...")
            json_data = json.loads(payload)
            logger.info(f"Parsed JSON keys: {list(json_data.keys())}")
            logger.info(f"Full parsed JSON: {json.dumps(json_data, indent=2)}")
            
            # Extract session_id and process
            session_id = json_data.get('session_id', 'default')
            logger.info(f"Extracted session_id: '{session_id}'")
            
            # Extract first 2 characters for prefix (better distribution vs performance)
            if len(session_id) >= 2:
                prefix = session_id[:2]
                # Sanitize prefix - replace non-alphanumeric with 'x'
                prefix = re.sub(r'[^a-zA-Z0-9]', 'x', prefix)
            else:
                prefix = 'de'  # Default 2-char prefix
            
            logger.info(f"Generated 2-char prefix: {prefix}")
            
            partition_keys = {"prefix": prefix}
            logger.info(f"Partition keys: {partition_keys}")
            
            # Create successful output
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'Ok',
                'metadata': { 'partitionKeys': partition_keys }
            }
            logger.info(f"Successfully processed record {i+1} with prefix '{prefix}'")
            
        except json.JSONDecodeError as e:
            logger.error(f"JSON decode error for record {i+1}: {str(e)}")
            logger.error(f"Failed payload: {payload}")
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'Ok',
                'metadata': { 'partitionKeys': {"prefix": "er"} }
            }
            
        except Exception as e:
            logger.error(f"Unexpected error processing record {i+1}: {str(e)}")
            logger.error(f"Error type: {type(e).__name__}")
            logger.error(f"Record data: {record}")
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'Ok',
                'metadata': { 'partitionKeys': {"prefix": "er"} }
            }
        
        output['records'].append(firehose_record_output)
        logger.info(f"Added record {i+1} to output with result: {firehose_record_output['result']}")
    
    # Log final output summary
    successful_records = sum(1 for r in output['records'] if r['result'] == 'Ok')
    error_records = len(output['records']) - successful_records
    
    logger.info(f"=== Processing Complete ===")
    logger.info(f"Total records processed: {len(output['records'])}")
    logger.info(f"Successful: {successful_records}")
    logger.info(f"Errors: {error_records}")
    logger.info(f"Final output structure: {json.dumps(output, indent=2)}")
    
    return output
      `),
    });

    // 7. Firehose Delivery Stream to S3 (Parquet-ready)
    const firehoseRole = new iam.Role(this, "FirehoseRole", {
      assumedBy: new iam.ServicePrincipal("firehose.amazonaws.com"),
    });
    chatBucket.grantWrite(firehoseRole);
    chatProcessorLambda.grantInvoke(firehoseRole);

    // CloudWatch Logs for Firehose
    const firehoseLogGroup = new logs.LogGroup(this, "ChatFirehoseLogGroup", {
      logGroupName: "/aws/kinesisfirehose/orpheus-chat-firehose",
      retention: logs.RetentionDays.ONE_WEEK,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
    });
    const firehoseLogStream = new logs.LogStream(this, "ChatFirehoseLogStream", {
      logGroup: firehoseLogGroup,
      logStreamName: "S3Delivery",
    });

    // Allow Firehose to write to CloudWatch Logs
    firehoseRole.addToPolicy(new iam.PolicyStatement({
      actions: ["logs:PutLogEvents", "logs:CreateLogStream", "logs:DescribeLogStreams"],
      resources: [
        firehoseLogGroup.logGroupArn,
        `${firehoseLogGroup.logGroupArn}:*`,
      ],
    }));

    new firehose.CfnDeliveryStream(this, "ChatFirehose", {
      deliveryStreamType: "DirectPut",
      deliveryStreamName: "orpheus-chat-firehose",
      extendedS3DestinationConfiguration: {
        bucketArn: chatBucket.bucketArn,
        roleArn: firehoseRole.roleArn,
        prefix: "chat_messages/prefix=!{partitionKeyFromLambda:prefix}/",
        errorOutputPrefix: "errors/!{firehose:error-output-type}/",
        cloudWatchLoggingOptions: {
          enabled: true,
          logGroupName: firehoseLogGroup.logGroupName,
          logStreamName: firehoseLogStream.logStreamName,
        },
        bufferingHints: { intervalInSeconds: 60, sizeInMBs: 64 },
        compressionFormat: "UNCOMPRESSED",
        processingConfiguration: {
          enabled: true,
          processors: [
            {
              type: "Lambda",
              parameters: [
                {
                  parameterName: "LambdaArn",
                  parameterValue: chatProcessorLambda.functionArn,
                },
                {
                  parameterName: "BufferSizeInMBs",
                  parameterValue: "3",
                },
                {
                  parameterName: "BufferIntervalInSeconds",
                  parameterValue: "60",
                }
              ]
            }
          ]
        },
        dynamicPartitioningConfiguration: {
          enabled: true,
          retryOptions: {
            durationInSeconds: 300,
          }
        }
      },
    });

    // ===================== Session Metadata Pipeline (Full Duplication) =====================
    // 1. S3 Bucket for session metadata
    const sessionBucket = new s3.Bucket(this, "SessionMetadataBucket", {
      bucketName: "orpheus-session-metadata-bucket" + (props?.accountStage === 'prod' ? "-prod" : ""),
      versioned: true,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    });

    // 2. Glue Database for session metadata
    const sessionDb = new glue.CfnDatabase(this, "SessionMetadataDatabase", {
      catalogId: this.account,
      databaseInput: { name: "orpheus_session_metadata_db" },
    });

    // 3. Glue Table with 2-Character Partition Projection
    new glue.CfnTable(this, "SessionMetadataTable", {
      catalogId: this.account,
      databaseName: 'orpheus_session_metadata_db',
      tableInput: {
        name: "session_metadata",
        tableType: "EXTERNAL_TABLE",
        parameters: { 
          classification: "json",
          "projection.enabled": "true",
          "projection.prefix.type": "enum",
          "projection.prefix.values": generateHexPrefixes().join(","),
          "storage.location.template": `s3://${sessionBucket.bucketName}/session_metadata/prefix=\${prefix}/`
        },
        storageDescriptor: {
          columns: [
            { name: "id", type: "string" },
            { name: "name", type: "string" },
            { name: "created_at", type: "timestamp" },
            { name: "updated_at", type: "timestamp" },
            { name: "user_id", type: "string" },
            { name: "clip_ids", type: "string" },
          ],
          location: `s3://${sessionBucket.bucketName}/session_metadata/`,
          inputFormat: "org.apache.hadoop.mapred.TextInputFormat",
          outputFormat: "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
          serdeInfo: {
            serializationLibrary: "org.openx.data.jsonserde.JsonSerDe",
          },
        },
        partitionKeys: [
          { name: "prefix", type: "string" },
        ],
      },
    });

    // 4. Glue Crawler (auto schema evolution)
    new glue.CfnCrawler(this, "SessionMetadataCrawler", {
      role: new iam.Role(this, "SessionCrawlerRole", {
        assumedBy: new iam.ServicePrincipal("glue.amazonaws.com"),
        managedPolicies: [
          iam.ManagedPolicy.fromAwsManagedPolicyName("service-role/AWSGlueServiceRole"),
        ],
      }).roleArn,
      databaseName: 'orpheus_session_metadata_db',
      targets: { s3Targets: [{ path: `s3://${sessionBucket.bucketName}/session_metadata/` }] },
      schedule: { scheduleExpression: "cron(0 0 * * ? *)" },
    });

    // 5. Athena WorkGroup & Results Bucket for session metadata
    const sessionAthenaResultsBucket = new s3.Bucket(this, "SessionAthenaResultsBucket", {
      bucketName: "orpheus-session-metadata-athena-results" + (props?.accountStage === 'prod' ? "-prod" : ""),
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    });

    new athena.CfnWorkGroup(this, "SessionAthenaWorkGroup", {
      name: "orpheus_session_metadata_wg",
      workGroupConfiguration: {
        resultConfiguration: {
          outputLocation: `s3://${sessionAthenaResultsBucket.bucketName}/results/`,
        },
      },
    });

    // 6. Lambda function for session metadata processing and partitioning
    const sessionProcessorLambda = new lambda.Function(this, "SessionProcessorLambda", {
      functionName: "orpheus-session-metadata-processor",
      runtime: lambda.Runtime.PYTHON_3_9,
      handler: "index.lambda_handler",
      timeout: cdk.Duration.seconds(120),
      code: lambda.Code.fromInline(`
import base64
import json
import logging
import re

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def lambda_handler(event, context):
    logger.info("=== SessionProcessor Lambda Started ===")
    logger.info(f"Request ID: {context.aws_request_id}")
    logger.info(f"Function Name: {context.function_name}")
    logger.info(f"Remaining Time: {context.get_remaining_time_in_millis()}ms")
    
    logger.info(f"Event keys: {list(event.keys())}")
    if 'records' in event:
        logger.info(f"Number of records to process: {len(event['records'])}")
    else:
        logger.error("No 'records' key in event!")
        return {'records': []}
    
    output = {'records': []}
    
    for i, record in enumerate(event['records']):
        logger.info(f"--- Processing Record {i+1}/{len(event['records'])} ---")
        logger.info(f"Record ID: {record.get('recordId', 'unknown')}")
        logger.info(f"Record keys: {list(record.keys())}")
        
        try:
            raw_data = record.get('data', '')
            logger.info(f"Raw data length: {len(raw_data)} characters")
            logger.info(f"Raw data preview: {raw_data[:100]}...")
            
            logger.info("Attempting to decode base64 data...")
            payload = base64.b64decode(raw_data)
            logger.info(f"Decoded payload length: {len(payload)} bytes")
            logger.info(f"Decoded payload preview: {payload[:200]}")
            
            logger.info("Attempting to parse JSON...")
            json_data = json.loads(payload)
            logger.info(f"Parsed JSON keys: {list(json_data.keys())}")
            logger.info(f"Full parsed JSON: {json.dumps(json_data, indent=2)}")
            
            # Extract id and process
            record_id = json_data.get('id', 'default')
            logger.info(f"Extracted id: '{record_id}'")
            
            if len(record_id) >= 2:
                prefix = record_id[:2]
                prefix = re.sub(r'[^a-zA-Z0-9]', 'x', prefix)
            else:
                prefix = 'de'
            
            logger.info(f"Generated 2-char prefix: {prefix}")
            
            partition_keys = {"prefix": prefix}
            logger.info(f"Partition keys: {partition_keys}")
            
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'Ok',
                'metadata': { 'partitionKeys': partition_keys }
            }
            logger.info(f"Successfully processed record {i+1} with prefix '{prefix}'")
            
        except json.JSONDecodeError as e:
            logger.error(f"JSON decode error for record {i+1}: {str(e)}")
            logger.error(f"Failed payload: {payload}")
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'Ok',
                'metadata': { 'partitionKeys': {"prefix": "er"} }
            }
            
        except Exception as e:
            logger.error(f"Unexpected error processing record {i+1}: {str(e)}")
            logger.error(f"Error type: {type(e).__name__}")
            logger.error(f"Record data: {record}")
            firehose_record_output = {
                'recordId': record['recordId'],
                'data': record['data'],
                'result': 'Ok',
                'metadata': { 'partitionKeys': {"prefix": "er"} }
            }
        
        output['records'].append(firehose_record_output)
        logger.info(f"Added record {i+1} to output with result: {firehose_record_output['result']}")
    
    successful_records = sum(1 for r in output['records'] if r['result'] == 'Ok')
    error_records = len(output['records']) - successful_records
    
    logger.info(f"=== Processing Complete ===")
    logger.info(f"Total records processed: {len(output['records'])}")
    logger.info(f"Successful: {successful_records}")
    logger.info(f"Errors: {error_records}")
    logger.info(f"Final output structure: {json.dumps(output, indent=2)}")
    
    return output
      `),
    });

    // 7. Firehose Delivery Stream to S3 for session metadata
    const sessionFirehoseRole = new iam.Role(this, "SessionFirehoseRole", {
      assumedBy: new iam.ServicePrincipal("firehose.amazonaws.com"),
    });
    sessionBucket.grantWrite(sessionFirehoseRole);
    sessionProcessorLambda.grantInvoke(sessionFirehoseRole);

    // CloudWatch Logs for session Firehose
    const sessionFirehoseLogGroup = new logs.LogGroup(this, "SessionFirehoseLogGroup", {
      logGroupName: "/aws/kinesisfirehose/orpheus-session-metadata-firehose",
      retention: logs.RetentionDays.ONE_WEEK,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
    });
    const sessionFirehoseLogStream = new logs.LogStream(this, "SessionFirehoseLogStream", {
      logGroup: sessionFirehoseLogGroup,
      logStreamName: "S3Delivery",
    });
    sessionFirehoseRole.addToPolicy(new iam.PolicyStatement({
      actions: ["logs:PutLogEvents", "logs:CreateLogStream", "logs:DescribeLogStreams"],
      resources: [
        sessionFirehoseLogGroup.logGroupArn,
        `${sessionFirehoseLogGroup.logGroupArn}:*`,
      ],
    }));

    new firehose.CfnDeliveryStream(this, "SessionMetadataFirehose", {
      deliveryStreamType: "DirectPut",
      deliveryStreamName: "orpheus-session-metadata-firehose",
      extendedS3DestinationConfiguration: {
        bucketArn: sessionBucket.bucketArn,
        roleArn: sessionFirehoseRole.roleArn,
        prefix: "session_metadata/prefix=!{partitionKeyFromLambda:prefix}/",
        errorOutputPrefix: "errors/!{firehose:error-output-type}/",
        cloudWatchLoggingOptions: {
          enabled: true,
          logGroupName: sessionFirehoseLogGroup.logGroupName,
          logStreamName: sessionFirehoseLogStream.logStreamName,
        },
        bufferingHints: { intervalInSeconds: 60, sizeInMBs: 64 },
        compressionFormat: "UNCOMPRESSED",
        processingConfiguration: {
          enabled: true,
          processors: [
            {
              type: "Lambda",
              parameters: [
                {
                  parameterName: "LambdaArn",
                  parameterValue: sessionProcessorLambda.functionArn,
                },
                {
                  parameterName: "BufferSizeInMBs",
                  parameterValue: "3",
                },
                {
                  parameterName: "BufferIntervalInSeconds",
                  parameterValue: "60",
                }
              ]
            }
          ]
        },
        dynamicPartitioningConfiguration: {
          enabled: true,
          retryOptions: {
            durationInSeconds: 300,
          }
        }
      },
    });

    // 7. Grant Athena permissions to read from S3
    chatBucket.addToResourcePolicy(
      new iam.PolicyStatement({
        actions: ["s3:GetObject"],
        resources: [`${chatBucket.bucketArn}/*`],
        principals: [new iam.ServicePrincipal("athena.amazonaws.com")],
      })
    );

    athenaResultsBucket.addToResourcePolicy(
      new iam.PolicyStatement({
        actions: ["s3:PutObject"],
        resources: [`${athenaResultsBucket.bucketArn}/*`],
        principals: [new iam.ServicePrincipal("athena.amazonaws.com")],
      })
    );
  }
}