# 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 json
import re
import uuid
from abc import ABC
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from functools import singledispatch
from http import HTTPStatus
from typing import Generic, List, NamedTuple, Optional, Tuple, TypeVar, Union, overload

from boto3.dynamodb.conditions import Attr, Key
from botocore.exceptions import ClientError
from pydantic import BaseModel
from retrying import retry
from toolz import get_in

from czc import dynamodb as dyn
from czc import s3
from czc.aws_environment import aws_service
from czc.constants import DYNAMODB_TTL_OVERRIDE
from czc.dynamodb import COMPOSITE_CHAR, DEFAULT_TTL_KEY, DynamoDbTable
from czc.environ import environ_aware
from czc.helpers import coalesce
from czc.logging import get_logger
from czc.type_util import coerce, deserializable, require, serializable

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')
RawItemData = Union[str, dict]


@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 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=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 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)


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, composite_char=COMPOSITE_CHAR) -> str:
    return dyn.prefix('organization', cz_organization_id, composite_char=composite_char)


def _extract_org_id_from_pk(pk: str) -> str:
    pk_parts = pk.split('#')
    if len(pk_parts) <= 1:
        raise ValueError(f'Invalid PK format, cannot extract organization ID: {pk}')
    return pk_parts[1]


@deserializable()
class DynamoDBKey(NamedTuple):
    PK: str
    SK: str


class RecordType(Enum):
    object_store = 'object_store'
    legacy_feature_data = 'legacy_feature_data'
    legacy_billing = 'legacy_billing'


def _dyn_key(
    cz_organization_id: str,
    keyspace: str,
    key_parts: Tuple[str, ...],
    internal: bool = False,
    record_type: RecordType = RecordType.object_store,
) -> DynamoDBKey:
    def _legacy_billing_key() -> DynamoDBKey:
        return DynamoDBKey(
            PK=_prefix_organization(cz_organization_id, composite_char='-'), SK=dyn.compose_key(keyspace, *key_parts)
        )

    def _legacy_key() -> DynamoDBKey:
        return DynamoDBKey(PK=_prefix_organization(cz_organization_id), SK=dyn.compose_key(keyspace, *key_parts))

    def _key() -> DynamoDBKey:
        return DynamoDBKey(
            PK=dyn.compose_key(_prefix_organization(cz_organization_id), keyspace),
            SK=dyn.compose_key('internal' if internal else 'record', *key_parts),
        )

    return {
        RecordType.object_store: _key,
        RecordType.legacy_feature_data: _legacy_key,
        RecordType.legacy_billing: _legacy_billing_key,
    }[record_type]()


def _legacy_dyn_key(cz_organization_id: str, keyspace: str, *key: str) -> DynamoDBKey:
    return _dyn_key(cz_organization_id, keyspace, key, record_type=RecordType.legacy_feature_data)


class HierarchicalKey(BaseModel, ABC):
    @property
    def key_parts(self) -> tuple[str, ...]:
        return tuple(str(getattr(self, field)) for field in type(self).model_fields.keys())


KeyT = TypeVar('KeyT', bound=HierarchicalKey)
DataT = TypeVar('DataT')


class MigrationMode(Enum):
    none = 'none'
    from_feature_data = 'from_feature_data'
    from_billing = 'from_billing'
    all = 'all'


@dataclass(frozen=True)
class _ObjectStoreCommonParams(Generic[DataT]):
    conn: FeatureDataConnection
    keyspace: str
    obj_type: type[DataT]
    migration_mode: MigrationMode


