# 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, abstractmethod
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 inspect import isclass
from types import UnionType
from typing import (
    Generic,
    List,
    NamedTuple,
    Optional,
    Tuple,
    TypeVar,
    Union,
    get_args,
    get_origin,
    overload,
)

from boto3.dynamodb.conditions import Attr, Key
from botocore.exceptions import ClientError
from pydantic import BaseModel, ConfigDict, Field, model_validator
from retrying import retry
from toolz import concat, get_in, unique

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'
INTERNAL_KEY = 'internal'
RECORD_KEY = 'record'
SNAPSHOT_REGISTRY_KEY = 'snapshot_registry'
LEGACY_BILLING_COMPOSITE_CHAR = '-'

T = TypeVar('T')
RawItemData = Union[str, dict]


class SnapshotReadOnlyError(Exception):
    def __init__(self):
        super().__init__('Snapshots are read-only and cannot be modified')


class SnapshotNotSupportedError(Exception):
    def __init__(
        self,
    ):
        super().__init__('Snapshots do not support this operation')


class InvalidSnapshotIdError(Exception):
    pass


class SnapshotCreationError(Exception):
    pass


class SnapshotDeletionError(Exception):
    pass


@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_with_keys(
        self, conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str
    ) -> Iterable[tuple[tuple[str, ...], T]]:
        for key_parts, x in list_objects_with_keys(conn, cz_organization_id, keyspace, *key):
            yield key_parts, 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:
    if (unfixed := dyn.unfix('organization', pk)) == pk:
        raise ValueError(f'Invalid PK format, cannot extract organization ID: {pk}')
    return unfixed.split(COMPOSITE_CHAR)[0]


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


class DynamoDBKey(NamedTuple):
    PK: str
    SK: str
    record_type: RecordType

    def to_key_dict(self) -> dict:
        return {'PK': self.PK, 'SK': self.SK}


def _dyn_key(
    cz_organization_id: str,
    keyspace: str,
    key_parts: Tuple[str, ...],
    internal: bool = False,
    record_type: RecordType = RecordType.object_store,
    snapshot_id: Optional[str] = None,
) -> DynamoDBKey:
    def _legacy_billing_alt_key() -> DynamoDBKey:
        if len(key_parts) > 1:
            raise ValueError('Legacy billing alt key can have at most one key part')

        return DynamoDBKey(
            PK=_prefix_organization(cz_organization_id, composite_char=LEGACY_BILLING_COMPOSITE_CHAR),
            SK=dyn.compose_key(keyspace, *key_parts, composite_char=LEGACY_BILLING_COMPOSITE_CHAR),
            record_type=RecordType.legacy_billing_alt,
        )

    def _legacy_billing_key() -> DynamoDBKey:
        return DynamoDBKey(
            PK=_prefix_organization(cz_organization_id, composite_char=LEGACY_BILLING_COMPOSITE_CHAR),
            SK=dyn.compose_key(keyspace, *key_parts),
            record_type=RecordType.legacy_billing,
        )

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

    def _key() -> DynamoDBKey:
        pk = dyn.compose_key(_prefix_organization(cz_organization_id), keyspace)
        return DynamoDBKey(
            PK=dyn.compose_key(snapshot_id, pk) if snapshot_id else pk,
            SK=dyn.compose_key(INTERNAL_KEY if internal else RECORD_KEY, *key_parts, ''),
            record_type=RecordType.object_store,
        )

    return {
        RecordType.object_store: _key,
        RecordType.legacy_feature_data: _legacy_key,
        RecordType.legacy_billing: _legacy_billing_key,
        RecordType.legacy_billing_alt: _legacy_billing_alt_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):
    model_config = ConfigDict(frozen=True, extra='forbid')

    @classmethod
    def __pydantic_init_subclass__(cls, **kwargs):
        super().__pydantic_init_subclass__(**kwargs)

        for field_name, field_info in cls.model_fields.items():
            annotation = field_info.annotation

            # Check if the annotation is Optional (Union with None)
            # Support both typing.Union (Optional[str]) and types.UnionType (str | None)
            origin = get_origin(annotation)
            if origin not in (Union, UnionType):
                raise TypeError(f'Field {field_name} in {cls.__name__} must be Optional')

            args = get_args(annotation)
            if type(None) not in args:
                raise TypeError(f'Field {field_name} in {cls.__name__} must be Optional (include None in Union)')

            # Get all non-None types
            inner_types = [arg for arg in args if arg is not type(None)]

            # All non-None types must be str or str subclasses
            for inner_type in inner_types:
                if not (isclass(inner_type) and issubclass(inner_type, str)):
                    raise TypeError(
                        f'Field {field_name} in {cls.__name__} must be Optional[str] or Optional[StrSubclass], '
                        f'but found {inner_type}'
                    )

            # Check that the field has None as default value
            if field_info.default is not None:
                raise TypeError(f'Field {field_name} in {cls.__name__} must have None as default value')

    @property
    def _field_names(self) -> tuple[str, ...]:
        return tuple(type(self).model_fields.keys())

    @property
    def _length(self) -> int:
        return next((i for i, name in enumerate(self._field_names) if getattr(self, name) is None), len(self._field_names))

    @model_validator(mode='after')
    def _prefix_rule(self):
        fields_after_none = {
            self._field_names[i]
            for i in range(self._length, len(self._field_names))
            if getattr(self, self._field_names[i]) is not None
        }

        if fields_after_none:
            raise ValueError('Fields provided must be a prefix')

        return self

    @property
    def is_complete(self) -> bool:
        return self._length == len(self._field_names)

    @property
    def key_parts(self) -> tuple[str, ...]:
        return tuple(str(getattr(self, name)) for name in self._field_names[: self._length])


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


