# 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 datetime import datetime, timedelta
from functools import lru_cache, singledispatch
from typing import Generic, NamedTuple, Optional, TypeVar, Union

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.aws_environment import aws_service
from czc.constants import DYNAMODB_TTL_OVERRIDE
from czc.dynamodb import DEFAULT_TTL_KEY, DynamoDbTable
from czc.environ import environ_aware
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.
#

# max number of retries when resolving listed feature data
MAX_RESOLVE_DATA_TRIES = 3

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


@dataclass(frozen=True)
class FeatureDataConnection:
    _table: DynamoDbTable
    _s3_bucket: str
    _ttl: Optional[timedelta]
    _s3_temp_prefix: 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 org_id, x in list_all_objects_with_organization(conn, keyspace, *key):
            yield org_id, 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)


@lru_cache(maxsize=16)
def connect(ttl: Optional[timedelta] = None, s3_temp_prefix: str = 'tmp') -> FeatureDataConnection:
    table_name = dyn.get_table_name()
    return FeatureDataConnection(
        _table=dyn.table(table_name), _s3_bucket=s3.get_bucket_name(), _ttl=ttl, _s3_temp_prefix=s3_temp_prefix
    )


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


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


class _S3Ref(NamedTuple):
    ref: str

    @staticmethod
    def create(base_key: str) -> '_S3Ref':
        return _S3Ref(f'{S3_OBJECT_REF}{base_key}')

    @staticmethod
    def create_temp(base_key: str, temp_prefix: str) -> '_S3Ref':
        return _S3Ref(f'{S3_TEMP_OBJECT_REF}{temp_prefix}/{S3_OBJECT_STORE_ROOT}/{base_key}')

    @property
    def is_temp(self) -> bool:
        return self.ref.startswith(S3_TEMP_OBJECT_REF)

    @property
    def is_valid(self) -> bool:
        return self.is_temp or self.ref.startswith(S3_OBJECT_REF)

    @property
    def key(self) -> str:
        if not self.is_valid:
            raise ValueError(f'String data must be REF or TEMP_REF.  Got: {self.ref}')

        return (
            self.ref[len(S3_TEMP_OBJECT_REF) :] if self.is_temp else f'{S3_OBJECT_STORE_ROOT}/{self.ref[len(S3_OBJECT_REF) :]}'
        )


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


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

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

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

    @load.register
    def _(data: str) -> Optional[dict]:
        s3_ref = _S3Ref(data)
        try:
            return s3.get_json_object(conn._s3_bucket, s3_ref.key)
        except ClientError as e:
            if _is_no_such_key_exception(e) and s3_ref.is_temp:
                logger.warning(
                    f'Temporary data is missing from S3 for key: {s3_ref.key}. '
                    f'Confirm that the "delete" Lifecycle Policy for this bucket/key '
                    f'is >= the TTL being used.'
                )
                return None
            raise

    @retry(
        retry_on_exception=_is_no_such_key_exception,
        stop_max_attempt_number=MAX_RESOLVE_DATA_TRIES,
        wait_random_min=1000,
        wait_random_max=2000,
    )
    def get_with_retry(key):
        logger.warning(f'Feature data reference is missing for key: {key}.  Retrying')
        data = _get_data(conn, key)
        if data is None:
            return None
        return load(data)

    try:
        return load(item_data)
    except ClientError as e:
        if _is_no_such_key_exception(e):
            return get_with_retry(item_key)
        raise


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):
        s3_ref = _S3Ref(data)
        if not s3_ref.is_valid:
            logger.warning(f'Unexpected string data: {str(data)}')
            return

        aws_service.s3_client.delete_object(Bucket=conn._s3_bucket, Key=s3_ref.key)

    delete(old_item_data)


@environ_aware(optional=[DYNAMODB_TTL_OVERRIDE])
def _ttl_attr(ttl: Optional[timedelta], **kwargs) -> dict:
    if not ttl:
        return {}
    ttl_key = kwargs.get(DYNAMODB_TTL_OVERRIDE) or DEFAULT_TTL_KEY
    return {ttl_key: int((datetime.now() + ttl).timestamp())}


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, **_ttl_attr(conn._ttl), data=dyn.type_for_dynamo(item_data)),
        )
        _delete_old_data(conn, old_item)

    def put_s3(data):
        s3_base_key = f'{_clean(dyn_key["PK"])}/{_clean(dyn_key["SK"])}_{uuid.uuid4()}.json'
        s3_ref = _S3Ref.create_temp(s3_base_key, conn._s3_temp_prefix) if conn._ttl else _S3Ref.create(s3_base_key)
        s3.put_json_object(conn._s3_bucket, s3_ref.key, data)
        return s3_ref

    try:
        put(data)
    except ClientError as err:
        if _is_max_item_size_exception(err):
            put(put_s3(data).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, dyn_key)
    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:
    dyn_key = _dyn_key(cz_organization_id, keyspace, *key)
    data = _get_data(conn, dyn_key)
    return _load_data(conn, data, dyn_key) or {}


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:
        item_key = {'PK': item['PK'], 'SK': item['SK']}
        data = _load_data(conn, item['data'], item_key)
        if data is None:
            continue
        yield 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:
        item_key = {'PK': item['PK'], 'SK': item['SK']}
        data = _load_data(conn, item['data'], item_key)
        if data is None:
            # skip if there was actually no data
            continue
        yield _unfix_organization(item['PK']), data


def list_all_objects(conn: FeatureDataConnection, keyspace: str, *key: str) -> Iterable[dict]:
    for _, x in list_all_objects_with_organization(conn, keyspace, *key):
        yield x


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)
