# 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 collections.abc import Callable, Iterable
from dataclasses import dataclass
from functools import singledispatch
from typing import Generic, Optional, TypeVar, Union

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

logger = get_logger()

#
# 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 list_all_with_organization(self, conn: FeatureDataConnection, keyspace: str, *key: str) -> Iterable[tuple[str, T]]:
        for x, y in list_all_objects_with_organization(conn, keyspace, *key):
            yield x, self._coerce(y)

    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 _unfix_organization(key: str) -> str:
    return dyn.unfix('organization', key)


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 list_all_objects_with_organization(conn: FeatureDataConnection, keyspace: str, *key: str) -> Iterable[tuple[str, 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 _unfix_organization(item['PK']), _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)
