import { Aspects, Tags, IAspect, Stack, StackProps } from 'aws-cdk-lib';
import { Bucket } from 'aws-cdk-lib/aws-s3';
import { Construct, IConstruct } from 'constructs';

export class CdkIssueStack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);
    
    Tags.of(this).add('test-tag', 'test-value');
    new MyConstruct(this, 'myConstruct');
  }
}

class MyConstruct extends Construct {

  constructor(scope: IConstruct, id: string) {
    super(scope, id);

    const stack = Stack.of(scope);
    const s3Bucket = new Bucket(stack, 'bucket-with-tags');
    
    Aspects.of(this).add(new MyDynAspect());
  }
}

class MyDynAspect implements IAspect {
  private processed = false;
  
  public visit(node: IConstruct): void {
    if (!this.processed) {
      //IMPORTANT: The bucket is added to the parent of the construct where the aspect it was initialized
      const stack = Stack.of(node);
      const s3Bucket = new Bucket(stack, 'bucket-without-tags-that-should-have');
      
      this.processed = true;
    }
  }
}

invokeAspects = function (root: IConstruct): void {

    const invokedByPath: { [nodePath: string]: IAspect[] } = {};

    let nestedAspectWarning = false;
    
    // CHANGE: repeat full recursion until no more invocations happen
    let noOfInvocationsInLastFullRecursion;
    do {
        noOfInvocationsInLastFullRecursion = recurse(root, []);
    } while(noOfInvocationsInLastFullRecursion > 0);

    function recurse(construct: IConstruct, inheritedAspects: IAspect[]): number {
        let noOfInvocations = 0;
        const node = construct.node;
        const aspects = Aspects.of(construct);
        const allAspectsHere = [...inheritedAspects ?? [], ...aspects.all];
        const nodeAspectsCount = aspects.all.length;
        for (const aspect of allAspectsHere) {
            let invoked = invokedByPath[node.path];
            if (!invoked) {
                invoked = invokedByPath[node.path] = [];
            }

            if (invoked.includes(aspect)) { continue; }

            aspect.visit(construct);
            noOfInvocations++;

            // if an aspect was added to the node while invoking another aspect it will not be invoked, emit a warning
            // the `nestedAspectWarning` flag is used to prevent the warning from being emitted for every child
            if (!nestedAspectWarning && nodeAspectsCount !== aspects.all.length) {
                Annotations.of(construct).addWarning('We detected an Aspect was added via another Aspect, and will not be applied');
                nestedAspectWarning = true;
            }

            // mark as invoked for this node
            invoked.push(aspect);
        }

        for (const child of construct.node.children) {
            if (!Stage.isStage(child)) {
                noOfInvocations += recurse(child, allAspectsHere);
            }
        }

        return noOfInvocations;
    }
}