class _ObjectStoreBase(Generic[DataT]):
    def __init__(self, conn: FeatureDataConnection, params: _ObjectStoreCommonParams):
        self._conn = conn
        self._keyspace = params.keyspace
        self._obj_type = params.obj_type
        self._migration_mode = params.migration_mode

    def _get_dyn_key(
        self, cz_organization_id: str, key: Optional[HierarchicalKey] = None, record_type: RecordType = RecordType.object_store
    ) -> DynamoDBKey:
        return _dyn_key(
            cz_organization_id,
            self._keyspace,
            key.key_parts if key else tuple(),
            record_type=record_type,
        )

    def _get_migration_keys(self, cz_organization_id: str, key: Optional[HierarchicalKey] = None) -> List[DynamoDBKey]:
        migration_path = {
            MigrationMode.none: [RecordType.object_store],
            MigrationMode.from_feature_data: [RecordType.object_store, RecordType.legacy_feature_data],
            MigrationMode.from_billing: [RecordType.object_store, RecordType.legacy_billing],
            MigrationMode.all: [
                RecordType.object_store,
                RecordType.legacy_feature_data,
                RecordType.legacy_billing,
            ],
        }[self._migration_mode]

        return [self._get_dyn_key(cz_organization_id, key, record_type=rt) for rt in migration_path]

    def _get(self, cz_organization_id: str, key: Optional[HierarchicalKey] = None) -> Optional[DataT]:
        result = _get_object(self._conn, self._get_migration_keys(cz_organization_id, key))
        return coerce(self._obj_type, result.data) if result else None

    def _set(self, cz_organization_id: str, data: DataT, key: Optional[HierarchicalKey] = None) -> None:
        dyn_key = self._get_dyn_key(cz_organization_id, key)
        _put_item(self._conn, serializable(dyn_key), serializable(data))

    def _update(
        self,
        cz_organization_id: str,
        handler: Callable[[Optional[DataT]], Optional[DataT]],
        key: Optional[HierarchicalKey] = None,
    ) -> Optional[DataT]:
        def update_handler(data: dict) -> Optional[dict]:
            current_typed = coerce(self._obj_type, data) if data else None
            new_typed = handler(current_typed)
            return serializable(new_typed) if new_typed is not None else None

        dyn_keys = self._get_migration_keys(cz_organization_id, key)
        result = _update_object_atomic(self._conn, dyn_keys, update_handler)
        return coerce(self._obj_type, result) if result else None

    def _delete(self, cz_organization_id: str, key: Optional[HierarchicalKey] = None) -> None:
        dyn_keys = self._get_migration_keys(cz_organization_id, key)
        for dyn_key in reversed(dyn_keys):
            _delete_object(self._conn, dyn_key)


class _ObjectStore(Generic[KeyT, DataT], _ObjectStoreBase[DataT]):
    def __init__(self, conn: FeatureDataConnection, key_type: type[KeyT], params: _ObjectStoreCommonParams):
        super().__init__(conn, params)
        self._key_type = key_type

    def get(self, cz_organization_id: str, key: KeyT) -> Optional[DataT]:
        return self._get(cz_organization_id, key)

    def set(self, cz_organization_id: str, key: KeyT, data: DataT) -> None:
        self._set(cz_organization_id, data, key)

    def update(
        self, cz_organization_id: str, key: KeyT, handler: Callable[[Optional[DataT]], Optional[DataT]]
    ) -> Optional[DataT]:
        return self._update(cz_organization_id, handler, key)

    def delete(self, cz_organization_id: str, key: KeyT) -> None:
        self._delete(cz_organization_id, key)


class _ObjectStoreNoKey(Generic[DataT], _ObjectStoreBase[DataT]):
    def __init__(self, conn: FeatureDataConnection, params: _ObjectStoreCommonParams):
        super().__init__(conn, params)

    def get(self, cz_organization_id: str) -> Optional[DataT]:
        return self._get(cz_organization_id)

    def set(self, cz_organization_id: str, data: DataT) -> None:
        self._set(cz_organization_id, data)

    def update(self, cz_organization_id: str, handler: Callable[[Optional[DataT]], Optional[DataT]]) -> Optional[DataT]:
        return self._update(cz_organization_id, handler)

    def delete(self, cz_organization_id: str) -> None:
        self._delete(cz_organization_id)


@overload
def create_object_store(
    obj_type: type[DataT],
    keyspace: str,
    *,
    key_type: type[KeyT],
    ttl: Optional[timedelta] = None,
    s3_temp_prefix: str = 'tmp',
    migration_mode: MigrationMode = MigrationMode.none,
) -> _ObjectStore[KeyT, DataT]: ...


@overload
def create_object_store(
    obj_type: type[DataT],
    keyspace: str,
    *,
    ttl: Optional[timedelta] = None,
    s3_temp_prefix: str = 'tmp',
    migration_mode: MigrationMode = MigrationMode.none,
) -> _ObjectStoreNoKey[DataT]: ...


