import json
import logging
import os
import re
import subprocess
import shutil
import tempfile
import zipfile
import boto3

logger = logging.getLogger()
logger.setLevel(logging.INFO)

# these are coming from the kubectl layer
os.environ['PATH'] = '/opt/helm:/opt/awscli:' + os.environ['PATH']

outdir = os.environ.get('TEST_OUTDIR', '/tmp')
kubeconfig = os.path.join(outdir, 'kubeconfig')

def get_chart_asset_from_url(chart_asset_url):
    chart_zip = os.path.join(outdir, 'chart.zip')
    shutil.rmtree(chart_zip, ignore_errors=True)
    subprocess.check_call(['aws', 's3', 'cp', chart_asset_url, chart_zip])
    chart_dir = os.path.join(outdir, 'chart')
    shutil.rmtree(chart_dir, ignore_errors=True)
    os.mkdir(chart_dir)
    with zipfile.ZipFile(chart_zip, 'r') as zip_ref:
        zip_ref.extractall(chart_dir)
    return chart_dir

def is_ecr_public_available(region):
    s = boto3.Session()
    return s.get_partition_for_region(region) == 'aws'

def helm_handler(event, context):
    logger.info(json.dumps(dict(event, ResponseURL='...')))

    request_type = event['RequestType']
    props = event['ResourceProperties']

    # resource properties
    cluster_name     = props['ClusterName']
    role_arn         = props['RoleArn']
    release          = props['Release']
    chart            = props.get('Chart', None)
    chart_asset_url  = props.get('ChartAssetURL', None)
    version          = props.get('Version', None)
    wait             = props.get('Wait', False)
    timeout          = props.get('Timeout', None)
    namespace        = props.get('Namespace', None)
    create_namespace = props.get('CreateNamespace', None)
    repository       = props.get('Repository', None)
    values_text      = props.get('Values', None)
    skip_crds        = props.get('SkipCrds', False)

    # "log in" to the cluster
    subprocess.check_call([ 'aws', 'eks', 'update-kubeconfig',
        '--role-arn', role_arn,
        '--name', cluster_name,
        '--kubeconfig', kubeconfig
    ])

    if os.path.isfile(kubeconfig):
        os.chmod(kubeconfig, 0o600)

    # Write out the values to a file and include them with the install and upgrade
    values_file = None
    if not request_type == "Delete" and not values_text is None:
        values = json.loads(values_text)
        values_file = os.path.join(outdir, 'values.yaml')
        with open(values_file, "w") as f:
            f.write(json.dumps(values, indent=2))

    if request_type == 'Create' or request_type == 'Update':
        # Ensure chart or chart_asset_url are set
        if chart == None and chart_asset_url == None:
            raise RuntimeError(f'chart or chartAsset must be specified')

        if chart_asset_url != None:
            assert(chart==None)
            assert(repository==None)
            assert(version==None)
            if not chart_asset_url.startswith('s3://'):
              raise RuntimeError(f'ChartAssetURL must point to as s3 location but is {chart_asset_url}')
            # future work: support versions from s3 assets
            chart = get_chart_asset_from_url(chart_asset_url)

        if repository is not None and repository.startswith('oci://'):
            tmpdir = tempfile.TemporaryDirectory()
            chart_dir = get_chart_from_oci(tmpdir.name, repository, version)
            chart = chart_dir

        helm('upgrade', release, chart, repository, values_file, namespace, version, wait, timeout, create_namespace)
    elif request_type == "Delete":
        try:
            helm('uninstall', release, namespace=namespace, timeout=timeout)
        except Exception as e:
            logger.info("delete error: %s" % e)