class ItemRecord(NamedTuple, Generic[KeyT, DataT]):
    key: KeyT
    data: DataT


class ListRecord(NamedTuple, Generic[DataT]):
    organization_id: str
    data: DataT


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


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


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
        self._snapshot_id = params.snapshot_id

    def _get_dyn_key(
        self,
        cz_organization_id: str,
        key: Optional[HierarchicalKey] = None,
        record_type: RecordType = RecordType.object_store,
        for_list: bool = False,
    ) -> Optional[DynamoDBKey]:
        keyparts = key.key_parts if key else tuple()

        needs_trailing_char = for_list and record_type != RecordType.object_store
        all_keyparts = (*keyparts, '') if needs_trailing_char else keyparts
        if record_type == RecordType.legacy_billing_alt and len(all_keyparts) > 1:
            return None

        return _dyn_key(
            cz_organization_id,
            self._keyspace,
            all_keyparts,
            record_type=record_type,
            snapshot_id=self._snapshot_id,
        )

    def _get_migration_keys(
        self, cz_organization_id: str, key: Optional[HierarchicalKey] = None, for_list: bool = False
    ) -> 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, RecordType.legacy_billing_alt],
        }[self._migration_mode]

        migration_keys = [
            self._get_dyn_key(cz_organization_id, key, record_type=rt, for_list=for_list) for rt in migration_path
        ]
        return [x for x in migration_keys if x is not None]

    def _get(self, cz_organization_id: str, key: Optional[HierarchicalKey] = None) -> Optional[DataT]:
        if self._snapshot_id:
            _validate_snapshot(self._conn, cz_organization_id, self._snapshot_id, self._keyspace)
        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:
        if self._snapshot_id:
            raise SnapshotReadOnlyError()
        dyn_key = require(self._get_dyn_key(cz_organization_id, key))
        _put_item(self._conn, dyn_key.to_key_dict(), serializable(data))

    def _update(
        self,
        cz_organization_id: str,
        handler: Callable[[Optional[DataT]], Optional[DataT]],
        key: Optional[HierarchicalKey] = None,
    ) -> Optional[DataT]:
        if self._snapshot_id:
            raise SnapshotReadOnlyError()

        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:
        if self._snapshot_id:
            raise SnapshotReadOnlyError()
        dyn_keys = self._get_migration_keys(cz_organization_id, key)
        for dyn_key in reversed(dyn_keys):
            _delete_object(self._conn, dyn_key)

    def _list_all(self, key: Optional[HierarchicalKey] = None) -> Iterable[ListRecord[DataT]]:
        if self._snapshot_id:
            raise SnapshotNotSupportedError()
        dyn_keys = self._get_migration_keys('', key, for_list=True)
        for org_id, obj in _list_all_objects(self._conn, self._keyspace, dyn_keys):
            yield ListRecord(organization_id=org_id, data=coerce(self._obj_type, obj))


