# Copyright (c) 2023 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from functools import lru_cache

import boto3
from aws_lambda_powertools import Logger
from pyfaaster.aws import handlers_decorators_v2 as faaster

from czc import aws, json

from .constants import DEFAULT_NAMESPACE, PROJECT_NAME

logger = Logger(service=PROJECT_NAME)
CORE_ACCOUNT_IDS = {'931830253929', '618300337335'}


@lru_cache(maxsize=1)
def _get_sm_client():
    return boto3.client('secretsmanager')


@lru_cache(maxsize=1)
def _get_sts_client():
    return boto3.client('sts')


def get_account_id() -> str:
    return _get_sts_client().get_caller_identity()['Account']


@faaster.namespace_aware
def get_core_credentials(*, NAMESPACE: str = DEFAULT_NAMESPACE):
    sm = _get_sm_client()
    secret = sm.get_secret_value(SecretId=f'cz-{NAMESPACE}-feature-setup-secrets')
    secret_dict = json.loads(secret['SecretString'])
    logger.debug(f'retrieved core credentials {secret_dict}')
    return secret_dict['core_role_arn'], secret_dict['core_external_id']


@lru_cache(maxsize=1)
def get_session():
    account_id = get_account_id()
    if account_id not in CORE_ACCOUNT_IDS:
        core_role_arn, core_external_id = get_core_credentials()
        return aws.assume_role(role_arn=core_role_arn, external_id=core_external_id)
    return boto3.session.Session()
