# 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 re
import uuid
from dataclasses import dataclass
from functools import singledispatch
from typing import Callable, Optional, Union, Iterable, Generic, TypeVar, Type

import boto3
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, s3, type_util
from czc.constants import PROJECT_NAME
from czc.dynamodb import DynamoDbTable

logger = faaster_aws_tools.setup_logging(PROJECT_NAME)

#
# This module provides a simple interface for a common use case: storing and retrieving a data object for an
# organization and specified key. The data object can be any serializable data class.
#
# Caveats:
# - This is intended to be used for "configuration" type data for which access is infrequent and performance
#   is not a significant concern.
# - Atomic update mechanism is *SLOW* when a collision occurs and is only suitable for
#   low volume updates for which collisions are rare.
#

S3_OBJECT_REF = 'REF:'
S3_OBJECT_STORE_ROOT = 'feature_data_object_store'
T = TypeVar('T')

s3_client = boto3.client('s3')


@dataclass(frozen=True)
class FeatureDataConnection:
    _table: DynamoDbTable
    _s3_bucket: str


class _FeatureData(Generic[T]):
    def __init__(self, obj_cls: Type[T]):
        self._obj_cls = obj_cls

    def _coerce(self, data: dict) -> T:
        return type_util.coerce(self._obj_cls, data)

    def get(self, conn: FeatureDataConnection, cz_organization_id: str,
            keyspace: str, *key: str) -> Optional[T]:
        obj = get_object(conn, cz_organization_id, keyspace, *key)
        return self._coerce(obj) if obj else None

    def set(self, conn: FeatureDataConnection, cz_organization_id: str,
            keyspace: str, *key: str, data: T) -> None:
        set_object(conn, cz_organization_id, keyspace, *key, data=type_util.serializable(data))

    def update(self, conn: FeatureDataConnection, cz_organization_id: str,
               keyspace: str, *key: str, handler: Callable[[Optional[T]], Optional[T]]) -> Optional[T]:
        def update_handler(data: dict) -> dict:
            return type_util.serializable(handler(self._coerce(data) if data else None))

        result = update_object_atomic(conn, cz_organization_id, keyspace, *key,
                                      update_handler=update_handler)
        return self._coerce(result) if result else None

    def list(self, conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str) -> Iterable[T]:
        for x in list_objects(conn, cz_organization_id, keyspace, *key):
            yield self._coerce(x)

    def list_all(self, conn: FeatureDataConnection, keyspace: str, *key: str) -> Iterable[T]:
        for x in list_all_objects(conn, keyspace, *key):
            yield self._coerce(x)

    def delete(self, conn: FeatureDataConnection, cz_organization_id: str,
               keyspace: str, *key: str) -> None:
        delete_object(conn, cz_organization_id, keyspace, *key)


def for_type(cls: Type[T]) -> _FeatureData[T]:
    return _FeatureData(cls)


def connect() -> FeatureDataConnection:
    table_name = dyn.get_table_name()
    return FeatureDataConnection(
        _table=dyn.table(table_name),
        _s3_bucket=s3.get_bucket_name()
    )


def _prefix_organization(cz_organization_id: str) -> str:
    return dyn.prefix('organization', cz_organization_id)


def _dyn_key(cz_organization_id: str, keyspace: str, *key: str) -> dict:
    return {
        'PK': _prefix_organization(cz_organization_id),
        'SK': dyn.compose_key(keyspace, *key)
    }


def _s3_key(ref: str):
    return f'{S3_OBJECT_STORE_ROOT}/{ref[len(S3_OBJECT_REF):]}'


def _is_ref(item_data: str) -> bool:
    return item_data.startswith(S3_OBJECT_REF)


def _load_data(conn: FeatureDataConnection, item_data: Optional[Union[str, dict]]) -> dict:
    @singledispatch
    def load(data):
        raise ValueError(f'Invalid object data of type {type(data)}: {str(data)}')

    @load.register
    def _(data: None) -> dict:
        return {}

    @load.register
    def _(data: dict) -> dict:
        return data

    @load.register
    def _(data: str) -> dict:
        if _is_ref(data):
            return s3.get_json_object(conn._s3_bucket, _s3_key(data))

        raise ValueError(f'String data must be REF.  Got: {str(data)}')

    return load(item_data)


