import collections
import logging
import os
import re
from pathlib import Path
from typing import Any, Dict, List

import boto3
import git

from conductor.cli.cmd.deploy import clean_build
from conductor.internal.utils import get_dags_s3_path, load_project


class BranchManager:
    def __init__(
        self,
        project_root,
        s3_bucket,
        s3_prefix,
        airflow_bucket,
        airflow_prefix,
        delimiter,
        dry_run,
    ):
        self.project_root = project_root
        self.s3_client = boto3.client("s3")
        self.s3_resource = boto3.resource("s3")
        self.ecr_client = boto3.client("ecr")

        # parameters
        self.s3_bucket = s3_bucket
        self.s3_prefix = s3_prefix
        self.airflow_bucket = airflow_bucket
        self.airflow_prefix = airflow_prefix
        self.delimiter = delimiter
        self.dry_run = dry_run

        self.branch_manager: Dict[str, List] = {"delete": [], "not_delete": []}
        self.hanging_s3_resources: Dict[str, List] = collections.defaultdict(list)
        self.hanging_airflow_resources: Dict[str, List] = collections.defaultdict(list)
        self.hanging_ecr_resources: Dict[str, List] = collections.defaultdict(list)
        self.hanging_branches = set()

        # prune remote closed branches
        self._prune_remote_branch()
        self.live_branches = self._get_live_branches()
        self.to_be_deleted_branch = []

    def _get_live_branches(self) -> List[str]:
        # Create repo for current directory
        repo = git.Repo(self.project_root)

        # Run "git branch -r" and collect results into array
        remote_live_branches = []
        for ref in repo.git.branch("-r").split("\n"):
            remote_live_branches.append(ref.split("/")[-1])

        return remote_live_branches

    def stage_for_deletion(self, branch):
        # Add branch to deletion list, and gather all resources that should be deleted.
        self.to_be_deleted_branch.append(branch)

    def _prune_remote_branch(self):
        repo = git.Repo(self.project_root)
        for branch in repo.remote("origin").stale_refs:
            if isinstance(branch, git.refs.remote.RemoteReference):
                type(branch).delete(repo, branch)

    def get_hanging_resources(self):
        self._collect_s3_resources()
        self._collect_airflow_resources()
        self._collect_ecr_resources()
        return self.hanging_branches

    def _collect_s3_resources(self):
        response = self._get_s3_bucket_response(
            self.s3_bucket, prefix=self.s3_prefix, delimiter=self.delimiter
        )
        for content in response.get("CommonPrefixes", []):
            branch_on_s3 = content.get("Prefix").replace("/", "")
            if branch_on_s3 not in self.live_branches:
                self.hanging_s3_resources[branch_on_s3].append(branch_on_s3)
                self.hanging_branches.add(branch_on_s3)

    def _collect_airflow_resources(self):
        airflow_response = self._get_s3_bucket_response(
            self.airflow_bucket,
            prefix=self.airflow_prefix,
            delimiter=self.delimiter,
        )
        for content in airflow_response.get("Contents", []):
            clean_content = content["Key"].lstrip(self.airflow_prefix)
            if clean_content != "":
                branch_name_re = re.findall(r"[\w-]+(?=\.py$)", clean_content)
                repo_name_re = re.findall(r"^[\w-]+\.[\w-]+", clean_content)
                if len(branch_name_re) != 1 or len(repo_name_re) != 1:
                    raise Exception(
                        f"{content['Key']} could not find the appropriate s3 path key."
                    )
                elif (
                    repo_name_re[0] == self.s3_bucket
                    and branch_name_re[0] not in self.live_branches
                ):
                    self.hanging_airflow_resources[branch_name_re[0]].append(
                        content["Key"]
                    )
                    self.hanging_branches.add(branch_name_re[0])

    def _collect_ecr_resources(self):
        response = self._get_ecr_response()["imageIds"]

        for img in response:
            if "imageTag" in img:
                branch_name_re = re.findall(r"^[\w-]+(?=-\w+)", img["imageTag"])
                if len(branch_name_re) != 1:
                    raise Exception(
                        f"""{img['imageTag']} could not find the appropriate branch name.
                        imageTag has to be in the format of <branch_name>-<commit-hash>"""
                    )
                elif branch_name_re[0] not in self.live_branches:
                    self.hanging_ecr_resources[branch_name_re[0]].append(
                        img["imageDigest"]
                    )
                    self.hanging_branches.add(branch_name_re[0])

    def _print_resources(self, branch):
        print(f"{branch} contains the following resources:")
        for s3_resource in self.hanging_s3_resources[branch]:
            print(f"s3://{self.s3_bucket}/{s3_resource}")
        for airflow_resource in self.hanging_airflow_resources[branch]:
            print(f"s3://{self.airflow_bucket}/{airflow_resource}")
        for ecr_resource in self.hanging_ecr_resources[branch]:
            print(f"ecr image: {ecr_resource}")

    def prompt_for_branch(self, branch) -> bool:
        self._print_resources(branch)
        print(
            """If you have local branch that haven't pushed to remote, please do so.
         If you don't know this branch, ask others before deleting all the resources associated with it"""
        )
        answer = input(f"Do you want to delete this branch {branch}?[Y/N]")
        if answer.lower() in ["y", "yes"]:
            return True
        else:
            return False

    def _delete_s3_object(self, bucket, prefix):
        ct = 0
        for obj in self.s3_resource.Bucket(bucket).objects.filter(Prefix=prefix):
            ct += 1
            self.s3_resource.Object(bucket, obj.key).delete()
        logging.info(f"Successfully deleted {ct} files")

    def _get_s3_bucket_response(self, bucket, prefix, delimiter) -> Dict[str, Any]:
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=bucket, Prefix=prefix, Delimiter=delimiter
            )
        except self.s3_client.exceptions.NoSuchBucket:
            logging.info(f"Bucket {bucket} doesn't exist")
        return response

    def _delete_ecr_image(self, image_id):
        self.ecr_client.batch_delete_image(
            repositoryName=self.s3_bucket, imageIds=image_id
        )
        logging.info(f"Successfully deleted image {image_id} in {self.s3_bucket}")

    def _get_ecr_response(self) -> Dict[str, Any]:
        response = self.ecr_client.list_images(repositoryName=self.s3_bucket)
        return response

    def delete_staged_branches(self):
        if not self.dry_run:
            for branch in self.to_be_deleted_branch:
                if branch in self.hanging_s3_resources:
                    for key in self.hanging_s3_resources[branch]:
                        self._delete_s3_object(self.s3_bucket, key)

                if branch in self.hanging_airflow_resources:
                    for key in self.hanging_airflow_resources[branch]:
                        self._delete_s3_object(self.airflow_bucket, key)

                if branch in self.hanging_ecr_resources:
                    for key in self.hanging_ecr_resources[branch]:
                        self._delete_ecr_image(key)
        else:
            print(
                """
                ----------------------------------------------------------------------
                |                                                                    |
                |The following branches will be deleted if --dry-run flag is removed |
                |                                                                    |
                ----------------------------------------------------------------------
                """
            )
            for branch in self.to_be_deleted_branch:
                self._print_resources(branch)


