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 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 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 { 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;
}

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'] },
          { ResourceType: 'dashboard', 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`],
    }));

    // OpenSearch Serverless data access policy
    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, `arn:aws:iam::${this.account}:root`],
      }]),
    });

    // Bedrock Knowledge Base
    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',
          },
        },
      },
    });

    // 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,
      },
      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],
    }));

    // 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 apiIntegration = new apigatewayv2Integrations.HttpLambdaIntegration(
      'ApiIntegration',
      apiFunction
    );

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

    // 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: httpApi.apiEndpoint,
        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: httpApi.apiEndpoint,
        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: httpApi.apiEndpoint,
        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: httpApi.apiEndpoint,
        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,
      filters: [
        lambda.FilterCriteria.filter({
          body: { jobType: lambda.FilterRule.isEqual('suggestions') },
        }),
      ],
    }));

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

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

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

    // 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 });
  }
}