import * as cdk from 'aws-cdk-lib';
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as s3deploy from 'aws-cdk-lib/aws-s3-deployment';
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
import * as sqs from 'aws-cdk-lib/aws-sqs';
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 secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
import * as ssm from 'aws-cdk-lib/aws-ssm';
import * as events from 'aws-cdk-lib/aws-events';
import * as eventsTargets from 'aws-cdk-lib/aws-events-targets';
import { SqsEventSource } from 'aws-cdk-lib/aws-lambda-event-sources';
import { Construct } from 'constructs';
import * as path from 'path';
import { execSync } from 'child_process';

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

    // ── Context parameters ─────────────────────────────────────────────────
    const telegramBotTokenParamName = this.node.tryGetContext('telegramBotTokenParamName') as string | undefined;
    const braveApiKeyParamName = this.node.tryGetContext('braveApiKeyParamName') as string | undefined;
    const googleOAuthClientParamName = this.node.tryGetContext('googleOAuthClientParamName') as string | undefined;
    const existingWorkspaceBucketName = this.node.tryGetContext('workspaceBucketName') as string | undefined;
    const runtime1Arn = this.node.tryGetContext('runtime1Arn') as string | undefined;

    if (!telegramBotTokenParamName) {
      throw new Error('Context param required: telegramBotTokenParamName');
    }
    if (!braveApiKeyParamName) {
      throw new Error('Context param required: braveApiKeyParamName');
    }
    if (!googleOAuthClientParamName) {
      throw new Error('Context param required: googleOAuthClientParamName');
    }

    // ── SSM Parameters (reference existing SecureString params) ────────────
    const ssmParamArns = [
      `arn:aws:ssm:${this.region}:${this.account}:parameter${telegramBotTokenParamName}`,
      `arn:aws:ssm:${this.region}:${this.account}:parameter${braveApiKeyParamName}`,
      `arn:aws:ssm:${this.region}:${this.account}:parameter${googleOAuthClientParamName}`,
    ];

    const ssmReadPolicy = new iam.PolicyStatement({
      actions: ['ssm:GetParameter'],
      resources: ssmParamArns,
    });

    // ── S3 workspace bucket ────────────────────────────────────────────────
    const workspaceBucket = existingWorkspaceBucketName
      ? s3.Bucket.fromBucketName(this, 'WorkspaceBucket', existingWorkspaceBucketName)
      : new s3.Bucket(this, 'WorkspaceBucket', {
          bucketName: `agent-claw-workspace-${this.account}`,
          removalPolicy: cdk.RemovalPolicy.RETAIN,
          versioned: false,
          encryption: s3.BucketEncryption.S3_MANAGED,
        });

    // Seed workspace files on deploy (only if bucket was created by us)
    if (!existingWorkspaceBucketName) {
      new s3deploy.BucketDeployment(this, 'WorkspaceFiles', {
        sources: [s3deploy.Source.asset(path.join(__dirname, '../../workspace'))],
        destinationBucket: workspaceBucket as s3.Bucket,
      });
    }

    // ── DynamoDB session store ─────────────────────────────────────────────
    const sessionTable = new dynamodb.Table(this, 'SessionStore', {
      tableName: 'agent-claw-sessions',
      partitionKey: { name: 'actor_id', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      timeToLiveAttribute: 'ttl',
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // ── DynamoDB user registry ─────────────────────────────────────────────
    const usersTable = new dynamodb.Table(this, 'UsersTable', {
      tableName: 'agent-claw-users',
      partitionKey: { name: 'actor_id', type: dynamodb.AttributeType.STRING },
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // ── SQS FIFO message queue ─────────────────────────────────────────────
    const messageQueue = new sqs.Queue(this, 'MessageQueue', {
      queueName: 'agent-claw-messages.fifo',
      fifo: true,
      contentBasedDeduplication: false,
      visibilityTimeout: cdk.Duration.seconds(900),
      receiveMessageWaitTime: cdk.Duration.seconds(20),
    });

    // ── Lambda: tg-ingest ─────────────────────────────────────────────────
    const tgIngestFn = new lambda.Function(this, 'TgIngest', {
      functionName: 'agent-claw-tg-ingest',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/tg-ingest')),
      timeout: cdk.Duration.seconds(10),
      memorySize: 128,
      environment: {
        MESSAGE_QUEUE_URL: messageQueue.queueUrl,
        TELEGRAM_BOT_TOKEN_SSM_PARAM: telegramBotTokenParamName,
        TELEGRAM_WEBHOOK_SECRET: '', // set via SSM or direct env after deploy
        ATTACHMENTS_BUCKET_NAME: workspaceBucket.bucketName,
      },
    });
    messageQueue.grantSendMessages(tgIngestFn);
    tgIngestFn.addToRolePolicy(ssmReadPolicy);
    workspaceBucket.grantWrite(tgIngestFn);

    // ── Lambda: agent-runner ───────────────────────────────────────────────
    const agentRunnerFn = new lambda.Function(this, 'AgentRunner', {
      functionName: 'agent-claw-agent-runner',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/agent-runner')),
      timeout: cdk.Duration.seconds(900),
      memorySize: 256,
      environment: {
        SESSION_TABLE_NAME: sessionTable.tableName,
        WORKSPACE_BUCKET_NAME: workspaceBucket.bucketName,
        TELEGRAM_BOT_TOKEN_SSM_PARAM: telegramBotTokenParamName,
        BRAVE_API_KEY_SSM_PARAM: braveApiKeyParamName,
        RUNTIME_1_ARN: runtime1Arn ?? 'PLACEHOLDER_SET_AFTER_RUNTIME_DEPLOY',
        AWS_REGION_NAME: 'us-east-1',
      },
    });

    sessionTable.grantReadWriteData(agentRunnerFn);
    usersTable.grantReadWriteData(agentRunnerFn);
    agentRunnerFn.addEnvironment('USERS_TABLE_NAME', usersTable.tableName);
    workspaceBucket.grantRead(agentRunnerFn);
    agentRunnerFn.addToRolePolicy(ssmReadPolicy);
    messageQueue.grantConsumeMessages(agentRunnerFn);

    // AgentCore invoke permission
    agentRunnerFn.addToRolePolicy(new iam.PolicyStatement({
      actions: ['bedrock-agentcore:InvokeAgentRuntime'],
      resources: ['*'],
    }));

    // SQS event source
    agentRunnerFn.addEventSource(new SqsEventSource(messageQueue, {
      batchSize: 10,
      enabled: true,
    }));

    // ── API Gateway HTTP ───────────────────────────────────────────────────
    const httpApi = new apigatewayv2.HttpApi(this, 'WebhookApi', {
      apiName: 'agent-claw-webhook',
    });

    httpApi.addRoutes({
      path: '/telegram',
      methods: [apigatewayv2.HttpMethod.POST],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'TgIngestIntegration', tgIngestFn
      ),
    });

    // ── AgentCore Runtime 1 ────────────────────────────────────────────────
    // NOTE: AgentCore CDK L2 constructs are in preview. Using CfnResource.
    // The runtime1Arn output below needs to be fed back as context param
    // on subsequent deploys so agent-runner can invoke it.

    // IAM execution role for Runtime 1
    const runtime1Role = new iam.Role(this, 'Runtime1Role', {
      assumedBy: new iam.ServicePrincipal('bedrock-agentcore.amazonaws.com'),
      description: 'Execution role for agent-claw Runtime 1 (main assistant)',
    });
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      actions: [
        'bedrock:InvokeModel',
        'bedrock:InvokeModelWithResponseStream',
      ],
      resources: ['*'],
    }));
    workspaceBucket.grantRead(runtime1Role);
    runtime1Role.addToPolicy(ssmReadPolicy);
    usersTable.grantReadWriteData(runtime1Role);
    // Google secret grants added after workspace_mcp section below
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      actions: [
        'bedrock-agentcore:CreateEvent',
        'bedrock-agentcore:ListEvents',
        'bedrock-agentcore:RetrieveMemoryRecords',
      ],
      resources: ['*'],
    }));

    // AgentCore Runtime 1 resource (CfnResource — L2 in preview)
    // CodeZip: packages src/runtime-1/ as a zip and uploads to S3
    // TODO: Replace with L2 construct when aws-cdk-lib includes it stable
    // For now, the runtime ARN must be created manually or via CLI after first synth
    // and fed back as context param runtime1Arn.

    // ── Google Workspace MCP ──────────────────────────────────────────────

    // workspace-mcp Lambda execution role (import existing — created during initial setup)
    const _workspaceMcpRole = iam.Role.fromRoleName(
      this, 'WorkspaceMcpRole', 'agent-claw-workspace-mcp-role'
    );
    _workspaceMcpRole.addToPrincipalPolicy?.(ssmReadPolicy);
    // Grant workspace-mcp role read access to all per-user Google credential secrets
    (_workspaceMcpRole as iam.Role).addToPrincipalPolicy?.(new iam.PolicyStatement({
      sid: 'PerUserGoogleCredentialsRead',
      actions: ['secretsmanager:GetSecretValue'],
      resources: [`arn:aws:secretsmanager:${this.region}:${this.account}:secret:agent-claw/google-credentials/*`],
    }));

    // workspace-mcp Lambda — import existing (created with zip + layer, no Docker)
    const workspaceMcpFn = lambda.Function.fromFunctionName(
      this, 'WorkspaceMcp', 'agent-claw-workspace-mcp'
    );

    // Function URL — AWS_IAM auth (already created, reference for policy attachment)
    const workspaceMcpFunctionUrl = 'https://25hugrzw4uwtueeg77jsmft6lq0wunmd.lambda-url.us-east-1.on.aws';
    const workspaceMcpMcpUrl = workspaceMcpFunctionUrl + '/mcp';

    // AgentCore execution role — grant InvokeFunctionUrl identity policy
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'WorkspaceMcpInvoke',
      actions: ['lambda:InvokeFunctionUrl'],
      resources: [workspaceMcpFn.functionArn],
      conditions: { StringEquals: { 'lambda:FunctionUrlAuthType': 'AWS_IAM' } },
    }));

    // Grant AgentCore execution role read access to per-user Google credentials (stays in Secrets Manager)
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'PerUserGoogleCredentialsReadRuntime',
      actions: ['secretsmanager:GetSecretValue'],
      resources: [`arn:aws:secretsmanager:${this.region}:${this.account}:secret:agent-claw/google-credentials/*`],
    }));
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'GoogleCredentialsListRuntime',
      actions: ['secretsmanager:ListSecrets'],
      resources: ['*'],
    }));

    // Pass workspace_mcp MCP URL to agent-runner (informational)
    agentRunnerFn.addEnvironment('WORKSPACE_MCP_URL', workspaceMcpMcpUrl);

    // ── OAuth Handler Lambda ───────────────────────────────────────────────
    const oauthHandlerFn = new lambda.Function(this, 'OAuthHandler', {
      functionName: 'agent-claw-oauth-handler',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/oauth-handler')),
      timeout: cdk.Duration.seconds(30),
      memorySize: 128,
      environment: {
        GOOGLE_OAUTH_CLIENT_SSM_PARAM: googleOAuthClientParamName,
        USERS_TABLE_NAME: usersTable.tableName,
        TELEGRAM_BOT_TOKEN_SSM_PARAM: telegramBotTokenParamName,
        // OAUTH_REDIRECT_URI set after API GW URL is known — injected via addEnvironment below
        OAUTH_REDIRECT_URI: 'PLACEHOLDER',
      },
    });
    oauthHandlerFn.addToRolePolicy(ssmReadPolicy);
    usersTable.grantReadWriteData(oauthHandlerFn);
    // Grant OAuth handler write access to per-user credential secrets
    oauthHandlerFn.addToRolePolicy(new iam.PolicyStatement({
      sid: 'PerUserGoogleCredentialsWrite',
      actions: [
        'secretsmanager:CreateSecret',
        'secretsmanager:PutSecretValue',
        'secretsmanager:GetSecretValue',
      ],
      resources: [`arn:aws:secretsmanager:${this.region}:${this.account}:secret:agent-claw/google-credentials/*`],
    }));

    // Add OAuth routes to existing HTTP API
    httpApi.addRoutes({
      path: '/oauth/start',
      methods: [apigatewayv2.HttpMethod.GET],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'OAuthStartIntegration', oauthHandlerFn
      ),
    });
    httpApi.addRoutes({
      path: '/oauth/callback',
      methods: [apigatewayv2.HttpMethod.GET],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'OAuthCallbackIntegration', oauthHandlerFn
      ),
    });

    // workspace-mcp proxy route — no auth (SCP blocks Lambda Function URLs)
    httpApi.addRoutes({
      path: '/workspace/{proxy+}',
      methods: [apigatewayv2.HttpMethod.ANY],
      integration: new apigatewayv2integrations.HttpLambdaIntegration(
        'WorkspaceMcpIntegration', workspaceMcpFn
      ),
    });

    // Set OAUTH_REDIRECT_URI now that we have the API URL
    const oauthRedirectUri = `${httpApi.url}oauth/callback`;
    oauthHandlerFn.addEnvironment('OAUTH_REDIRECT_URI', oauthRedirectUri);

    // Pass OAuth start URL to AgentCore runtime so agent can generate connect links
    const oauthStartUrl = `${httpApi.url}oauth/start`;
    // NOTE: AgentCore runtime env vars are set in agentcore.json / agentcore deploy, not CDK.
    // Output the URL so it can be manually set in agentcore.json OAUTH_START_URL.


    // ── Lambda: heartbeat-runner ──────────────────────────────────────────
    const heartbeatRunnerFn = new lambda.Function(this, 'HeartbeatRunner', {
      functionName: 'agent-claw-heartbeat-runner',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/heartbeat-runner')),
      timeout: cdk.Duration.seconds(60),
      memorySize: 128,
      environment: {
        MESSAGE_QUEUE_URL: messageQueue.queueUrl,
        USERS_TABLE_NAME: usersTable.tableName,
      },
    });
    messageQueue.grantSendMessages(heartbeatRunnerFn);
    usersTable.grantReadData(heartbeatRunnerFn);

    // EventBridge rule: every 30 minutes
    const heartbeatRule = new events.Rule(this, 'HeartbeatRule', {
      ruleName: 'agent-claw-heartbeat',
      schedule: events.Schedule.rate(cdk.Duration.minutes(30)),
    });
    heartbeatRule.addTarget(new eventsTargets.LambdaFunction(heartbeatRunnerFn));

    // ── Lambda: scheduler ─────────────────────────────────────────────────
    const schedulerFn = new lambda.Function(this, 'Scheduler', {
      functionName: 'agent-claw-scheduler',
      runtime: lambda.Runtime.PYTHON_3_12,
      handler: 'handler.handler',
      code: lambda.Code.fromAsset(path.join(__dirname, '../../src/lambdas/scheduler')),
      timeout: cdk.Duration.seconds(30),
      memorySize: 128,
      environment: {
        TELEGRAM_BOT_TOKEN_SSM_PARAM: telegramBotTokenParamName,
      },
    });
    schedulerFn.addToRolePolicy(ssmReadPolicy);
    // Allow EventBridge to invoke the scheduler Lambda
    schedulerFn.addPermission('EventBridgeInvoke', {
      principal: new iam.ServicePrincipal('events.amazonaws.com'),
      sourceArn: `arn:aws:events:${this.region}:${this.account}:rule/agent-claw-reminder-*`,
    });
    // Allow scheduler Lambda to delete its own EventBridge rule
    schedulerFn.addToRolePolicy(new iam.PolicyStatement({
      actions: ['events:RemoveTargets', 'events:DeleteRule'],
      resources: [`arn:aws:events:${this.region}:${this.account}:rule/agent-claw-reminder-*`],
    }));

    // AgentCore runtime: EventBridge + Lambda permissions for scheduling
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'EventBridgeScheduler',
      actions: [
        'events:PutRule', 'events:PutTargets',
        'events:ListRules', 'events:ListTargetsByRule',
        'events:RemoveTargets', 'events:DeleteRule',
      ],
      resources: [`arn:aws:events:us-east-1:*:rule/agent-claw-reminder-*`],
    }));
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'SchedulerLambdaPermission',
      actions: ['lambda:AddPermission', 'lambda:RemovePermission'],
      resources: [schedulerFn.functionArn],
    }));

    // ── AgentCore Runtime 1 — extended permissions ───────────────────────

    // Compute/build
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'ComputeBuild',
      actions: ['codebuild:*', 'ecr:*', 'ecs:*', 'logs:*'],
      resources: ['*'],
    }));

    // Broad read-only across account
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'BroadReadOnly',
      actions: [
        's3:List*', 's3:GetObject',
        'lambda:List*', 'lambda:Get*',
        'cloudformation:Describe*', 'cloudformation:List*',
        'sqs:List*', 'sqs:GetQueueAttributes',
        'ec2:Describe*',
        'ssm:Describe*', 'ssm:List*',
        'ce:GetCostAndUsage', 'ce:GetCostForecast',
      ],
      resources: ['*'],
    }));

    // IAM self-modify — scoped to own role only
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'IamSelfModify',
      actions: ['iam:PutRolePolicy', 'iam:AttachRolePolicy', 'iam:DetachRolePolicy', 'iam:DeleteRolePolicy'],
      resources: [runtime1Role.roleArn],
    }));

    // IAM policy management
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'IamPolicyManagement',
      actions: ['iam:CreatePolicy', 'iam:GetPolicy', 'iam:ListPolicies'],
      resources: ['*'],
    }));

    // SSM read for AWS MCP URL
    runtime1Role.addToPolicy(new iam.PolicyStatement({
      sid: 'AwsMcpUrlSsmRead',
      actions: ['ssm:GetParameter', 'ssm:GetParameters'],
      resources: [`arn:aws:ssm:${this.region}:${this.account}:parameter/agent-claw/aws-mcp-url`],
    }));

    // ── Outputs ────────────────────────────────────────────────────────────

    new cdk.CfnOutput(this, 'WorkspaceMcpFunctionUrl', {
      value: workspaceMcpFunctionUrl,
      description: 'workspace-mcp Lambda Function URL (MCP endpoint for Gmail/Calendar)',
    });
    new cdk.CfnOutput(this, 'OAuthStartUrl', {
      value: oauthStartUrl,
      description: 'Google OAuth start URL — set as OAUTH_START_URL in agentcore.json',
    });
    new cdk.CfnOutput(this, 'OAuthRedirectUri', {
      value: oauthRedirectUri,
      description: 'Google OAuth redirect URI — register in Google Cloud Console',
    });

    new cdk.CfnOutput(this, 'WebhookUrl', {
      value: `${httpApi.url}telegram`,
      description: 'Register this URL with Telegram BotFather as webhook endpoint',
    });

    new cdk.CfnOutput(this, 'WorkspaceBucketName', {
      value: workspaceBucket.bucketName,
      description: 'S3 bucket containing agent workspace files',
    });

    new cdk.CfnOutput(this, 'SessionTableName', {
      value: sessionTable.tableName,
      description: 'DynamoDB table for session mapping',
    });

    new cdk.CfnOutput(this, 'UsersTableName', {
      value: usersTable.tableName,
      description: 'DynamoDB user registry table',
    });

    new cdk.CfnOutput(this, 'MessageQueueUrl', {
      value: messageQueue.queueUrl,
      description: 'SQS FIFO queue for incoming messages',
    });

    new cdk.CfnOutput(this, 'Runtime1RoleArn', {
      value: runtime1Role.roleArn,
      description: 'IAM execution role ARN for AgentCore Runtime 1',
    });

    new cdk.CfnOutput(this, 'SchedulerLambdaArn', {
      value: schedulerFn.functionArn,
      description: 'Scheduler Lambda ARN — set as SCHEDULER_LAMBDA_ARN in agentcore.json',
    });
  }
}