# 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 List, Any, Iterable, Callable, Optional, NamedTuple

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

logger = faaster_aws_tools.setup_logging(PROJECT_NAME)


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


def _organization_key(organization_id: str) -> str:
    return _compound_key(['organization', organization_id])


def _compound_key(parts: Iterable[str]) -> str:
    return '#'.join(parts)


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


def _feature_data_table(namespace: str) -> DynamoDbTable:
    return dyn.table(_feature_data_table_name(namespace))


def for_organization(namespace: str, organization_id: str) -> FeatureDataConnection:
    return FeatureDataConnection(
        table=_feature_data_table(namespace),
        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'))

    table.put_item(
        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 _safe_load_data(pk: str, sk: str, object_data) -> dict:
    if object_data is None:
        return {}
    if not isinstance(object_data, str):
        return object_data
    try:
        result = json.loads(object_data)
        return result or {}
    except Exception:
        logger.warning(f'Bad json data for PK={pk} SK={sk}: {object_data}')
        return {}


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 = _compound_key(key)

    data = _get_data(table, pk, sk)
    object_data = _safe_load_data(conn.primary_key, sk, 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 = _compound_key(key)
    data = _get_data(conn.table, conn.primary_key, sk)
    return _safe_load_data(conn.primary_key, sk, data)


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


def set_object(conn: FeatureDataConnection, *key: str, data: dict, serialize=False):
    conn.table.put_item(
        Item=dict(
            PK=conn.primary_key,
            SK=_compound_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=_compound_key(key)
    ))