class ObjectStore(ABC, Generic[KeyT, DataT]):
    """A keyed object store for persisting typed data objects per organization.

    Each record is identified by a combination of organization ID and a hierarchical key.
    Use ``create_object_store`` with a ``key_type`` argument to obtain an instance.
    """

    @abstractmethod
    def get(self, cz_organization_id: str, key: KeyT) -> Optional[DataT]:
        """Retrieve a single object by its complete key.

        Returns ``None`` if no record exists for the given organization and key.
        """

    @abstractmethod
    def set(self, cz_organization_id: str, key: KeyT, data: DataT) -> None:
        """Store an object, replacing any existing record at the same key."""

    @abstractmethod
    def update(
        self, cz_organization_id: str, key: KeyT, handler: Callable[[Optional[DataT]], Optional[DataT]]
    ) -> Optional[DataT]:
        """Atomically read-modify-write an object.

        ``handler`` receives the current value (or ``None`` if absent) and returns the new
        value to persist. Return ``None`` from the handler to delete the record.
        """

    @abstractmethod
    def delete(self, cz_organization_id: str, key: KeyT) -> None:
        """Delete the object at the given key."""

    @abstractmethod
    def list(self, cz_organization_id: str, key: Optional[KeyT] = None) -> Iterable[DataT]:
        """List data objects under a key prefix for a single organization.

        Pass a partial key to filter results, or ``None`` to list all objects in the keyspace.
        """

    @abstractmethod
    def items(self, cz_organization_id: str, key: Optional[KeyT] = None) -> Iterable[ItemRecord[KeyT, DataT]]:
        """List key/data pairs under a key prefix for a single organization.

        Like ``list``, but yields ``ItemRecord`` tuples containing both the reconstructed
        key and the data object.
        """

    @abstractmethod
    def list_all(self, key: Optional[KeyT] = None) -> Iterable[ListRecord[DataT]]:
        """List data objects across all organizations.

        Yields ``ListRecord`` tuples containing the organization ID and data object.
        """


class ObjectStoreNoKey(ABC, Generic[DataT]):
    """A keyless object store for persisting a single typed data object per organization.

    Each organization has at most one record in the keyspace.
    Use ``create_object_store`` without a ``key_type`` argument to obtain an instance.
    """

    @abstractmethod
    def get(self, cz_organization_id: str) -> Optional[DataT]:
        """Retrieve the object for the given organization.

        Returns ``None`` if no record exists.
        """

    @abstractmethod
    def set(self, cz_organization_id: str, data: DataT) -> None:
        """Store an object, replacing any existing record for the organization."""

    @abstractmethod
    def update(self, cz_organization_id: str, handler: Callable[[Optional[DataT]], Optional[DataT]]) -> Optional[DataT]:
        """Atomically read-modify-write the object for an organization.

        ``handler`` receives the current value (or ``None`` if absent) and returns the new
        value to persist. Return ``None`` from the handler to delete the record.
        """

    @abstractmethod
    def delete(self, cz_organization_id: str) -> None:
        """Delete the object for the given organization."""

    @abstractmethod
    def list_all(self) -> Iterable[ListRecord[DataT]]:
        """List data objects across all organizations.

        Yields ``ListRecord`` tuples containing the organization ID and data object.
        """


