# Copyright (c) 2021 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
import re
from datetime import datetime, timedelta
from decimal import Decimal
from enum import Enum
from typing import Any, Iterable, Mapping

import boto3
import boto3.dynamodb.types
import pyfaaster.aws.handlers_decorators_v2 as faaster
import pyfaaster.aws.tools as faaster_aws_tools

from czc.api.v2.constants import CLOUDZERO_TEST_KEY
from czc.constants import DYNAMODB_TTL_OVERRIDE, FEATURE_NAME, PROJECT_NAME

# Lazy-eval the dynamo client so we can use the serializer
DEFAULT_TTL_KEY = 'ttl'
dynamodb = boto3.resource('dynamodb')
deserializer = boto3.dynamodb.types.TypeDeserializer()
logger = faaster_aws_tools.setup_logging(PROJECT_NAME)


class DynamoDbTable:
    get_item: Any
    put_item: Any
    delete_item: Any
    query: Any
    scan: Any


@faaster.namespace_aware
def get_table_name(NAMESPACE='alfa'):
    return f'cz-{NAMESPACE}-{FEATURE_NAME}-data'


def get_expiry_timestamp(seconds_from_now=60) -> int:
    now = datetime.now()
    period = timedelta(seconds=seconds_from_now)
    return int((now + period).timestamp())


def table(name) -> DynamoDbTable:
    logger.debug(f'Connecting to {name} table using default endpoint')
    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)
    else:
        logger.debug(f'Connecting to {name} table using default endpoint')
        return dynamodb.Table(name)


def type_for_dynamo(v):
    if isinstance(v, str) and not v:
        return None
    if isinstance(v, datetime):
        return v.isoformat()
    if isinstance(v, float):
        return Decimal(str(v)).quantize(Decimal('1.0000000'))
    if isinstance(v, Decimal):
        return v.quantize(Decimal('1.0000000'))
    if isinstance(v, Enum):
        return str(v)
    if isinstance(v, dict):
        return {type_for_dynamo(k): type_for_dynamo(vv)
                for (k, vv) in v.items()}
    if isinstance(v, list):
        return [type_for_dynamo(vv) for vv in v]

    return v


def index_from_table(table, index_name):
    [index_name] = [x['IndexName']
                    for x in table.global_secondary_indexes if index_name in x['IndexName']]
    return index_name


COMPOSITE_CHAR = '#'


@faaster.environ_aware(optional=[CLOUDZERO_TEST_KEY])
def prefix(prefix: str, value: str = '', *, composite_char=COMPOSITE_CHAR, **kwargs) -> str:
    cloudzero_test_key = kwargs.get(CLOUDZERO_TEST_KEY)
    _prefix = f'{cloudzero_test_key}{COMPOSITE_CHAR}{prefix}' if cloudzero_test_key else prefix
    return value if value.startswith(compose_key(_prefix, '', composite_char=composite_char)) or value is None else compose_key(_prefix, value, composite_char=composite_char)


@faaster.environ_aware(optional=[CLOUDZERO_TEST_KEY])
def unfix(unfix: str, value: str = '', *, composite_char=COMPOSITE_CHAR, **kwargs) -> str:
    cloudzero_test_key = kwargs.get(CLOUDZERO_TEST_KEY)
    _unfix = f'{cloudzero_test_key}{COMPOSITE_CHAR}{unfix}' if cloudzero_test_key else unfix
    return value.replace(compose_key(_unfix, '', composite_char=composite_char), '') if value is not None and value.startswith(compose_key(_unfix, '', composite_char=composite_char)) else value


def compose_key(*args, composite_char=COMPOSITE_CHAR) -> str:
    return composite_char.join(args)


def decompose_key(key: str, *, composite_char=COMPOSITE_CHAR) -> Iterable:
    return key.split(composite_char)


pattern = re.compile(r'[\W_]+', re.UNICODE)


def dynamo_safe(value):
    if isinstance(value, str) and not value:
        return None
    if isinstance(value, float):
        return Decimal(str(value)).quantize(Decimal('1.0000000'))
    if isinstance(value, Decimal):
        return value.quantize(Decimal('1.0000000'))
    if isinstance(value, Enum):
        return str(value)
    if isinstance(value, dict):
        return {k: dynamo_safe(vv)
                for k, vv in value.items()}
    if isinstance(value, list):
        return [dynamo_safe(vv) for vv in value]

    return value


@faaster.environ_aware(optional=[CLOUDZERO_TEST_KEY, DYNAMODB_TTL_OVERRIDE])
def update_item_from_dict(table, key: Mapping, dictionary: Mapping, **kwargs) -> Mapping:
    """Update the item identified by `key` in the DynamoDB `table` by adding
    all of the attributes in the `dictionary`.
    """
    cloudzero_test_key = kwargs.get(CLOUDZERO_TEST_KEY)
    ttl_key = kwargs.get(DYNAMODB_TTL_OVERRIDE) or DEFAULT_TTL_KEY
    dictionary_with_ttl = {**dictionary, ttl_key: get_expiry_timestamp()} if cloudzero_test_key else dictionary
    updates_string = ', '.join([f'#{pattern.sub("", k)} = :{pattern.sub("", k)}' for k in dictionary_with_ttl.keys()])
    update_expression = f'SET {updates_string}'
    attribute_names = {f'#{pattern.sub("", k)}': k for k in dictionary_with_ttl.keys()}
    attribute_values = {f':{pattern.sub("", k)}': dynamo_safe(v)
                        for k, v in dictionary_with_ttl.items()}

    try:
        item = table.update_item(
            Key=key,
            UpdateExpression=update_expression,
            ExpressionAttributeNames=attribute_names,
            ExpressionAttributeValues=attribute_values,
            ReturnValues='ALL_NEW',
        )
        return item.get('Attributes', {}) if item else {}
    except Exception as err:
        logger.error(f'Could not commit {key}')
        raise err


@faaster.environ_aware(optional=[CLOUDZERO_TEST_KEY, DYNAMODB_TTL_OVERRIDE])
def put_item(table, Item=None, ConditionExpression=None, **kwargs):
    cloudzero_test_key = kwargs.get(CLOUDZERO_TEST_KEY)
    ttl_key = kwargs.get(DYNAMODB_TTL_OVERRIDE) or DEFAULT_TTL_KEY
    item = Item or {}
    item_with_ttl = {**item, ttl_key: get_expiry_timestamp()} if cloudzero_test_key else item
    condition_args = dict(ConditionExpression=ConditionExpression) if ConditionExpression is not None else {}
    return table.put_item(Item=item_with_ttl, **condition_args)


def _paginate(paginator_fn, pagination_kwargs):
    try:
        response = paginator_fn(**pagination_kwargs)
        items = response.get('Items', [])
        while 'LastEvaluatedKey' in response:
            response = paginator_fn(**{
                **pagination_kwargs,
                'ExclusiveStartKey': response['LastEvaluatedKey'],
            })
            items += response.get('Items', [])
        return items
    except Exception as error:
        logger.exception(f'Error while performing {paginator_fn} with {pagination_kwargs}')
        raise error


def scan(table, scan_kwargs):
    return _paginate(table.scan, scan_kwargs)


def query(table, query_kwargs):
    return _paginate(table.query, query_kwargs)
