import * as cdk from 'aws-cdk-lib';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as apigatewayv2 from 'aws-cdk-lib/aws-apigatewayv2';
import * as apigatewayv2Authorizers from 'aws-cdk-lib/aws-apigatewayv2-authorizers';
import * as apigatewayv2Integrations from 'aws-cdk-lib/aws-apigatewayv2-integrations';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as sqs from 'aws-cdk-lib/aws-sqs';
import * as sns from 'aws-cdk-lib/aws-sns';
import * as cognito from 'aws-cdk-lib/aws-cognito';
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
import * as lambdaEventSources from 'aws-cdk-lib/aws-lambda-event-sources';
import * as bedrock from 'aws-cdk-lib/aws-bedrock';
import * as opensearchserverless from 'aws-cdk-lib/aws-opensearchserverless';
import * as logs from 'aws-cdk-lib/aws-logs';
import * as cloudwatch from 'aws-cdk-lib/aws-cloudwatch';
import * as cloudwatchActions from 'aws-cdk-lib/aws-cloudwatch-actions';
import * as cr from 'aws-cdk-lib/custom-resources';
import { Construct } from 'constructs';

export interface ComputeStackProps extends cdk.StackProps {
  vpc: ec2.IVpc;
  lambdaSecurityGroup: ec2.ISecurityGroup;
  dbSecret: secretsmanager.ISecret;
  uploadsBucket: s3.IBucket;
  generatedBucket: s3.IBucket;
  libraryBucket: s3.IBucket;
  jobsQueue: sqs.IQueue;
  userPool: cognito.IUserPool;
  alarmTopic: sns.ITopic;
  webClientId: string;
  mobileClientId: string;
}

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

    const privateSubnets = { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS };

    // Internal API key for Lambda-to-API calls (stored in Secrets Manager)
    const internalApiKeySecret = new secretsmanager.Secret(this, 'InternalApiKeySecret', {
      secretName: 'proposal-system/internal-api-key',
      generateSecretString: {
        excludePunctuation: true,
        passwordLength: 48,
      },
    });

    // OpenSearch Serverless collection for Bedrock KB vector store
    const ossEncryptionPolicy = new opensearchserverless.CfnSecurityPolicy(this, 'OssEncryptionPolicy', {
      name: 'proposal-system-kb-enc',
      type: 'encryption',
      policy: JSON.stringify({
        Rules: [{ ResourceType: 'collection', Resource: ['collection/proposal-system-kb'] }],
        AWSOwnedKey: true,
      }),
    });

    const ossNetworkPolicy = new opensearchserverless.CfnSecurityPolicy(this, 'OssNetworkPolicy', {
      name: 'proposal-system-kb-net',
      type: 'network',
      policy: JSON.stringify([{
        Rules: [
          { ResourceType: 'collection', Resource: ['collection/proposal-system-kb'] },
        ],
        AllowFromPublic: true,
      }]),
    });

    const ossCollection = new opensearchserverless.CfnCollection(this, 'OssCollection', {
      name: 'proposal-system-kb',
      type: 'VECTORSEARCH',
    });
    ossCollection.addDependency(ossEncryptionPolicy);
    ossCollection.addDependency(ossNetworkPolicy);

    // Bedrock KB execution role
    const kbRole = new iam.Role(this, 'KnowledgeBaseRole', {
      roleName: 'proposal-system-kb-role',
      assumedBy: new iam.ServicePrincipal('bedrock.amazonaws.com'),
    });

    kbRole.addToPolicy(new iam.PolicyStatement({
      actions: ['s3:GetObject', 's3:ListBucket'],
      resources: [props.libraryBucket.bucketArn, `${props.libraryBucket.bucketArn}/*`],
    }));

    kbRole.addToPolicy(new iam.PolicyStatement({
      actions: ['aoss:APIAccessAll'],
      resources: [ossCollection.attrArn],
    }));

    kbRole.addToPolicy(new iam.PolicyStatement({
      actions: ['bedrock:InvokeModel'],
      resources: [`arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v2:0`],
    }));

    // Lambda to pre-create the vector index (retries until AOSS access policy propagates)
    const indexCreatorFn = new lambda.Function(this, 'OssIndexCreator', {
      functionName: 'proposal-system-oss-index-creator',
      runtime: lambda.Runtime.PYTHON_3_12,
      architecture: lambda.Architecture.ARM_64,
      handler: 'app.handler',
      code: lambda.Code.fromAsset('../lambdas/oss-index-creator', {
        bundling: {
          image: lambda.Runtime.PYTHON_3_12.bundlingImage,
          command: [
            'bash', '-c',
            'pip install -r requirements.txt -t /asset-output && cp -au . /asset-output',
          ],
        },
      }),
      timeout: cdk.Duration.minutes(6),
      logRetention: logs.RetentionDays.TWO_MONTHS,
    });

    indexCreatorFn.addToRolePolicy(new iam.PolicyStatement({
      actions: ['aoss:APIAccessAll'],
      resources: [ossCollection.attrArn],
    }));

    const ossDataAccessPolicy = new opensearchserverless.CfnAccessPolicy(this, 'OssDataAccessPolicy', {
      name: 'proposal-system-kb-access',
      type: 'data',
      policy: JSON.stringify([{
        Rules: [
          { ResourceType: 'collection', Resource: ['collection/proposal-system-kb'], Permission: ['aoss:*'] },
          { ResourceType: 'index', Resource: ['index/proposal-system-kb/*'], Permission: ['aoss:*'] },
        ],
        Principal: [kbRole.roleArn, indexCreatorFn.role!.roleArn],
      }]),
    });
    ossDataAccessPolicy.addDependency(ossCollection);

    const indexProvider = new cr.Provider(this, 'OssIndexProvider', {
      onEventHandler: indexCreatorFn,
    });

    const ossIndex = new cdk.CustomResource(this, 'OssIndex', {
      serviceToken: indexProvider.serviceToken,
      properties: {
        Endpoint: ossCollection.attrCollectionEndpoint,
        IndexName: 'proposal-system-index',
        VectorField: 'embedding',
        TextField: 'text',
        MetadataField: 'metadata',
      },
    });
    ossIndex.node.addDependency(ossDataAccessPolicy);

    const knowledgeBase = new bedrock.CfnKnowledgeBase(this, 'KnowledgeBase', {
      name: 'proposal-system-kb',
      roleArn: kbRole.roleArn,
      knowledgeBaseConfiguration: {
        type: 'VECTOR',
        vectorKnowledgeBaseConfiguration: {
          embeddingModelArn: `arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v2:0`,
        },
      },
      storageConfiguration: {
        type: 'OPENSEARCH_SERVERLESS',
        opensearchServerlessConfiguration: {
          collectionArn: ossCollection.attrArn,
          vectorIndexName: 'proposal-system-index',
          fieldMapping: {
            vectorField: 'embedding',
            textField: 'text',
            metadataField: 'metadata',
          },
        },
      },
    });
    knowledgeBase.node.addDependency(ossIndex);

    // KB Data Source (S3 library bucket)
    const dataSource = new bedrock.CfnDataSource(this, 'KbDataSource', {
      name: 'proposal-system-library',
      knowledgeBaseId: knowledgeBase.attrKnowledgeBaseId,
      dataSourceConfiguration: {
        type: 'S3',
        s3Configuration: {
          bucketArn: props.libraryBucket.bucketArn,
        },
      },
      vectorIngestionConfiguration: {
        chunkingConfiguration: {
          chunkingStrategy: 'FIXED_SIZE',
          fixedSizeChunkingConfiguration: {
            maxTokens: 512,
            overlapPercentage: 20,
          },
        },
      },
    });

    // .NET 8 API Lambda
    const apiFunction = new lambda.Function(this, 'ApiFunction', {
      functionName: 'proposal-system-api',
      runtime: lambda.Runtime.DOTNET_8,
      architecture: lambda.Architecture.ARM_64,
      handler: 'ProposalSystem.Api',
      code: lambda.Code.fromAsset('../api/src/ProposalSystem.Api/bin/Release/net8.0/linux-arm64/publish'),
      memorySize: 1024,
      timeout: cdk.Duration.seconds(30),
      vpc: props.vpc,
      vpcSubnets: privateSubnets,
      securityGroups: [props.lambdaSecurityGroup],
      environment: {
        ASPNETCORE_ENVIRONMENT: 'Production',
        DB_SECRET_ARN: props.dbSecret.secretArn,
        UPLOADS_BUCKET: props.uploadsBucket.bucketName,
        GENERATED_BUCKET: props.generatedBucket.bucketName,
        LIBRARY_BUCKET: props.libraryBucket.bucketName,
        JOBS_QUEUE_URL: props.jobsQueue.queueUrl,
        INTERNAL_API_KEY_SECRET_ARN: internalApiKeySecret.secretArn,
        Auth__Authority: `https://cognito-idp.${this.region}.amazonaws.com/${props.userPool.userPoolId}`,
        Auth__ClientId: props.webClientId,
        Auth__CognitoDomain: `proposal-system-seahaven.auth.${this.region}.amazoncognito.com`,
        COGNITO_WEB_CLIENT_ID: props.webClientId,
        COGNITO_MOBILE_CLIENT_ID: props.mobileClientId,
      },
      tracing: lambda.Tracing.ACTIVE,
      logRetention: logs.RetentionDays.TWO_MONTHS,
    });

    // API Lambda permissions
    props.dbSecret.grantRead(apiFunction);
    internalApiKeySecret.grantRead(apiFunction);
    props.uploadsBucket.grantReadWrite(apiFunction);
    props.generatedBucket.grantRead(apiFunction);
    props.jobsQueue.grantSendMessages(apiFunction);

    apiFunction.addToRolePolicy(new iam.PolicyStatement({
      actions: ['cognito-idp:AdminGetUser', 'cognito-idp:AdminListGroupsForUser'],
      resources: [props.userPool.userPoolArn],
    }));

    // Function URL for internal Lambda-to-API calls (bypasses API Gateway JWT authorizer)
    const apiFunctionUrl = apiFunction.addFunctionUrl({
      authType: lambda.FunctionUrlAuthType.NONE,
    });

    // API Gateway HTTP API
    const httpApi = new apigatewayv2.HttpApi(this, 'HttpApi', {
      apiName: 'proposal-system-gateway',
      corsPreflight: {
        allowOrigins: [
          'https://proposals.seahaven.com',
          'http://localhost:5173',
        ],
        allowMethods: [
          apigatewayv2.CorsHttpMethod.GET,
          apigatewayv2.CorsHttpMethod.POST,
          apigatewayv2.CorsHttpMethod.PUT,
          apigatewayv2.CorsHttpMethod.DELETE,
          apigatewayv2.CorsHttpMethod.OPTIONS,
        ],
        allowHeaders: ['Authorization', 'Content-Type', 'X-Requested-With'],
        maxAge: cdk.Duration.hours(1),
      },
    });

    const apiAccessLogGroup = new logs.LogGroup(this, 'ApiAccessLogs', {
      logGroupName: '/aws/apigateway/proposal-system',
      retention: logs.RetentionDays.TWO_MONTHS,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
    });

    const defaultStage = httpApi.defaultStage!.node.defaultChild as apigatewayv2.CfnStage;
    defaultStage.defaultRouteSettings = {
      throttlingBurstLimit: 50,
      throttlingRateLimit: 100,
    };
    defaultStage.accessLogSettings = {
      destinationArn: apiAccessLogGroup.logGroupArn,
      format: JSON.stringify({
        requestId: '$context.requestId',
        ip: '$context.identity.sourceIp',
        method: '$context.httpMethod',
        path: '$context.path',
        status: '$context.status',
        latency: '$context.responseLatency',
        userAgent: '$context.identity.userAgent',
      }),
    };

    const apiIntegration = new apigatewayv2Integrations.HttpLambdaIntegration(
      'ApiIntegration',
      apiFunction
    );

    const jwtAuthorizer = new apigatewayv2Authorizers.HttpJwtAuthorizer(
      'CognitoAuthorizer',
      `https://cognito-idp.${this.region}.amazonaws.com/${props.userPool.userPoolId}`,
      { jwtAudience: [props.webClientId, props.mobileClientId] },
    );

    httpApi.addRoutes({
      path: '/api/health',
      methods: [apigatewayv2.HttpMethod.GET],
      integration: apiIntegration,
    });

    httpApi.addRoutes({
      path: '/api/auth/{proxy+}',
      methods: [apigatewayv2.HttpMethod.POST],
      integration: apiIntegration,
    });

    httpApi.addRoutes({
      path: '/{proxy+}',
      methods: [apigatewayv2.HttpMethod.ANY],
      integration: apiIntegration,
      authorizer: jwtAuthorizer,
    });

    // Python Lambda: Suggestions Engine
    const suggestionsFunction = new lambda.Function(this, 'SuggestionsFunction', {
      functionName: 'proposal-system-suggestions',
      runtime: lambda.Runtime.PYTHON_3_12,
      architecture: lambda.Architecture.ARM_64,
      handler: 'app.handler',
      code: lambda.Code.fromAsset('../lambdas/suggestions'),
      memorySize: 512,
      timeout: cdk.Duration.seconds(60),
      vpc: props.vpc,
      vpcSubnets: privateSubnets,
      securityGroups: [props.lambdaSecurityGroup],
      environment: {
        KNOWLEDGE_BASE_ID: knowledgeBase.attrKnowledgeBaseId,
        MODEL_ID: 'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
        API_BASE_URL: apiFunctionUrl.url,
        INTERNAL_API_KEY_SECRET_ARN: internalApiKeySecret.secretArn,
      },
      logRetention: logs.RetentionDays.TWO_MONTHS,
    });

    internalApiKeySecret.grantRead(suggestionsFunction);
    suggestionsFunction.addToRolePolicy(new iam.PolicyStatement({
      actions: ['bedrock:InvokeModel'],
      resources: [`arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-*`],
    }));
    suggestionsFunction.addToRolePolicy(new iam.PolicyStatement({
      actions: ['bedrock:Retrieve'],
      resources: [knowledgeBase.attrKnowledgeBaseArn],
    }));

    // Python Lambda: PDF Extract
    const pdfExtractFunction = new lambda.Function(this, 'PdfExtractFunction', {
      functionName: 'proposal-system-pdf-extract',
      runtime: lambda.Runtime.PYTHON_3_12,
      architecture: lambda.Architecture.ARM_64,
      handler: 'app.handler',
      code: lambda.Code.fromAsset('../lambdas/pdf-extract'),
      memorySize: 1024,
      timeout: cdk.Duration.seconds(120),
      vpc: props.vpc,
      vpcSubnets: privateSubnets,
      securityGroups: [props.lambdaSecurityGroup],
      environment: {
        UPLOADS_BUCKET: props.uploadsBucket.bucketName,
        MODEL_ID: 'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
        API_BASE_URL: apiFunctionUrl.url,
        INTERNAL_API_KEY_SECRET_ARN: internalApiKeySecret.secretArn,
      },
      logRetention: logs.RetentionDays.TWO_MONTHS,
    });

    internalApiKeySecret.grantRead(pdfExtractFunction);
    props.uploadsBucket.grantRead(pdfExtractFunction);
    pdfExtractFunction.addToRolePolicy(new iam.PolicyStatement({
      actions: ['bedrock:InvokeModel'],
      resources: [`arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-*`],
    }));

    // Python Lambda: PDF Generate
    const pdfGenerateFunction = new lambda.Function(this, 'PdfGenerateFunction', {
      functionName: 'proposal-system-pdf-generate',
      runtime: lambda.Runtime.PYTHON_3_12,
      architecture: lambda.Architecture.ARM_64,
      handler: 'app.handler',
      code: lambda.Code.fromAsset('../lambdas/pdf-generate'),
      memorySize: 512,
      timeout: cdk.Duration.seconds(30),
      vpc: props.vpc,
      vpcSubnets: privateSubnets,
      securityGroups: [props.lambdaSecurityGroup],
      environment: {
        GENERATED_BUCKET: props.generatedBucket.bucketName,
        API_BASE_URL: apiFunctionUrl.url,
        INTERNAL_API_KEY_SECRET_ARN: internalApiKeySecret.secretArn,
      },
      logRetention: logs.RetentionDays.TWO_MONTHS,
    });

    internalApiKeySecret.grantRead(pdfGenerateFunction);
    props.generatedBucket.grantWrite(pdfGenerateFunction);

    // Python Lambda: Library Ingest
    const libraryIngestFunction = new lambda.Function(this, 'LibraryIngestFunction', {
      functionName: 'proposal-system-library-ingest',
      runtime: lambda.Runtime.PYTHON_3_12,
      architecture: lambda.Architecture.ARM_64,
      handler: 'app.handler',
      code: lambda.Code.fromAsset('../lambdas/library-ingest'),
      memorySize: 512,
      timeout: cdk.Duration.seconds(60),
      vpc: props.vpc,
      vpcSubnets: privateSubnets,
      securityGroups: [props.lambdaSecurityGroup],
      environment: {
        LIBRARY_BUCKET: props.libraryBucket.bucketName,
        KNOWLEDGE_BASE_ID: knowledgeBase.attrKnowledgeBaseId,
        DATA_SOURCE_ID: dataSource.attrDataSourceId,
        API_BASE_URL: apiFunctionUrl.url,
        INTERNAL_API_KEY_SECRET_ARN: internalApiKeySecret.secretArn,
      },
      logRetention: logs.RetentionDays.TWO_MONTHS,
    });

    internalApiKeySecret.grantRead(libraryIngestFunction);
    props.libraryBucket.grantWrite(libraryIngestFunction);
    libraryIngestFunction.addToRolePolicy(new iam.PolicyStatement({
      actions: ['bedrock:StartIngestionJob'],
      resources: [knowledgeBase.attrKnowledgeBaseArn],
    }));

    // SQS Event Sources with message filtering
    suggestionsFunction.addEventSource(new lambdaEventSources.SqsEventSource(props.jobsQueue, {
      batchSize: 1,
      reportBatchItemFailures: true,
      filters: [
        lambda.FilterCriteria.filter({
          body: { jobType: lambda.FilterRule.isEqual('suggestions') },
        }),
      ],
    }));

    pdfExtractFunction.addEventSource(new lambdaEventSources.SqsEventSource(props.jobsQueue, {
      batchSize: 1,
      reportBatchItemFailures: true,
      filters: [
        lambda.FilterCriteria.filter({
          body: { jobType: lambda.FilterRule.isEqual('pdf-extract') },
        }),
      ],
    }));

    pdfGenerateFunction.addEventSource(new lambdaEventSources.SqsEventSource(props.jobsQueue, {
      batchSize: 1,
      reportBatchItemFailures: true,
      filters: [
        lambda.FilterCriteria.filter({
          body: { jobType: lambda.FilterRule.isEqual('pdf-generate') },
        }),
      ],
    }));

    libraryIngestFunction.addEventSource(new lambdaEventSources.SqsEventSource(props.jobsQueue, {
      batchSize: 1,
      reportBatchItemFailures: true,
      filters: [
        lambda.FilterCriteria.filter({
          body: { jobType: lambda.FilterRule.isEqual('library-ingest') },
        }),
      ],
    }));

    // CloudWatch Alarms
    const alarmAction = new cloudwatchActions.SnsAction(props.alarmTopic);

    const lambdaFunctions = [
      { fn: apiFunction, name: 'api' },
      { fn: suggestionsFunction, name: 'suggestions' },
      { fn: pdfExtractFunction, name: 'pdf-extract' },
      { fn: pdfGenerateFunction, name: 'pdf-generate' },
      { fn: libraryIngestFunction, name: 'library-ingest' },
    ];

    for (const { fn, name } of lambdaFunctions) {
      const alarm = new cloudwatch.Alarm(this, `LambdaErrors-${name}`, {
        alarmName: `proposal-system-${name}-errors`,
        alarmDescription: `Lambda errors for ${name}`,
        metric: fn.metricErrors({ period: cdk.Duration.minutes(5) }),
        threshold: 1,
        evaluationPeriods: 1,
        treatMissingData: cloudwatch.TreatMissingData.NOT_BREACHING,
      });
      alarm.addAlarmAction(alarmAction);
    }

    const api5xxAlarm = new cloudwatch.Alarm(this, 'Api5xxAlarm', {
      alarmName: 'proposal-system-api-5xx',
      alarmDescription: 'API Gateway 5xx errors',
      metric: new cloudwatch.Metric({
        namespace: 'AWS/ApiGateway',
        metricName: '5xx',
        dimensionsMap: { ApiId: httpApi.httpApiId },
        statistic: 'Sum',
        period: cdk.Duration.minutes(5),
      }),
      threshold: 5,
      evaluationPeriods: 1,
      treatMissingData: cloudwatch.TreatMissingData.NOT_BREACHING,
    });
    api5xxAlarm.addAlarmAction(alarmAction);

    // Outputs
    new cdk.CfnOutput(this, 'ApiEndpoint', { value: httpApi.apiEndpoint });
    new cdk.CfnOutput(this, 'ApiFunctionArn', { value: apiFunction.functionArn });
    new cdk.CfnOutput(this, 'KnowledgeBaseId', { value: knowledgeBase.attrKnowledgeBaseId });
    new cdk.CfnOutput(this, 'DataSourceId', { value: dataSource.attrDataSourceId });
  }
}