# 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 enum import Enum
from functools import cache
from typing import TYPE_CHECKING, Literal, NamedTuple

import boto3
from cachetools import cachedmethod

if TYPE_CHECKING:
    from mypy_boto3_apigateway import APIGatewayClient
    from mypy_boto3_cloudformation import CloudFormationClient
    from mypy_boto3_cloudformation.service_resource import CloudFormationServiceResource
    from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource
    from mypy_boto3_events import EventBridgeClient
    from mypy_boto3_iam import IAMClient
    from mypy_boto3_lambda import LambdaClient
    from mypy_boto3_logs import CloudWatchLogsClient
    from mypy_boto3_resourcegroupstaggingapi import ResourceGroupsTaggingAPIClient
    from mypy_boto3_s3 import S3Client
    from mypy_boto3_s3.service_resource import S3ServiceResource
    from mypy_boto3_secretsmanager import SecretsManagerClient
    from mypy_boto3_sqs import SQSClient
    from mypy_boto3_sqs.service_resource import SQSServiceResource
    from mypy_boto3_stepfunctions import SFNClient
    from mypy_boto3_sts import STSClient

DEV_CORE_ACCOUNT_ID = '618300337335'
DEV_REACTOR_ACCOUNT_ID = '998146006915'
PROD_CORE_ACCOUNT_ID = '931830253929'
PROD_REACTOR_ACCOUNT_ID = '061190967865'

DEFAULT_NAMESPACE_DEV = 'alfa'
DEFAULT_NAMESPACE_PROD = 'live'


class _AwsService:
    def __init__(self):
        self._cache = {}

    @cachedmethod(lambda s: s._cache)
    def _service(self, service_type: Literal['client', 'resource'], service_name: str):
        return getattr(boto3, service_type)(service_name)

    @property
    def apigateway_client(self) -> 'APIGatewayClient':
        return self._service('client', 'apigateway')

    @property
    def cloudformation_client(self) -> 'CloudFormationClient':
        return self._service('client', 'cloudformation')

    @property
    def cloudformation_resource(self) -> 'CloudFormationServiceResource':
        return self._service('resource', 'cloudformation')

    @property
    def dynamodb_resource(self) -> 'DynamoDBServiceResource':
        return self._service('resource', 'dynamodb')

    @property
    def events_client(self) -> 'EventBridgeClient':
        return self._service('client', 'events')

    @property
    def iam_client(self) -> 'IAMClient':
        return self._service('client', 'iam')

    @property
    def lambda_client(self) -> 'LambdaClient':
        return self._service('client', 'lambda')

    @property
    def logs_client(self) -> 'CloudWatchLogsClient':
        return self._service('client', 'logs')

    @property
    def s3_client(self) -> 'S3Client':
        return self._service('client', 's3')

    @property
    def s3_resource(self) -> 'S3ServiceResource':
        return self._service('resource', 's3')

    @property
    def secretsmanager_client(self) -> 'SecretsManagerClient':
        return self._service('client', 'secretsmanager')

    @property
    def sqs_client(self) -> 'SQSClient':
        return self._service('client', 'sqs')

    @property
    def sqs_resource(self) -> 'SQSServiceResource':
        return self._service('resource', 'sqs')

    @property
    def stepfunctions_client(self) -> 'SFNClient':
        return self._service('client', 'stepfunctions')

    @property
    def sts_client(self) -> 'STSClient':
        return self._service('client', 'sts')

    @property
    def tagging_api_client(self) -> 'ResourceGroupsTaggingAPIClient':
        return self._service('client', 'resourcegroupstaggingapi')

    def cache_clear(self):
        self._cache.clear()


aws_service = _AwsService()


class Environment(str, Enum):
    dev = 'dev'
    prod = 'prod'


class Account(str, Enum):
    core = 'core'
    reactor = 'reactor'


def default_namespace(environment: Environment) -> str:
    return {Environment.dev: DEFAULT_NAMESPACE_DEV, Environment.prod: DEFAULT_NAMESPACE_PROD}[environment]


def account_id(environment: Environment, account: Account) -> str:
    return {
        (Environment.dev, Account.core): DEV_CORE_ACCOUNT_ID,
        (Environment.dev, Account.reactor): DEV_REACTOR_ACCOUNT_ID,
        (Environment.prod, Account.core): PROD_CORE_ACCOUNT_ID,
        (Environment.prod, Account.reactor): PROD_REACTOR_ACCOUNT_ID,
    }[(environment, account)]


def reactor_account_id(environment: Environment) -> str:
    return account_id(environment, Account.reactor)


def core_account_id(environment: Environment) -> str:
    return account_id(environment, Account.core)


def environment(namespace: str) -> Environment:
    return Environment.prod if namespace == DEFAULT_NAMESPACE_PROD else Environment.dev


class AccountInfo(NamedTuple):
    environment: Environment
    account: Account

    @property
    def account_id(self) -> str:
        return account_id(self.environment, self.account)

    @property
    def default_namespace(self) -> str:
        return default_namespace(self.environment)

    @property
    def reactor_account_id(self) -> str:
        return reactor_account_id(self.environment)

    @property
    def core_account_id(self) -> str:
        return core_account_id(self.environment)

    @classmethod
    def from_namespace(cls, namespace: str, account: Account) -> 'AccountInfo':
        return cls(environment(namespace), account)


@cache
def resolve_account_info() -> AccountInfo:
    sts_client = boto3.client('sts')
    account_id = sts_client.get_caller_identity()['Account']
    return {
        DEV_CORE_ACCOUNT_ID: AccountInfo(Environment.dev, Account.core),
        DEV_REACTOR_ACCOUNT_ID: AccountInfo(Environment.dev, Account.reactor),
        PROD_CORE_ACCOUNT_ID: AccountInfo(Environment.prod, Account.core),
        PROD_REACTOR_ACCOUNT_ID: AccountInfo(Environment.prod, Account.reactor),
    }[account_id]
