wrapper.ts

import { CONFIG_DIR } from '../../../lib';
import { CDK_APP_ENTRY, CDK_PROJECT_DIR } from '../../constants';
import { getErrorMessage, isChangesetInProgressError } from '../../errors';
import type { CdkToolkitWrapperOptions, DeployOptions, DestroyOptions, DiffOptions, ListOptions } from './types';
import {
  BaseCredentials,
  BootstrapEnvironments,
  BootstrapStackParameters,
  type ICloudAssemblySource,
  Toolkit,
} from '@aws-cdk/toolkit-lib';
import * as path from 'node:path';

// Type for the assembly returned by synth().produce() - has an async dispose method
interface DisposableAssembly {
  cloudAssembly: { directory: string; stacks: { stackName: string }[] };
  dispose(): Promise<void>;
}

// Type for synth result - may have dispose method
interface SynthResult {
  produce(): Promise<DisposableAssembly>;
  dispose?: () => Promise<void>;
}

/** Sleep helper for retry delays */
function sleep(ms: number): Promise<void> {
  return new Promise(resolve => setTimeout(resolve, ms));
}

/**
 * Wrap an async operation with enhanced error handling.
 * Captures and formats errors with full context for debugging.
 */
async function withErrorContext<T>(context: string, operation: () => Promise<T>): Promise<T> {
  try {
    return await operation();
  } catch (err) {
    const message = getErrorMessage(err);
    const stack = err instanceof Error ? err.stack : undefined;
    const cause = err instanceof Error ? err.cause : undefined;

    const error = new Error(`CDK ${context} failed: ${message}`);
    if (stack) {
      error.stack = `CDK ${context} failed: ${message}\n\nOriginal stack:\n${stack}`;
    }
    if (cause) {
      error.cause = cause;
    }
    throw error;
  }
}

/**
 * Command to execute the CDK app.
 * Matches the "app" command in cdk.json.
 * Requires the CDK project to be built first (tsc).
 */
const CDK_APP_COMMAND = 'node';

/**
 * Thin wrapper around @aws-cdk/toolkit-lib for AgentCore CLI.
 *
 * Provides a simplified interface to the CDK Toolkit, pre-configured
 * to work with the AgentCore CDK project generated by `agentcore create`.
 */
export class CdkToolkitWrapper {
  private readonly projectDir: string;
  private readonly options: CdkToolkitWrapperOptions;

  private toolkit: Toolkit | null = null;
  private cloudAssemblySource: ICloudAssemblySource | null = null;
  // After synth, we store both the synth result and the produced assembly
  // Both may hold locks that need releasing
  private synthResult: SynthResult | null = null;
  private synthesizedAssembly: DisposableAssembly | null = null;
  private synthesizedAssemblyDir: string | null = null;

  constructor(options: CdkToolkitWrapperOptions = {}) {
    this.projectDir = options.projectDir ?? path.join(process.cwd(), CONFIG_DIR, CDK_PROJECT_DIR);
    this.options = options;
  }

  /**
   * Get the path to the CDK project directory.
   */
  getProjectDir(): string {
    return this.projectDir;
  }

  /**
   * Get the CDK app command for this project.
   * Points to the compiled JS file in dist/ (requires tsc build first).
   */
  getCdkAppCommand(): string {
    return `${CDK_APP_COMMAND} ${CDK_APP_ENTRY}`;
  }

  /**
   * Initialize the toolkit and cloud assembly source.
   * Must be called before any CDK operations.
   */
  async initialize(): Promise<void> {
    return withErrorContext(`initialize (project: ${this.projectDir})`, async () => {
      // Use explicit profile, fall back to AWS_PROFILE env var per AWS SDK precedence
      const profile = this.options.profile ?? process.env.AWS_PROFILE;
      const sdkConfig = profile ? { baseCredentials: BaseCredentials.awsCliCompatible({ profile }) } : undefined;

      this.toolkit = new Toolkit({
        ioHost: this.options.ioHost,
        sdkConfig,
      });

      this.cloudAssemblySource = await this.toolkit.fromCdkApp(this.getCdkAppCommand(), {
        workingDirectory: this.projectDir,
      });
    });
  }

  /**
   * Ensure the toolkit is initialized.
   */
  private ensureInitialized(): { toolkit: Toolkit; source: ICloudAssemblySource } {
    if (!this.toolkit || !this.cloudAssemblySource) {
      throw new Error('CdkToolkitWrapper not initialized. Call initialize() first.');
    }
    return { toolkit: this.toolkit, source: this.cloudAssemblySource };
  }

  /**
   * Synthesize the CDK app.
   * Returns stack information and stores the assembly directory for subsequent operations.
   * After calling this, deploy/destroy/diff will use the synthesized assembly
   * instead of re-synthesizing (avoiding lock conflicts).
   */
  async synth(): Promise<{ stackNames: string[]; assemblyDirectory: string }> {
    const { toolkit, source } = this.ensureInitialized();

    // Dispose any previous synth result and assembly to release lock files
    await this.disposeSynthResult();

    const result = await withErrorContext('synth', () => toolkit.synth(source, { validateStacks: true }));

    // Store synth result - it may hold locks
    this.synthResult = result as SynthResult;

    // Produce the assembly to get the directory and stack info
    const assembly = await result.produce();
    this.synthesizedAssembly = assembly;
    this.synthesizedAssemblyDir = assembly.cloudAssembly.directory;

    const stackNames = assembly.cloudAssembly.stacks.map(s => s.stackName);

    return {
      stackNames,
      assemblyDirectory: assembly.cloudAssembly.directory,
    };
  }

