import os
from actions_logging.app_logging import logger
from github.env import (exit_on_error_and_write_summary, get_required_env_var,
                        parse_bool_env_var, write_github_env)
from aws.constants import DEFAULT_ROLE_SESSION_NAME, ENVS, CICD_CONFIG, CICD_NAMES, AUTH0_CONFIG, AUTH0_TENANTS, MARKETING_CONFIG, MARKETING_ENVS


class ConfigError(Exception):
    pass

def set_assume_role_config(env_name: str, aws_region: str = None, explicit_production_role: bool = False,
                           admin: bool = False):
    role_session_name = DEFAULT_ROLE_SESSION_NAME
    logger.info(
        f"configuring assume role with env_name: {env_name}, aws_region: {aws_region}, explicit_production_role: {explicit_production_role}, admin: {admin}")

    if env_name not in ENVS:
        if env_name in CICD_NAMES:
            env_config = CICD_CONFIG
        elif env_name in AUTH0_TENANTS:
            env_config = AUTH0_CONFIG
        elif env_name in MARKETING_ENVS:
            env_config = MARKETING_CONFIG
        else:
            exit_on_error_and_write_summary(f"Unsupported ENV_NAME for assume role: {env_name}, halting.")
    else:
        env_config = ENVS[env_name]
    logger.info(f"Assume role configuration for {env_name}: {env_config}")

    try:
        aws_account = env_config['aws_account']
    except KeyError:
        raise ConfigError(f"'aws_account' not defined for {env_name}")

    if aws_region:
        logger.info(f"aws_region '{aws_region}' supplied to use during assume")
        region = aws_region
    else:
        try:
            region = env_config['aws_region']
        except KeyError:
            raise ConfigError(f"'aws_region' not defined for {env_name}")

    if admin:
        logger.warning(f"assuming admin role for {env_name}")
        try:
            admin_config = env_config['admin_role']
        except KeyError:
            raise ConfigError(f"'admin_role' configuration not defined for {env_name} with ENV_NAME: {env_name}")
        try:
            role_name = admin_config['role_name']
        except KeyError:
            raise ConfigError(
                f"'role_name' configuration not defined in 'admin_role' for {env_name} with ENV_NAME: {env_name}")
        try:
            role_session_name = admin_config['role_session_name']
        except KeyError:
            raise ConfigError(
                f"'role_session_name' configuration not defined  in 'admin_role' for {env_name} with ENV_NAME: {env_name}")
    else:
        try:
            role_name = env_config['role_name']
        except KeyError:
            raise ConfigError(f"'role_name' configuration not defined for {env_name}")

    role = f"arn:aws:iam::{aws_account}:role/{role_name}"

    logger.warning(f"Will assume role {role} in {region}")

    write_github_env(region, 'ASSUME_ROLE_AWS_REGION')
    write_github_env(role_session_name, 'ASSUME_ROLE_SESSION_NAME')
    write_github_env(role, 'ROLE_TO_ASSUME')


def main():
    env_name = get_required_env_var('ENV_NAME')
    assume_role_aws_region = os.getenv('ASSUME_ROLE_AWS_REGION')
    aws_secret_region = os.getenv('AWS_SECRET_REGION')

    aws_region = None
    if assume_role_aws_region:
        aws_region = assume_role_aws_region
    # Backward compatibility with assume role v1
    elif aws_secret_region:
        aws_region = aws_secret_region

    explicit_production_role = parse_bool_env_var(os.getenv('EXPLICIT_PRODUCTION_ROLE', False))
    admin = parse_bool_env_var(os.getenv('ADMIN_ROLE', False))
    try:
        set_assume_role_config(env_name, aws_region, explicit_production_role, admin)
    except Exception as ce:
        exit_on_error_and_write_summary(f"Couldn't setup configuration for assume role action: {ce}")


if __name__ == "__main__":
    main()