# 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

from typing import Any, Callable, List, NamedTuple, Optional, overload

import pyfaaster.aws.tools as faaster_aws_tools
from boto3.dynamodb.conditions import Attr, Key
from botocore.exceptions import ClientError
from retrying import retry
from toolz import get_in

from czc import dynamodb as dyn
from czc import json
from czc.constants import FEATURE_NAME, PROJECT_NAME
from czc.dynamodb import DynamoDbTable
from deprecated import deprecated

logger = faaster_aws_tools.setup_logging(PROJECT_NAME)


#
# This module provides a simple interface for a common use case: storing and retrieving simple object data in DynamoDB.
# The "object data" can be any serializable dictionary. The data is stored by organization_id (PK) and a user defined
# compound key (SK). Object data can also be updated atomically.
#
# Caveats:
# - Atomic update mechanism is *SLOW* when a collision occurs and is only suitable for
#   low volume updates for which collisions are rare.
# - Object data is stored in a single DynamoDB attribute.  This is not suitable for complex use cases which require
#   filtering on multiple attributes.  Only use for simple per-org key (or key prefix) access.
#


class FeatureDataConnection(NamedTuple):
    table: DynamoDbTable
    primary_key: str


def _organization_key(organization_id: str) -> str:
    return dyn.prefix('organization', organization_id)


def _feature_data_table_name(namespace: str) -> str:
    return f'cz-{namespace}-{FEATURE_NAME}-data'


@overload
def for_organization(organization_id: str) -> FeatureDataConnection:
    ...


@overload
@deprecated(reason='Use for_organization(organization_id: str) to get NAMESPACE from the environment variable.')
def for_organization(namespace: str, organization_id: str) -> FeatureDataConnection:
    ...


def for_organization(*args, **kwargs) -> FeatureDataConnection:
    namespace_arg, organization_id_arg = (None, args[0]) if len(args) == 1 else args
    namespace = kwargs.get('namespace', namespace_arg)
    organization_id = kwargs.get('organization_id', organization_id_arg)
    table_name = _feature_data_table_name(namespace) if namespace is not None else dyn.get_table_name()
    return FeatureDataConnection(
        table=dyn.table(table_name),
        primary_key=_organization_key(organization_id)
    )


def _update_data(table: DynamoDbTable, pk: str, sk: str, current_data: dict, new_data: dict) -> None:
    data_attr = Attr('data')
    condition = (
        data_attr.eq(current_data)) if current_data is not None else (
            data_attr.not_exists() | data_attr.attribute_type('NULL'))

    dyn.put_item(table,
                 ConditionExpression=condition,
                 Item=dict(
                     PK=pk,
                     SK=sk,
                     data=new_data
                 ))


def _get_data(table: DynamoDbTable, pk: str, sk: str) -> Any:
    item = table.get_item(
        Key={
            'PK': pk,
            'SK': sk
        }
    )
    return get_in(['Item', 'data'], item)


def _load_data(object_data) -> dict:
    if object_data is None:
        return {}
    if not isinstance(object_data, str):
        return object_data

    result = json.loads(object_data)
    return result or {}


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


@retry(retry_on_exception=_is_conditional_check_exception,
       stop_max_attempt_number=3,
       wait_random_min=1000,
       wait_random_max=2000)
def update_object_atomic(conn: FeatureDataConnection, *key: str,
                         update_handler: Callable[[dict], Optional[dict]], serialize=False) -> Optional[dict]:
    table = conn.table
    pk = conn.primary_key
    sk = dyn.compose_key(*key)

    data = _get_data(table, pk, sk)
    object_data = _load_data(data)

    new_object_data = update_handler(object_data)
    if new_object_data is not None:
        _update_data(table, pk, sk, data,
                     json.dumps(new_object_data) if serialize else dyn.type_for_dynamo(new_object_data))
    return new_object_data


def get_object(conn: FeatureDataConnection, *key: str) -> dict:
    sk = dyn.compose_key(*key)
    data = _get_data(conn.table, conn.primary_key, sk)
    return _load_data(data)


def get_all_objects(conn: FeatureDataConnection, *key: str) -> List[dict]:
    sk = dyn.compose_key(key)
    cond = Key('PK').eq(conn.primary_key) & Key('SK').begins_with(sk)
    items = dyn.query(conn.table, dict(KeyConditionExpression=cond)).get('Items', [])
    return [_load_data(item['data']) for item in items]


def set_object(conn: FeatureDataConnection, *key: str, data: dict, serialize=False):
    dyn.put_item(conn.table,
                 Item=dict(
                     PK=conn.primary_key,
                     SK=dyn.compose_key(*key),
                     data=json.dumps(data) if serialize else dyn.type_for_dynamo(data)
                 )
                 )


def delete_object(conn: FeatureDataConnection, *key: str):
    conn.table.delete_item(Key=dict(
        PK=conn.primary_key,
        SK=dyn.compose_key(*key)
    ))
