# 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
from decimal import Decimal
from enum import Enum
from typing import Any, Iterable, Mapping

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

from czc.constants import PROJECT_NAME

# Lazy-eval the dynamo client so we can use the serializer
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


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 = '#'


def prefix(prefix: str, value: str = '', *, composite_char=COMPOSITE_CHAR) -> str:
    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)


def unfix(unfix: str, value: str = '', *, composite_char=COMPOSITE_CHAR) -> str:
    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 update_item_from_dict(table, key: Mapping, dictionary: Mapping) -> Mapping:
    """Update the item identified by `key` in the DynamoDB `table` by adding
    all of the attributes in the `dictionary`.
    """
    def fix(v):
        if isinstance(v, str) and not v:
            return None
        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 {k: fix(vv)
                    for k, vv in v.items()}
        if isinstance(v, list):
            return [fix(vv) for vv in v]

        return v

    updates_string = ', '.join([f'#{pattern.sub("", k)} = :{pattern.sub("", k)}' for k in dictionary.keys()])
    update_expression = f'SET {updates_string}'
    attribute_names = {f'#{pattern.sub("", k)}': k for k in dictionary.keys()}
    attribute_values = {f':{pattern.sub("", k)}': fix(v)
                        for k, v in dictionary.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 None
    except Exception as err:
        logger.error(f'Could not commit {key}')
        raise err