def get_oci_cmd(repository, version):
    # Generates OCI command based on pattern. Public ECR vs Private ECR are treated differently.
    private_ecr_pattern = 'oci://(?P<registry>\d+.dkr.ecr.(?P<region>[a-z]+-[a-z]+-\d).amazonaws.com)*'
    public_ecr_pattern = 'oci://(?P<registry>public.ecr.aws)*'

    private_registry = re.match(private_ecr_pattern, repository).groupdict()
    public_registry = re.match(public_ecr_pattern, repository).groupdict()

    if private_registry['registry'] is not None:
        logger.info("Found AWS private repository")
        cmnd = [
            f"aws ecr get-login-password --region {private_registry['region']} | " \
            f"helm registry login --username AWS --password-stdin {private_registry['registry']}; helm pull {repository} --version {version} --untar"
            ]
    elif public_registry['registry'] is not None:
        logger.info("Found AWS public repository, will use default region as deployment")
        region = os.environ.get('AWS_REGION', 'us-east-1')
        
        if is_ecr_public_available(region):
            cmnd = [
                f"aws ecr-public get-login-password --region us-east-1 | " \
                f"helm registry login --username AWS --password-stdin {public_registry['registry']}; helm pull {repository} --version {version} --untar"
                ]
        else:
            # `aws ecr-public get-login-password` and `helm registry login` not required as ecr public is not available in current region
            # see https://helm.sh/docs/helm/helm_registry_login/
            cmnd = [f"helm pull {repository} --version {version} --untar"]   
    else:
        logger.error("OCI repository format not recognized, falling back to helm pull")
        cmnd = [f"helm pull {repository} --version {version} --untar"]

    return cmnd


def get_chart_from_oci(tmpdir, repository = None, version = None):

    cmnd = get_oci_cmd(repository, version)

    maxAttempts = 3
    retry = maxAttempts
    while retry > 0:
        try:
            logger.info(cmnd)
            output = subprocess.check_output(cmnd, stderr=subprocess.STDOUT, cwd=tmpdir, shell=True)
            logger.info(output)

            # effectively returns "$tmpDir/$lastPartOfOCIUrl", because this is how helm pull saves OCI artifact. 
            # Eg. if we have oci://9999999999.dkr.ecr.us-east-1.amazonaws.com/foo/bar/pet-service repository, helm saves artifact under $tmpDir/pet-service
            return os.path.join(tmpdir, repository.rpartition('/')[-1])
        except subprocess.CalledProcessError as exc:
            output = exc.output
            if b'Broken pipe' in output:
                retry = retry - 1
                logger.info("Broken pipe, retries left: %s" % retry)
            else:
                raise Exception(output)
    raise Exception(f'Operation failed after {maxAttempts} attempts: {output}')


def helm(verb, release, chart = None, repo = None, file = None, namespace = None, version = None, wait = False, timeout = None, create_namespace = None, skip_crds = False):
    import subprocess

    cmnd = ['helm', verb, release]
    if not chart is None:
        cmnd.append(chart)
    if verb == 'upgrade':
        cmnd.append('--install')
    if create_namespace:
        cmnd.append('--create-namespace')
    if not repo is None:
        cmnd.extend(['--repo', repo])
    if not file is None:
        cmnd.extend(['--values', file])
    if not version is None:
        cmnd.extend(['--version', version])
    if not namespace is None:
        cmnd.extend(['--namespace', namespace])
    if wait:
        cmnd.append('--wait')
    if skip_crds:
        cmnd.append('--skip-crds')
    if not timeout is None:
        cmnd.extend(['--timeout', timeout])
    cmnd.extend(['--kubeconfig', kubeconfig])

    maxAttempts = 3
    retry = maxAttempts
    while retry > 0:
        try:
            output = subprocess.check_output(cmnd, stderr=subprocess.STDOUT, cwd=outdir)
            logger.info(output)
            return
        except subprocess.CalledProcessError as exc:
            output = exc.output
            if b'Broken pipe' in output:
                retry = retry - 1
                logger.info("Broken pipe, retries left: %s" % retry)
            else:
                raise Exception(output)
    raise Exception(f'Operation failed after {maxAttempts} attempts: {output}')