  /**
   * Dispose the synth result and produced assembly to release lock files.
   * The synth result may hold locks that the produced assembly doesn't release.
   */
  private async disposeSynthResult(): Promise<void> {
    // First dispose the produced assembly
    if (this.synthesizedAssembly) {
      try {
        await this.synthesizedAssembly.dispose();
      } catch {
        // Ignore disposal errors - best effort cleanup
      }
      this.synthesizedAssembly = null;
      this.synthesizedAssemblyDir = null;
    }

    // Then dispose the synth result itself (may hold the actual locks)
    if (this.synthResult) {
      if (typeof this.synthResult.dispose === 'function') {
        try {
          await this.synthResult.dispose();
        } catch {
          // Ignore disposal errors - best effort cleanup
        }
      }
      this.synthResult = null;
    }
  }

  /**
   * Clean up resources. Call this when done with the wrapper.
   * Disposes synth result, assembly, and source.
   */
  async dispose(): Promise<void> {
    await this.disposeSynthResult();
    await this.disposeSource();
  }

  /**
   * Dispose the cloud assembly source to release locks.
   * Some sources (like CachedCloudAssembly) have dispose methods.
   */
  private async disposeSource(): Promise<void> {
    if (this.cloudAssemblySource) {
      // Check if source has dispose method
      const source = this.cloudAssemblySource as { dispose?: () => Promise<void> };
      if (typeof source.dispose === 'function') {
        try {
          await source.dispose();
        } catch {
          // Ignore disposal errors - best effort cleanup
        }
      }
      this.cloudAssemblySource = null;
      this.toolkit = null;
    }
  }

  /**
   * Get the source to use for operations.
   * Uses synthesized assembly if available (avoids re-synth), otherwise uses original source.
   */
  private async getSourceForOperation(): Promise<ICloudAssemblySource> {
    const { toolkit } = this.ensureInitialized();

    // If we've already synthesized, use the assembly directory to avoid re-synth and lock conflicts
    if (this.synthesizedAssemblyDir) {
      return toolkit.fromAssemblyDirectory(this.synthesizedAssemblyDir);
    }

    // Otherwise use the original source (will trigger synth)
    return this.cloudAssemblySource!;
  }

  /**
   * Deploy stacks to AWS.
   * Includes retry logic for transient changeset conflicts that can occur
   * when multiple operations race on the same stack.
   */
  async deploy(options: DeployOptions = {}) {
    const { toolkit } = this.ensureInitialized();
    const source = await this.getSourceForOperation();

    const maxRetries = 3;
    const baseDelayMs = 5000; // 5 seconds base delay
    let lastError: Error | null = null;

    for (let attempt = 0; attempt < maxRetries; attempt++) {
      try {
        return await withErrorContext('deploy', () => toolkit.deploy(source, { stacks: options.stacks }));
      } catch (err) {
        lastError = err instanceof Error ? err : new Error(String(err));

        // Only retry on changeset-in-progress errors
        if (isChangesetInProgressError(err) && attempt < maxRetries - 1) {
          const delayMs = baseDelayMs * Math.pow(2, attempt); // Exponential backoff: 5s, 10s, 20s
          void this.options.ioHost?.notify?.({
            level: 'warn',
            code: 'CDK_TOOLKIT_W0001',
            message: `Changeset operation in progress, retrying in ${delayMs / 1000}s (attempt ${attempt + 1}/${maxRetries})...`,
            time: new Date(),
            action: 'deploy',
            data: undefined,
          });
          await sleep(delayMs);
          continue;
        }

        // For other errors or final retry, throw immediately
        throw lastError;
      }
    }

    // Should not reach here, but throw last error if we do
    throw lastError ?? new Error('Deploy failed after retries');
  }

  /**
   * Destroy stacks from AWS.
   */
  async destroy(options: DestroyOptions = {}) {
    const { toolkit } = this.ensureInitialized();
    const source = await this.getSourceForOperation();
    return withErrorContext('destroy', () => toolkit.destroy(source, { stacks: options.stacks }));
  }

  /**
   * Show diff between deployed and local stacks.
   */
  async diff(options: DiffOptions = {}) {
    const { toolkit } = this.ensureInitialized();
    const source = await this.getSourceForOperation();
    return withErrorContext('diff', () => toolkit.diff(source, { stacks: options.stacks }));
  }

  /**
   * List stacks in the CDK app.
   */
  async list(options: ListOptions = {}) {
    const { toolkit } = this.ensureInitialized();
    const source = await this.getSourceForOperation();
    return withErrorContext('list', () => toolkit.list(source, { stacks: options.stacks }));
  }

  /**
   * Bootstrap the CDK toolkit stack in the target environment.
   * Creates a KMS CMK for S3 bucket encryption by default.
   * @param environments List of environment strings like `['aws://012345678912/us-east-1']`
   * @param kmsKeyId Optional existing KMS key ID to use instead of creating a new one
   */
  async bootstrap(environments: string[], kmsKeyId?: string) {
    const { toolkit } = this.ensureInitialized();
    const params = kmsKeyId ? { kmsKeyId } : { createCustomerMasterKey: true };
    const options = { parameters: BootstrapStackParameters.withExisting(params) };
    return withErrorContext('bootstrap', () =>
      toolkit.bootstrap(BootstrapEnvironments.fromList(environments), options)
    );
  }
}

/**
 * Create and initialize a CdkToolkitWrapper instance.
 * Convenience function that handles initialization.
 */
export async function createCdkToolkitWrapper(options: CdkToolkitWrapperOptions = {}): Promise<CdkToolkitWrapper> {
  const cdkToolkitWrapper = new CdkToolkitWrapper(options);
  await cdkToolkitWrapper.initialize();
  return cdkToolkitWrapper;
}