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

import boto3
from cachetools import LRUCache, cachedmethod

from czc.logging import get_logger

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'

logger = get_logger()


class _AwsService:
    def __init__(self):
        self._cache = LRUCache(maxsize=1024)

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

    def client(self, service_name: str, *args, **kwargs):
        return self._service('client', service_name, *args, **kwargs)

    def resource(self, service_name: str, *args, **kwargs):
        return self._service('resource', service_name, *args, **kwargs)

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

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

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

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

    @property
    def ec2_client(self):
        return self._service('client', 'ec2')

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

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

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

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

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

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

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

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

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

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

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

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

    @cachedmethod(lambda s: s._cache, key=lambda s, n: ('dynamodb_table', n, os.environ.get('DYNAMODB_ENDPOINTURL')))
    def dynamodb_table(self, name: str):
        dynamodb_endpointurl = os.environ.get('DYNAMODB_ENDPOINTURL')
        if dynamodb_endpointurl:
            logger.debug(f'Connecting to {name} table at {dynamodb_endpointurl}')
            return boto3.resource('dynamodb', endpoint_url=dynamodb_endpointurl).Table(name)
        logger.debug(f'Connecting to {name} table using default endpoint')
        return self.dynamodb_resource.Table(name)

    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]


def cache_clear():
    boto3.setup_default_session()
    resolve_account_info.cache_clear()
    aws_service.cache_clear()
