import { App, Construct, Stack } from "@aws-cdk/core";
import * as s3 from "@aws-cdk/aws-s3";
import * as ecr from "@aws-cdk/aws-ecr";
import * as iam from "@aws-cdk/aws-iam";
import * as config from "./config";
import * as path from "path";
import * as crypto from "crypto";
import * as mwaaApi from "@aws-sdk/client-mwaa";
export * from "./config";

export interface ProjectResourcesProps {
  readonly name: string;
  readonly envName: string;
  readonly env: config.EnvironmentConfig;
  readonly branch: string;
  readonly commitHash: string;
  readonly cdkOutdir: string;
}

export class ProjectResources {
  /**
   * ProjectResources represents a Conductor projectResources, and manages remotely
   * provisioned AWS resources that are used by DAGs within the projectResources.
   */
  readonly name: string;
  readonly envName: string;
  readonly env: config.EnvironmentConfig;
  readonly branch: string;
  readonly app: App;
  readonly commitHash: string;
  private coreStack!: CoreStack;

  constructor(props: ProjectResourcesProps) {
    // Set instance properties.
    this.name = props.name;
    this.envName = props.envName;
    this.env = props.env;
    this.branch = props.branch;
    this.commitHash = props.commitHash;
    this.app = new App({ outdir: props.cdkOutdir });
  }

  /**
   * createCoreStack initializes the projectResources-level stack
   */
  public async createCoreStack(): Promise<ProjectResources> {
    this.coreStack = new CoreStack(this.app, this);
    await this.coreStack.addResources();
    return this;
  }

  /**
   * bucketName gets the bucket for this projectResources + environment.
   * The bucket stores project related data.
   * @returns {string} Bucket Name
   */
  public bucketName(): string {
    return this.coreStack.bucketName;
  }

  /**
   * getDagResources creates a DAGResource object, which points operators for the
   * given DAG to the appropriate AWS resources.
   * @param {string} name of the DAG
   * @returns {DAGResources} DAGResources instance for the given DAG name
   */
  public dagResources(name: string): DAGResources {
    return new DAGResources(name, this);
  }

  /**
   * ecrUrl gets the ECR+tag for the docker image for this projectResources + environment.
   * The ECR repository is common for the projectResources+environment, and is tagged with the
   * branch and commit hash.
   * @returns {string} ECR URL
   */
  public ecrUrl(): string {
    return this.coreStack.repositoryUriForTag(
      `${this.branch}-${this.commitHash}`
    );
  }

  /**
   * s3UrlForPath generates the s3 path for a given list of path segments by joining the projectResources's
   * auto-generated bucket with the git branch, finally appending the key.
   * @param {Array<string>} pathSegments to make an S3 path for
   * @returns {string} s3 path for the given key
   */
  public s3UrlForPath(pathSegments: Array<string>): string {
    return this.coreStack.s3UrlForObject(
      path.join(this.branch, ...pathSegments)
    );
  }

  /**
   * sagemakerExecutionRole provides the execution role, either one provided
   * manually by the user or one automatically generated.
   * @returns {string} exeuction role
   */
  public sagemakerExecutionRole(): string {
    if (this.env.sagemaker?.executionRole) {
      return this.env.sagemaker.executionRole;
    }
    return this.coreStack.sagemakerExecutionRoleArn();
  }

  /**
   * unloadRoleChain returns the automatically generated IAM role chain associated
   * with this projectResources. It is only usable if the unloadRole that is attached to your
   * Redshift cluster is provided in configuration. See config documentation for details.
   * @returns {string} autogenerated unload role chain
   */
  public unloadRoleChain(): string {
    if (!this.env.redshift?.unloadRole) {
      throw new Error(
        "Cannot unload from Redshift if no unload role is provided by your RedshiftConfig."
      );
    }
    return `${this.env.redshift?.unloadRole},${this.coreStack.unloadRoleArn()}`;
  }

  /**
   * synth is a wrapper for CDK app.synth()
   */
  public synth() {
    this.app.synth();
  }

  public cdkExecutablePath(): string {
    return path.join(__dirname, "..", "node_modules", "aws-cdk", "bin", "cdk");
  }
}

export class DAGResources {
  /**
   * DAGResources provides pointers for DAG level resources to operators that are
   * part of this DAG.
   * @param {string} name of the DAG
   * @param {ProjectResources} projectResources that the DAG belongs to
   */
  readonly name: string;
  readonly projectResources: ProjectResources;

  constructor(name: string, projectResources: ProjectResources) {
    this.name = name;
    this.projectResources = projectResources;
  }

  /**
   * Generates a DAG ID based on projectResources environment. Makes sure that DAGs within
   * the same environment do not have colliding names.
   * @returns {string} DAG ID composed of git_branch.dag_name.envName
   */
  public dagId(): string {
    return `${this.projectResources.name}.${this.projectResources.envName}.${this.name}.${this.projectResources.branch}`;
  }

  /**
   * Gets s3 path for a given of path segments.
   * @param {Array<string>} pathSegments to make an S3 path for.
   * @param {boolean} prefixRunId appends the run ID before any provided path segments.
   * @returns {string} S3 URL composed of the given key prefixed by branch, execution timestamp, and DAG name.
   *    example: s3://projectResources_bucket/branch/dag_name/key
   */
  public s3UrlForPath(
    pathSegments: Array<string>,
    prefixRunId: boolean = true
  ): string {
    if (prefixRunId) {
      return this.projectResources.s3UrlForPath([
        path.join(this.name, "{{run_id}}", ...pathSegments),
      ]);
    }
    return this.projectResources.s3UrlForPath([
      path.join(this.name, ...pathSegments),
    ]);
  }
}

