import boto3
import os
import time
from functools import reduce

from github.env import exit_on_error_and_write_summary, get_required_env_var
from actions_logging.app_logging import logger
from lambdas.lambda_info import get_lambda_function_name

def check_existing_version(lambda_client, function_name, version_description):
    try:
        response = lambda_client.list_versions_by_function(FunctionName=function_name)
        versions = response['Versions']
        versions_matched = [version['Description'] == version_description for version in versions]
        return versions_matched[0] or None
    except Exception as e:
        exit_on_error_and_write_summary(f"Failed to check existing version: {str(e)}")

def wait_until_condition_met(lambda_client, function_name, status_key, success_status, max_retries=20, delay=15):
    for attempt in range(max_retries):
        response = lambda_client.get_function(FunctionName=function_name)
        status = response
        # reduce extracts the nested value from the response using the dot-separated keys in status_key more here:https://docs.python.org/3/library/functools.html
        status = reduce(lambda d, key: d.get(key, None) if d else None, status_key.split("."), response)
        if status == success_status:
            return
        logger.info(f"Attempt {attempt + 1}/{max_retries}: {status_key}='{status}'. Retrying in {delay} seconds...")
        time.sleep(delay)
    exit_on_error_and_write_summary(
        f"Timeout while waiting for {status_key} to reach '{success_status}'."
    )


def publish_version(lambda_client, function_name, version_description):
    try:
        logger.info(f"Publishing new version: {version_description}")
        response = lambda_client.publish_version(
            FunctionName=function_name,
            Description=version_description
        )
        version = response['Version']
        wait_until_condition_met(lambda_client=lambda_client, function_name=f"{function_name}:{version}", status_key="Configuration.State", success_status="Active")
        logger.info(f"Published version: {version_description}")
        return version

    except Exception as e:
        exit_on_error_and_write_summary(f"Failed to publish a new version: {str(e)}")


def update_lambda_alias(lambda_client, function_name, alias_name, version, version_description, use_latest_version):
    version = "$LATEST" if use_latest_version else version
    return lambda_client.update_alias(
            FunctionName=function_name,
            Name=alias_name,
            FunctionVersion=version
        )

def main():
    svc_name = get_required_env_var('SVC_NAME')
    env_name = get_required_env_var("ENV_NAME")
    function_name = get_lambda_function_name(svc_name, env_name)
    region = get_required_env_var("AWS_REGION")
    version_description = os.getenv('DEPLOYED_ARTIFACT_VERSION')
    use_latest_version = os.getenv("USE_LATEST_VERSION") == "true"
    alias_name = "published-version"
    lambda_client = boto3.client('lambda', region_name=region)
    wait_until_condition_met(lambda_client=lambda_client, function_name=function_name,
                             status_key="Configuration.LastUpdateStatus", success_status="Successful")
    version = check_existing_version(lambda_client, function_name, version_description)
    if not version and not use_latest_version:
        logger.debug(f"Version {version_description} does not exist, publishing a new version")
        version = publish_version(lambda_client, function_name, version_description)
    response = update_lambda_alias(lambda_client, function_name, alias_name, version, version_description, use_latest_version)
    logger.info(response)

if __name__ == "__main__":
    main()

