# 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 functools import wraps
from retrying import retry
from typing import Any, Iterable, Mapping

import boto3
import boto3.dynamodb.types
from botocore.exceptions import ClientError
import pyfaaster.aws.handlers_decorators_v2 as faaster

from czc.api.v2.constants import CLOUDZERO_TEST_KEY
from czc.constants import DYNAMODB_TTL_DEFAULT_VALUE, DYNAMODB_TTL_OVERRIDE, DYNAMODB_TTL_OVERRIDE_VALUE, FEATURE_NAME
from czc.logging import get_logger


# 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 = get_logger()


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=None) -> int:
    now = datetime.now()
    period = timedelta(seconds=int(seconds_from_now or DYNAMODB_TTL_DEFAULT_VALUE))
    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)  # type: ignore
    else:
        logger.debug(f'Connecting to {name} table using default endpoint')
        return dynamodb.Table(name)  # type: ignore


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, use_enum_values=False):
    """
    `use_enum_values` -> Use enum value instead of string (i.e. value.value vs. str(value)).
    If set to false, then a string representation is used which will look like `EnumType.enum_value`.
    If set to true, then the enum value is used `enum_value`.

    This field is defaulted to false for backwards compatibility.
    """
    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 value.value if use_enum_values else str(value)
    if isinstance(value, dict):
        return {k: dynamo_safe(vv, use_enum_values=use_enum_values)
                for k, vv in value.items()}
    if isinstance(value, list):
        return [dynamo_safe(vv, use_enum_values=use_enum_values) for vv in value]

    return value


@faaster.environ_aware(optional=[CLOUDZERO_TEST_KEY, DYNAMODB_TTL_OVERRIDE, DYNAMODB_TTL_OVERRIDE_VALUE])
def update_item_from_dict(table, key: Mapping, dictionary: Mapping, ConditionExpression=None, use_enum_values=False, **kwargs) -> Mapping:
    """Update the item identified by `key` in the DynamoDB `table` by adding
    all of the attributes in the `dictionary`.

    `use_enum_values` -> Use enum value instead of string (i.e. value.value vs. str(value)).
    If set to false, then a string representation is used which will look like `EnumType.enum_value`.
    If set to true, then the enum value is used `enum_value`.

    This field is defaulted to false for backwards compatibility.
    """
    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(kwargs.get(DYNAMODB_TTL_OVERRIDE_VALUE))} 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, use_enum_values=use_enum_values)
                        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',
            **({'ConditionExpression': ConditionExpression} if ConditionExpression else {})
        )
        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, DYNAMODB_TTL_OVERRIDE_VALUE])
def put_item(table, Item=None, ConditionExpression=None, ReturnValues=None, **kwargs) -> dict:
    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(kwargs.get(DYNAMODB_TTL_OVERRIDE_VALUE))} if cloudzero_test_key else item
    condition_args = dict(ConditionExpression=ConditionExpression) if ConditionExpression is not None else {}
    return_values_args = dict(ReturnValues=ReturnValues) if ReturnValues is not None else {}
    return table.put_item(Item=item_with_ttl, **condition_args, **return_values_args)


def _is_conditional_check_exception(e: Exception) -> bool:
    return isinstance(e, ClientError) and e.response['Error']['Code'] == 'ConditionalCheckFailedException'


def is_max_item_size_exception(e: ClientError) -> bool:
    error = e.response['Error']
    return (error['Code'] == 'ValidationException' and
            ('exceeded the maximum allowed size' in error.get('Message', '') or error.get('Message', '').endswith('exceeds max size'))
            )


def atomic(stop_max_attempt_number=3, wait_random_min=1000, wait_random_max=2000):
    def atomic_handler(handler):
        @retry(
            retry_on_exception=_is_conditional_check_exception,
            stop_max_attempt_number=stop_max_attempt_number,
            wait_random_min=wait_random_min,
            wait_random_max=wait_random_max)
        @wraps(handler)
        def wrapper(*args, **kwargs):
            return handler(*args, **kwargs)
        return wrapper
    return atomic_handler


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)