def _get_data(conn: FeatureDataConnection, dyn_key: dict) -> Union[dict, str]:
    item = conn._table.get_item(Key=dyn_key, ConsistentRead=True)
    return get_in(['Item', 'data'], item)


def _clean(value: str):
    return re.sub(r'\W+', '_', value)


def _is_max_item_size_exception(e: ClientError) -> bool:
    error = e.response['Error']
    return (error['Code'] == 'ValidationException' and
            error.get('Message') == 'Item size has exceeded the maximum allowed size')


def _delete_old_data(conn: FeatureDataConnection, old_item: dict):
    old_item_data = get_in(['Attributes', 'data'], old_item)

    @singledispatch
    def delete(data):
        pass

    @delete.register
    def _(data: str):
        if _is_ref(data):
            s3_client.delete_object(Bucket=conn._s3_bucket, Key=_s3_key(data))
            return

        logger.warning(f'Unexpected string data: {str(data)}')

    delete(old_item_data)


def _put_item(conn: FeatureDataConnection, dyn_key: dict, data: dict, condition=None):
    condition_expr = dict(ConditionExpression=condition) if condition is not None else {}

    def put(item_data):
        old_item = dyn.put_item(conn._table,
                                **condition_expr,
                                ReturnValues='ALL_OLD',
                                Item=dict(
                                    **dyn_key,
                                    data=dyn.type_for_dynamo(item_data)
                                ))
        _delete_old_data(conn, old_item)

    def put_s3(data):
        s3_key = f'{_clean(dyn_key["PK"])}/{_clean(dyn_key["SK"])}_{uuid.uuid4()}.json'
        s3.put_json_object(conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{s3_key}', data)
        return s3_key

    try:
        put(data)
    except ClientError as err:
        if _is_max_item_size_exception(err):
            ref = put_s3(data)
            put(f'{S3_OBJECT_REF}{ref}')
        else:
            raise


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, cz_organization_id: str, keyspace: str, *key: str,
                         update_handler: Callable[[dict], Optional[dict]]) -> Optional[dict]:
    dyn_key = _dyn_key(cz_organization_id, keyspace, *key)

    item_data = _get_data(conn, dyn_key)

    current_object_data = _load_data(conn, item_data)
    new_object_data = update_handler(current_object_data or {})

    if new_object_data is not None:
        data_attr = Attr('data')
        condition = (
            data_attr.eq(item_data)) if item_data is not None else (
                data_attr.not_exists() | data_attr.attribute_type('NULL'))

        _put_item(conn, dyn_key, new_object_data, condition)

    return new_object_data


def get_object(conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str) -> dict:
    data = _get_data(conn, _dyn_key(cz_organization_id, keyspace, *key))
    return _load_data(conn, data)


def list_objects(conn: FeatureDataConnection, cz_organization_id: str,
                 keyspace: str, *key: str) -> Iterable[dict]:
    dyn_key = _dyn_key(cz_organization_id, keyspace, *key, '')
    cond = Key('PK').eq(dyn_key['PK']) & Key('SK').begins_with(dyn_key['SK'])
    items = dyn.query(conn._table, dict(KeyConditionExpression=cond, ConsistentRead=True))
    for item in items:
        yield _load_data(conn, item['data'])


def list_all_objects(conn: FeatureDataConnection,
                     keyspace: str, *key: str) -> Iterable[dict]:
    dyn_key = _dyn_key('', keyspace, *key, '')
    cond = Attr('PK').begins_with(dyn_key['PK']) & Attr('SK').begins_with(dyn_key['SK'])
    items = dyn.scan(conn._table, dict(FilterExpression=cond))
    for item in items:
        yield _load_data(conn, item['data'])


def set_object(conn: FeatureDataConnection, cz_organization_id: str,
               keyspace: str, *key: str, data: dict) -> None:
    dyn_key = _dyn_key(cz_organization_id, keyspace, *key)
    _put_item(conn, dyn_key, data)


def delete_object(conn: FeatureDataConnection, cz_organization_id: str,
                  keyspace: str, *key: str) -> None:
    old_item = conn._table.delete_item(
        Key=_dyn_key(cz_organization_id, keyspace, *key),
        ReturnValues='ALL_OLD'
    )
    _delete_old_data(conn, old_item)