class _ObjectStore(ObjectStore[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 _require_complete_key(self, key: KeyT) -> KeyT:
        if not key.is_complete:
            raise ValueError('Complete key required for this operation')
        return key

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

    def set(self, cz_organization_id: str, key: KeyT, data: DataT) -> None:
        self._set(cz_organization_id, data, self._require_complete_key(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, self._require_complete_key(key))

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

    def list(self, cz_organization_id: str, key: Optional[KeyT] = None) -> Iterable[DataT]:
        if self._snapshot_id:
            _validate_snapshot(self._conn, cz_organization_id, self._snapshot_id, self._keyspace)
        hk = key or self._key_type()
        dyn_keys = self._get_migration_keys(cz_organization_id, hk, for_list=True)
        for obj in _list_objects(self._conn, dyn_keys):
            yield coerce(self._obj_type, obj)

    def items(self, cz_organization_id: str, key: Optional[KeyT] = None) -> Iterable[ItemRecord[KeyT, DataT]]:
        if self._snapshot_id:
            _validate_snapshot(self._conn, cz_organization_id, self._snapshot_id, self._keyspace)
        hk = key or self._key_type()
        dyn_keys = self._get_migration_keys(cz_organization_id, hk, for_list=True)
        field_names = tuple(self._key_type.model_fields.keys())
        prefix_parts = hk.key_parts
        for key_parts, data in _list_objects_with_keys(self._conn, dyn_keys):
            all_parts = prefix_parts + key_parts
            reconstructed_key = self._key_type(**dict(zip(field_names, all_parts)))
            yield ItemRecord(reconstructed_key, coerce(self._obj_type, data))

    def list_all(self, key: Optional[KeyT] = None) -> Iterable[ListRecord[DataT]]:
        hk = key or self._key_type()
        return self._list_all(hk)


class _ObjectStoreNoKey(ObjectStoreNoKey[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)

    def list_all(self) -> Iterable[ListRecord[DataT]]:
        return self._list_all()


@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,
    snapshot_id: Optional[str] = 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,
    snapshot_id: Optional[str] = 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,
    snapshot_id: Optional[str] = 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,
        snapshot_id=snapshot_id,
    )
    if not key_type:
        return _ObjectStoreNoKey(conn, common_params)

    return _ObjectStore(conn, key_type, common_params)


@deserializable()
class SnapshotRegistry(BaseModel):
    model_config = ConfigDict(frozen=True)

    snapshot_id: str
    keyspaces: set[str]
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())


def _snapshot_registry_dyn_key(cz_organization_id: str, snapshot_id: str) -> DynamoDBKey:
    return DynamoDBKey(
        PK=_prefix_organization(cz_organization_id),
        SK=dyn.compose_key(INTERNAL_KEY, SNAPSHOT_REGISTRY_KEY, snapshot_id),
        record_type=RecordType.object_store,
    )


def _read_snapshot_registry(conn: FeatureDataConnection, cz_organization_id: str, snapshot_id: str) -> SnapshotRegistry:
    dyn_key = _snapshot_registry_dyn_key(cz_organization_id, snapshot_id)
    result = _get_object(conn, [dyn_key])
    if not result or not result.data:
        raise InvalidSnapshotIdError(f'Snapshot {snapshot_id} not found for organization {cz_organization_id}')
    return coerce(SnapshotRegistry, result.data)


def _validate_snapshot(conn: FeatureDataConnection, cz_organization_id: str, snapshot_id: str, keyspace: str) -> None:
    registry = _read_snapshot_registry(conn, cz_organization_id, snapshot_id)
    if keyspace not in registry.keyspaces:
        raise InvalidSnapshotIdError(
            f'Snapshot {snapshot_id} does not include keyspace "{keyspace}". '
            f'Available keyspaces: {", ".join(registry.keyspaces)}'
        )


def _write_snapshot_registry(
    conn: FeatureDataConnection, cz_organization_id: str, snapshot_id: str, keyspaces: set[str]
) -> None:
    dyn_key = _snapshot_registry_dyn_key(cz_organization_id, snapshot_id)
    registry = SnapshotRegistry(snapshot_id=snapshot_id, keyspaces=keyspaces)
    _put_item(conn, serializable(dyn_key), serializable(registry))


def create_object_store_snapshot(cz_organization_id: str, snapshot_id: str, keyspaces: set[str]) -> None:
    if not keyspaces:
        raise ValueError('At least one keyspace is required to create a snapshot')

    conn = connect()

    try:
        for keyspace in keyspaces:
            for dyn_key, data in _list_objects_in_keyspace(conn, cz_organization_id, keyspace):
                snapshot_key = DynamoDBKey(
                    PK=dyn.compose_key(snapshot_id, dyn_key.PK), SK=dyn_key.SK, record_type=RecordType.object_store
                )
                _put_item(conn, serializable(snapshot_key), data)

        _write_snapshot_registry(conn, cz_organization_id, snapshot_id, keyspaces)
    except Exception as e:
        raise SnapshotCreationError(f'Failed to create snapshot {snapshot_id}') from e


def delete_object_store_snapshot(cz_organization_id: str, snapshot_id: str) -> None:
    conn = connect()
    registry = _read_snapshot_registry(conn, cz_organization_id, snapshot_id)

    try:
        for keyspace in registry.keyspaces:
            pk = dyn.compose_key(snapshot_id, _prefix_organization(cz_organization_id), keyspace)
            items = dyn.query(conn._table, dict(KeyConditionExpression=Key('PK').eq(pk)))
            for item in items:
                dyn_key = DynamoDBKey(PK=item['PK'], SK=item['SK'], record_type=RecordType.object_store)
                _delete_object(conn, dyn_key)

        registry_key = _snapshot_registry_dyn_key(cz_organization_id, snapshot_id)
        _delete_object(conn, registry_key)
    except Exception as e:
        raise SnapshotDeletionError(f'Failed to delete snapshot {snapshot_id}') from e


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, primary_dyn_key.to_key_dict(), 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 = dyn_key.to_key_dict()
        item_data = _get_data(conn, serialized_key)
        data = _load_data(conn, item_data, serialized_key)
        if data is None or item_data is None:
            return None
        return ObjectResult(item_data, data)

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


def _sk_suffix(dyn_key: DynamoDBKey, full_sk: str) -> str:
    suffix = full_sk[len(dyn_key.SK) :]
    return suffix.rstrip(COMPOSITE_CHAR)


def _list_objects_with_keys(conn: FeatureDataConnection, dyn_keys: List[DynamoDBKey]) -> Iterable[tuple[tuple[str, ...], dict]]:
    class ItemData(NamedTuple):
        sk_suffix: str
        key_parts: tuple[str, ...]
        data: dict

    def list_item_data(dyn_key: DynamoDBKey) -> Iterable[ItemData]:
        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_sk = item['SK']
            item_key = {'PK': item['PK'], 'SK': item_sk}
            data = _load_data(conn, item['data'], item_key)
            if data is None:
                continue
            sk_suffix = _sk_suffix(dyn_key, item_sk)
            yield ItemData(sk_suffix=sk_suffix, key_parts=tuple(dyn.decompose_key(sk_suffix)), data=data)

    for item_data in unique(concat(list_item_data(dyn_key) for dyn_key in dyn_keys), key=lambda x: x.sk_suffix):
        yield item_data.key_parts, item_data.data


def _list_objects(conn: FeatureDataConnection, dyn_keys: List[DynamoDBKey]) -> Iterable[dict]:
    for _, data in _list_objects_with_keys(conn, dyn_keys):
        yield data


def _legacy_list_all_objects(conn: FeatureDataConnection, dyn_key: DynamoDBKey) -> Iterable[tuple[DynamoDBKey, 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 DynamoDBKey(PK=item['PK'], SK=item['SK'], record_type=dyn_key.record_type), data


def _list_all_objects(
    conn: FeatureDataConnection,
    keyspace: str,
    dyn_keys: List[DynamoDBKey],
) -> Iterable[tuple[str, dict]]:
    class ItemData(NamedTuple):
        organization_id: str
        sk_suffix: str
        data: dict

    def list_all_item_data_for_object_store(dyn_key: DynamoDBKey) -> Iterable[ItemData]:
        org_prefix = _prefix_organization('')
        keyspace_suffix = f'{COMPOSITE_CHAR}{keyspace}'
        cond = Attr('PK').begins_with(org_prefix) & Attr('PK').contains(keyspace_suffix) & Attr('SK').begins_with(dyn_key.SK)
        items = dyn.scan(conn._table, dict(FilterExpression=cond))
        for item in items:
            item_pk = item['PK']
            item_sk = item['SK']
            if not item_pk.endswith(keyspace_suffix):
                continue
            item_key = {'PK': item_pk, 'SK': item_sk}
            data = _load_data(conn, item['data'], item_key)
            if data is None:
                continue
            yield ItemData(organization_id=_extract_org_id_from_pk(item_pk), sk_suffix=_sk_suffix(dyn_key, item_sk), data=data)

    def list_all_item_data_for_legacy(dyn_key: DynamoDBKey) -> Iterable[ItemData]:
        def unfix_org(pk: str) -> str:
            for comp_char in [COMPOSITE_CHAR, LEGACY_BILLING_COMPOSITE_CHAR]:
                unfixed = dyn.unfix('organization', pk, composite_char=comp_char)
                if unfixed != pk:
                    return unfixed
            raise ValueError(f'Invalid PK format, cannot extract organization ID: {pk}')

        for result_key, data in _legacy_list_all_objects(conn, dyn_key):
            yield ItemData(organization_id=unfix_org(result_key.PK), sk_suffix=_sk_suffix(dyn_key, result_key.SK), data=data)

    def list_all_item_data(dyn_key: DynamoDBKey) -> Iterable[ItemData]:
        return (
            list_all_item_data_for_object_store(dyn_key)
            if dyn_key.record_type == RecordType.object_store
            else list_all_item_data_for_legacy(dyn_key)
        )

    for item_data in unique(
        concat(list_all_item_data(dyn_key) for dyn_key in dyn_keys), key=lambda x: (x.organization_id, x.sk_suffix)
    ):
        yield item_data.organization_id, item_data.data


def _list_objects_in_keyspace(
    conn: FeatureDataConnection, cz_organization_id: str, keyspace: str
) -> Iterable[tuple[DynamoDBKey, dict]]:
    pk = dyn.compose_key(_prefix_organization(cz_organization_id), keyspace)
    items = dyn.query(conn._table, dict(KeyConditionExpression=Key('PK').eq(pk)))

    for item in items:
        item_key = DynamoDBKey(PK=item['PK'], SK=item['SK'], record_type=RecordType.object_store)
        data = _load_data(conn, item.get('data'), serializable(item_key))
        if data is not None:
            yield item_key, data


def _delete_object(conn: FeatureDataConnection, dyn_key: DynamoDBKey) -> None:
    old_item = conn._table.delete_item(Key=dyn_key.to_key_dict(), 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]:
    for _, data in list_objects_with_keys(conn, cz_organization_id, keyspace, *key):
        yield data


def list_objects_with_keys(
    conn: FeatureDataConnection, cz_organization_id: str, keyspace: str, *key: str
) -> Iterable[tuple[tuple[str, ...], dict]]:
    """Like list_objects, but yields (key_parts, data) tuples.

    key_parts is a tuple of strings representing the decomposed sort key
    components after the keyspace prefix.
    """
    dyn_key = _legacy_dyn_key(cz_organization_id, keyspace, *key, '')
    yield from _list_objects_with_keys(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, '')
    for result_key, data in _legacy_list_all_objects(conn, dyn_key):
        yield _extract_org_id_from_pk(result_key.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 = _legacy_dyn_key(cz_organization_id, keyspace, *key)
    _put_item(conn, dyn_key.to_key_dict(), 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)