class CoreStack extends Stack {
  /**
   * CoreStack is a CDK stack that generates projectResources level resources for a given Conductor projectResources.
   * The resources created by this stack are shared across all DAGs within the projectResources, and collisons are
   * avoided by sub-naming the resources, i.e. s3 bucket paths, ECR tags.
   * The stack governs provisioning and naming of the resources, but not usage.
   */
  readonly bucketName: string;
  readonly ecrRepoName: string;
  readonly sagemakerExecutionRoleName: string;
  readonly unloadRoleName: string;
  private readonly projectResources: ProjectResources;

  /**
   * @param {Construct} scope this stack belongs to
   * @param {ProjectResources} projectResources that the stack belongs to
   */
  constructor(scope: Construct, projectResources: ProjectResources) {
    super(scope, `${projectResources.name}-core-${projectResources.envName}`, {
      env: {
        account: projectResources.env.accountId,
        region: projectResources.env.defaultRegion,
      },
    });
    this.projectResources = projectResources;
    this.bucketName = `${projectResources.name}.${projectResources.envName}.${this.account}`;
    this.ecrRepoName = `${projectResources.name}.${projectResources.envName}`;
    const hashedBucketName = crypto
      .createHash("sha256")
      .update(this.bucketName)
      .digest("base64")
      .slice(0, 16);
    this.sagemakerExecutionRoleName = `${hashedBucketName}-sm-execution-role-${this.region}`;
    this.unloadRoleName = `${hashedBucketName}-unload-role`;
  }

  public async addResources() {
    const projectResources = this.projectResources;
    const bucket = new s3.Bucket(this, "core-bucket", {
      bucketName: this.bucketName,
    });
    new ecr.Repository(this, "core-ecr-repo", {
      repositoryName: this.ecrRepoName,
    });

    // Create a SageMaker execution role if one isn't already provided.
    if (!projectResources.env.sagemaker?.executionRole) {
      new iam.Role(this, "sagemaker-execution-role", {
        roleName: this.sagemakerExecutionRoleName,
        assumedBy: new iam.CompositePrincipal(
          new iam.ServicePrincipal("sagemaker.amazonaws.com")
        ),
        managedPolicies: [
          iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonS3FullAccess"),
          iam.ManagedPolicy.fromAwsManagedPolicyName(
            "AmazonRedshiftFullAccess"
          ),
          iam.ManagedPolicy.fromAwsManagedPolicyName(
            "AmazonSageMakerFullAccess"
          ),
          iam.ManagedPolicy.fromAwsManagedPolicyName("CloudWatchFullAccess"),
        ],
      });
    }

    // Create a role for Redshift to assume if unloadRole is provided.
    if (projectResources.env.redshift?.unloadRole) {
      const unloadRole = new iam.Role(this, `${this.bucketName}-unload-role`, {
        roleName: this.unloadRoleName,
        assumedBy: new iam.CompositePrincipal(
          new iam.ServicePrincipal("redshift.amazonaws.com"),
          new iam.ArnPrincipal(projectResources.env.redshift.unloadRole)
        ),
      });
      const policyDocument = new iam.PolicyDocument();
      const listBucketStatement = new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
      });
      listBucketStatement.addActions("s3:ListBucket");
      listBucketStatement.addResources(bucket.bucketArn);
      const readBucketStatement = new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
      });
      readBucketStatement.addActions("s3:GetObject");
      readBucketStatement.addResources(`${bucket.bucketArn}/*`);
      const writeBucketStatement = new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
      });
      writeBucketStatement.addActions(
        "s3:AbortMultipartUpload",
        "s3:PutObject"
      );
      writeBucketStatement.addResources(`${bucket.bucketArn}/*`);
      policyDocument.addStatements(
        listBucketStatement,
        readBucketStatement,
        writeBucketStatement
      );
      new iam.Policy(this, `${this.bucketName}-unload-policy`, {
        document: policyDocument,
      }).attachToRole(unloadRole);
    }

    if (projectResources.env.projectBucket) {
      const projectBucketConfig = projectResources.env.projectBucket;
      if (projectBucketConfig.airflowExecutionRoleAccess) {
        const mwaaEnv =
          projectResources.env.airflow.mwaaEnvironmentName &&
          (await CoreStack.getMwaaEnvironment(projectResources.env));
        if (mwaaEnv && mwaaEnv.ExecutionRoleArn) {
          const role = iam.Role.fromRoleArn(
            this,
            "airflow-execution-role",
            mwaaEnv.ExecutionRoleArn
          );
          bucket.grantReadWrite(role);
        }
      }
      projectBucketConfig.readAccessRoles.map((roleArn) => {
        const role = iam.Role.fromRoleArn(this, roleArn, roleArn);
        bucket.grantRead(role);
      });
    }
  }

  private static async getMwaaEnvironment(
    env: config.EnvironmentConfig
  ): Promise<mwaaApi.Environment> {
    const client = new mwaaApi.MWAA({ region: env.defaultRegion });
    const resp = await client.getEnvironment({
      Name: env.airflow.mwaaEnvironmentName,
    });
    if (!resp.Environment) {
      throw new Error(
        `unable to retrieve MWAA env named ${env.airflow.mwaaEnvironmentName}`
      );
    }
    return resp.Environment;
  }

  public sagemakerExecutionRoleArn(): string {
    return `arn:aws:iam::${this.account}:role/${this.sagemakerExecutionRoleName}`;
  }

  public unloadRoleArn(): string {
    return `arn:aws:iam::${this.account}:role/${this.unloadRoleName}`;
  }

  public s3UrlForObject(key: string): string {
    return `s3://${this.bucketName}/${key}`;
  }

  public repositoryUriForTag(tag: string): string {
    return `${this.account}.dkr.ecr.us-west-2.amazonaws.com/${this.ecrRepoName}:${tag}`;
  }
}