def clean(project_root, dry_run=False, remote=True) -> None:
    # Ensure build directory is a directory.
    build_dir = clean_build(project_root)

    # Set CDK_OUTDIR before projects are loaded.
    cdk_outdir = Path(build_dir, "cdk")
    os.environ["CDK_OUTDIR"] = str(cdk_outdir)

    # Load the project.
    project = load_project(project_root)
    os.environ["AWS_PROFILE"] = project.project_resources.env.account_name

    # Bucket name
    airflow_config = project.project_resources.env.airflow
    airflow_bucket, airflow_s3_path = get_dags_s3_path(airflow_config)
    s3_bucket = project.project_resources.bucket_name()

    # Get live branches names
    branch_manager = BranchManager(
        project_root=project_root,
        s3_bucket=s3_bucket,
        s3_prefix="",
        airflow_bucket=airflow_bucket,
        airflow_prefix=f"{airflow_s3_path}/",
        delimiter="/",
        dry_run=dry_run,
    )

    for branch in branch_manager.get_hanging_resources():
        if branch_manager.prompt_for_branch(branch):
            branch_manager.stage_for_deletion(branch)

    branch_manager.delete_staged_branches()


class ResourcesTypeNotFoundError(Exception):
    """ Exceptions for branch manager. """

    pass