def create_object_store(
    obj_type: type[DataT],
    keyspace: str,
    *,
    key_type: Optional[type[KeyT]] = None,
    ttl: Optional[timedelta] = None,
    s3_temp_prefix: str = 'tmp',
    migration_mode: MigrationMode = MigrationMode.none,
):
    conn = connect(ttl=ttl, s3_temp_prefix=s3_temp_prefix)
    common_params = _ObjectStoreCommonParams(
        conn=conn,
        keyspace=keyspace,
        obj_type=obj_type,
        migration_mode=migration_mode,
    )
    if not key_type:
        return _ObjectStoreNoKey(conn, common_params)

    return _ObjectStore(conn, key_type, common_params)


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[RawItemData], 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 None

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

    @load.register
    def _(data: str) -> Optional[dict]:
        is_serialized_obj = data.startswith('{') and data.endswith('}')
        if is_serialized_obj:
            return json.loads(data)

        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) -> Optional[RawItemData]:
    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_http_request_too_large(e: ClientError) -> bool:
    http_status = e.response.get('ResponseMetadata', {}).get('HTTPStatusCode')
    return http_status == HTTPStatus.REQUEST_ENTITY_TOO_LARGE


def _is_dynamo_item_limit_exceeded(e: ClientError) -> bool:
    error = e.response.get('Error', {})
    return error.get('Code') == 'ValidationException' and error.get('Message', '') in (
        'Item size has exceeded the maximum allowed size',
        'Nesting Levels have exceeded supported limits: Attributes in the item have nested levels beyond supported limit',
    )


def _should_offload_to_s3(e: ClientError) -> bool:
    return _is_http_request_too_large(e) or _is_dynamo_item_limit_exceeded(e)


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 _should_offload_to_s3(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,
    dyn_keys: List[DynamoDBKey],
    update_handler: Callable[[dict], Optional[dict]],
) -> Optional[dict]:
    primary_dyn_key = dyn_keys[0]
    current_object_data = _get_object(conn, [primary_dyn_key])
    fallback_object_data = _get_object(conn, dyn_keys[1:]) if current_object_data is None else None
    handler_data = require(coalesce(current_object_data, fallback_object_data, ObjectResult(item_data={}, data={})))

    new_object_data = update_handler(handler_data.data)

    if new_object_data is not None:
        item_data = current_object_data.item_data if current_object_data else 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, serializable(primary_dyn_key), new_object_data, condition)

    return new_object_data


class ObjectResult(NamedTuple):
    item_data: RawItemData
    data: dict


def _get_object(conn: FeatureDataConnection, dyn_keys: List[DynamoDBKey]) -> Optional[ObjectResult]:
    def get_obj(dyn_key: DynamoDBKey) -> Optional[ObjectResult]:
        serialized_key = serializable(dyn_key)
        item_data = _get_data(conn, serialized_key)
        data = _load_data(conn, item_data, serialized_key)
        return ObjectResult(require(item_data), data) if data is not None else None

    return next((x for k in dyn_keys if (x := get_obj(k)) is not None), None)


def _list_objects(conn: FeatureDataConnection, dyn_key: DynamoDBKey) -> Iterable[dict]:
    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, dyn_key: DynamoDBKey) -> Iterable[tuple[str, dict]]:
    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:
            continue
        yield _extract_org_id_from_pk(item['PK']), data


def _delete_object(conn: FeatureDataConnection, dyn_key: DynamoDBKey) -> None:
    old_item = conn._table.delete_item(Key=serializable(dyn_key), ReturnValues='ALL_OLD')
    _delete_old_data(conn, old_item)


def update_object_atomic(
    conn: FeatureDataConnection,
    cz_organization_id: str,
    keyspace: str,
    *key: str,
    update_handler: Callable[[dict], Optional[dict]],
) -> Optional[dict]:
    dyn_key = _legacy_dyn_key(cz_organization_id, keyspace, *key)
    return _update_object_atomic(conn, [dyn_key], update_handler)


def get_object(conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str) -> dict:
    dyn_key = _legacy_dyn_key(cz_organization_id, keyspace, *key)
    result = _get_object(conn, [dyn_key])
    return result.data if result else {}


def list_objects(conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str) -> Iterable[dict]:
    dyn_key = _legacy_dyn_key(cz_organization_id, keyspace, *key, '')
    return _list_objects(conn, dyn_key)


def list_all_objects_with_organization(conn: FeatureDataConnection, keyspace: str, *key: str) -> Iterable[tuple[str, dict]]:
    dyn_key = _legacy_dyn_key('', keyspace, *key, '')
    return _list_all_objects_with_organization(conn, dyn_key)


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 = _legacy_dyn_key(cz_organization_id, keyspace, *key)
    _put_item(conn, serializable(dyn_key), data)


def delete_object(conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str) -> None:
    dyn_key = _legacy_dyn_key(cz_organization_id, keyspace, *key)
    _delete_object(conn, dyn_key)
