# Copyright (c) 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
from datetime import datetime, timedelta, timezone
from enum import Enum, StrEnum
from typing import Any, Callable, NamedTuple, Optional
from unittest.mock import ANY, Mock, call
from uuid import uuid4

import freezegun
import pytest
from boto3.dynamodb.conditions import Attr, Key
from botocore.exceptions import ClientError
from pydantic import BaseModel

from czc import dynamodb as dyn
from czc import feature_data
from czc.feature_data import (
    S3_OBJECT_STORE_ROOT,
    SNAPSHOT_REGISTRY_KEY,
    FeatureDataConnection,
    HierarchicalKey,
    HKPart,
    ListRecord,
    MigrationMode,
)
from czc.type_util import deserializable
from czt.unittest import MATCH, Context, MockEx

UTC = timezone.utc

EXCEED_MAX_SIZE_ERROR = ClientError(
    error_response=dict(Error=dict(Code='ValidationException', Message='Item size has exceeded the maximum allowed size')),
    operation_name='PutItem',
)
EXCEED_NESTING_LEVELS_ERROR = ClientError(
    error_response=dict(
        Error=dict(
            Code='ValidationException',
            Message='Nesting Levels have exceeded supported limits: Attributes in the item have nested levels beyond supported limit',
        )
    ),
    operation_name='PutItem',
)
HTTP_413_ERROR = ClientError(
    error_response=dict(
        Error=dict(Code='RequestEntityTooLarge', Message='Request entity too large'),
        ResponseMetadata=dict(HTTPStatusCode=413),
    ),
    operation_name='PutItem',
)
NOT_FOUND_ERROR = ClientError(
    error_response=dict(Error=dict(Code='NoSuchKey', Message='The specified key does not exist.')),
    operation_name='GetObject',
)
CONDITIONAL_CHECK_FAILED_ERROR = ClientError(
    error_response=dict(Error=dict(Code='ConditionalCheckFailedException', Message='The conditional request failed')),
    operation_name='PutItem',
)


@pytest.fixture(scope='function')
def context(mocker):
    mock_table = Mock()
    mock_aws_service = mocker.patch.object(feature_data, 'aws_service', autospec=False)
    extra = dict(
        conn=FeatureDataConnection(_table=mock_table, _s3_bucket=str(uuid4()), _ttl=None, _s3_temp_prefix='tmp'),
        org_id=str(uuid4()),
        keyspace=str(uuid4()),
        table=mock_table,
        mock_aws_service=mock_aws_service,
    )
    with Context(mocker, feature_data, ['dyn', 's3', 'uuid', 'connect'], extra=extra) as context:
        context.mock_connect.return_value = context.conn
        # Allow prefix/compose_key/unfix to use real implementation
        context.mock_dyn.prefix = dyn.prefix
        context.mock_dyn.compose_key = dyn.compose_key
        context.mock_dyn.unfix = dyn.unfix
        context.mock_dyn.type_for_dynamo = dyn.type_for_dynamo
        context.mock_uuid.uuid4.return_value = 'UUID'
        mocker.patch('retrying.time.sleep')
        yield context


@deserializable()
class DataObject(NamedTuple):
    val: int
    name: str
    data: list[str]


class Enum1(Enum):
    A = 1
    B = 2
    C = 3


class Enum2(StrEnum):
    X = 'x'
    Y = 'y'
    Z = 'z'


@deserializable()
class Model1(BaseModel):
    a: int
    b: float
    c: str
    d: dict[str, str]
    e: list[str]
    f: datetime
    g: Enum1
    h: DataObject
    i: tuple[str, str]


@deserializable()
class Model2(BaseModel):
    a: str
    b: str


class Key1(HierarchicalKey):
    region: HKPart
    account: HKPart
    service: HKPart

    def expected_sk(self) -> str:
        parts = filter(None, [self.region, self.account, self.service])
        return '#'.join(['record', *parts])


class Key2(HierarchicalKey):
    env: HKPart
    some_enum: HKPart

    def expected_sk(self) -> str:
        parts = filter(None, [self.env, self.some_enum])
        return '#'.join(['record', *parts])


def _key1(region='us-east-1', account='123456789', service='s3') -> Key1:
    return Key1(region=region, account=account, service=service)


def _key2(env='prod', some_enum=Enum2.X) -> Key2:
    return Key2(env=env, some_enum=some_enum)


TEST_DATA = [
    (dict(val=1, name='foo', data=['one', 'two']), DataObject(val=1, name='foo', data=['one', 'two'])),
    (dict(val=2, name='bar', data=['three', 'four']), DataObject(val=2, name='bar', data=['three', 'four'])),
]


def _expected_pk(org_id: str, keyspace: str) -> str:
    return f'organization#{org_id}#{keyspace}'


def _expected_object_store_key(org_id: str, keyspace: str, key: Optional[HierarchicalKey] = None) -> dict:
    """Build expected DynamoDB key for object_store format."""
    key_parts = '#'.join(key.key_parts) if key else ''
    return {
        'PK': f'organization#{org_id}#{keyspace}',
        'SK': f'record#{key_parts}' if key else 'record',
    }


def _expected_legacy_feature_data_key(org_id: str, keyspace: str, key: Optional[HierarchicalKey] = None) -> dict:
    """Build expected DynamoDB key for legacy_feature_data format."""
    key_parts = '#'.join(key.key_parts) if key else ''
    return {
        'PK': f'organization#{org_id}',
        'SK': f'{keyspace}#{key_parts}' if key else keyspace,
    }


def _expected_legacy_billing_key(org_id: str, keyspace: str, key: Optional[HierarchicalKey] = None) -> dict:
    """Build expected DynamoDB key for legacy_billing format."""
    key_parts = '#'.join(key.key_parts) if key else ''
    return {
        'PK': f'organization-{org_id}',
        'SK': f'{keyspace}#{key_parts}' if key else keyspace,
    }


# Helper functions for test code reuse


def _create_store(obj_type: type, keyspace: str, key: Optional[HierarchicalKey] = None, **kwargs: Any) -> Any:
    """Helper to create object store with or without key."""
    if key:
        return feature_data.create_object_store(obj_type, keyspace, key_type=type(key), **kwargs)
    return feature_data.create_object_store(obj_type, keyspace, **kwargs)


def _get(store: Any, org_id: str, key: Optional[HierarchicalKey] = None) -> Any:
    """Helper to call store.get with or without key."""
    if key:
        return store.get(org_id, key)
    return store.get(org_id)


def _set(store: Any, org_id: str, data: Any, key: Optional[HierarchicalKey] = None) -> None:
    """Helper to call store.set with or without key."""
    if key:
        store.set(org_id, key, data)
    else:
        store.set(org_id, data)


def _update(store: Any, org_id: str, handler: Callable[[Any], Any], key: Optional[HierarchicalKey] = None) -> Any:
    """Helper to call store.update with or without key."""
    if key:
        return store.update(org_id, key, handler)
    return store.update(org_id, handler)


def _delete(store: Any, org_id: str, key: Optional[HierarchicalKey] = None) -> None:
    """Helper to call store.delete with or without key."""
    if key:
        store.delete(org_id, key)
    else:
        store.delete(org_id)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('data,expected_result', [({}, None), *TEST_DATA])
def test_object_store_get_should_coerce_result(context, key, data, expected_result):
    context.table.get_item.return_value = {'Item': {'data': data}} if data else {}
    store = _create_store(DataObject, context.keyspace, key)
    result = _get(store, context.org_id, key)

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), _key2(), None])
def test_object_store_get_should_call_get_item_with_correct_key(context, key):
    context.table.get_item.return_value = {}
    store = _create_store(DataObject, context.keyspace, key)
    _get(store, context.org_id, key)

    expected_key = {
        'PK': _expected_pk(context.org_id, context.keyspace),
        'SK': key.expected_sk() if key else 'record',
    }
    context.table.get_item.assert_called_once_with(Key=expected_key, ConsistentRead=True)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), _key2(), None])
@pytest.mark.parametrize('serializable_data,object_data', TEST_DATA)
def test_object_store_set_should_call_put_item_with_correct_item(context, key, serializable_data, object_data):
    context.mock_dyn.put_item.return_value = {}
    store = _create_store(DataObject, context.keyspace, key)
    _set(store, context.org_id, object_data, key)

    def validate_item(item):
        expected_pk = _expected_pk(context.org_id, context.keyspace)
        expected_sk = key.expected_sk() if key else 'record'
        return item['PK'] == expected_pk and item['SK'] == expected_sk and item['data'] == serializable_data

    context.mock_dyn.put_item.assert_called_once_with(context.table, ReturnValues='ALL_OLD', Item=MATCH(validate_item))


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
def test_object_store_update_should_correctly_update_data(context, key, serializable_data, expected_object_data):
    context.table.get_item.return_value = {'Item': {'data': serializable_data}}
    context.mock_dyn.put_item.return_value = {}
    store = _create_store(DataObject, context.keyspace, key)

    new_val = 99
    new_object_data = expected_object_data._replace(val=new_val)
    update_handler = Mock(return_value=new_object_data)
    result = _update(store, context.org_id, update_handler, key)

    update_handler.assert_called_once_with(expected_object_data)
    assert result == new_object_data


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_handler_param', [({}, None)])
def test_object_store_update_should_handle_empty_object(context, key, serializable_data, expected_handler_param):
    context.table.get_item.return_value = {'Item': {'data': serializable_data}}
    context.mock_dyn.put_item.return_value = {}
    store = _create_store(DataObject, context.keyspace, key)

    update_handler = Mock(side_effect=lambda x: x)
    result = _update(store, context.org_id, update_handler, key)

    update_handler.assert_called_once_with(expected_handler_param)
    assert result is None


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('handler_return', [DataObject(val=1, name='new', data=['a']), None])
def test_object_store_update_should_handle_none_result(context, key, handler_return):
    context.table.get_item.return_value = {'Item': {'data': {}}}
    context.mock_dyn.put_item.return_value = {}
    store = _create_store(DataObject, context.keyspace, key)

    update_handler = Mock(return_value=handler_return)
    result = _update(store, context.org_id, update_handler, key)

    assert result == handler_return


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), _key2(), None])
def test_object_store_delete_should_call_delete_item_with_correct_key(context, key):
    context.table.delete_item.return_value = {}
    store = _create_store(DataObject, context.keyspace, key)
    _delete(store, context.org_id, key)

    expected_key = {
        'PK': _expected_pk(context.org_id, context.keyspace),
        'SK': key.expected_sk() if key else 'record',
    }
    context.table.delete_item.assert_called_once_with(Key=expected_key, ReturnValues='ALL_OLD')


@pytest.mark.unit
@pytest.mark.parametrize('ttl,s3_temp_prefix', [(None, 'tmp'), (timedelta(days=7), 'custom')])
def test_create_object_store_should_pass_connection_params(context, ttl, s3_temp_prefix):
    feature_data.create_object_store(DataObject, context.keyspace, key_type=Key1, ttl=ttl, s3_temp_prefix=s3_temp_prefix)

    context.mock_connect.assert_called_once_with(ttl=ttl, s3_temp_prefix=s3_temp_prefix)


@pytest.mark.unit
@pytest.mark.parametrize(
    'key,record',
    [
        (
            _key1(),
            Model1(
                a=1,
                b=1.2,
                c='abc',
                d={'x': 'X', 'y': 'Y'},
                e=['a', 'b'],
                f=datetime(2020, 3, 14, 17, 54, 57, 415748, UTC),
                g=Enum1.B,
                h=DataObject(val=1, name='foo', data=['one', 'two']),
                i=('three', 'four'),
            ),
        ),
        (
            None,
            Model1(
                a=1,
                b=1.2,
                c='abc',
                d={'x': 'X', 'y': 'Y'},
                e=['a', 'b'],
                f=datetime(2020, 3, 14, 17, 54, 57, 415748, UTC),
                g=Enum1.B,
                h=DataObject(val=1, name='foo', data=['one', 'two']),
                i=('three', 'four'),
            ),
        ),
    ],
)
def test_object_store_roundtrip(mock_table, mocker, key, record):
    conn = FeatureDataConnection(_table=mock_table, _s3_bucket=str(uuid4()), _ttl=None, _s3_temp_prefix='tmp')
    mocker.patch.object(feature_data, 'connect', return_value=conn)
    org_id = str(uuid4())
    keyspace = str(uuid4())
    if key:
        store = feature_data.create_object_store(Model1, keyspace, key_type=type(key))
        store.set(org_id, key, record)
        assert store.get(org_id, key) == record
    else:
        store = feature_data.create_object_store(Model1, keyspace)
        store.set(org_id, record)
        assert store.get(org_id) == record


@pytest.mark.unit
@pytest.mark.parametrize('key,record', [(_key2(), Model2(a='x', b='')), (None, Model2(a='x', b=''))])
def test_object_store_roundtrip_strings(mock_table, mocker, key, record):
    conn = FeatureDataConnection(_table=mock_table, _s3_bucket=str(uuid4()), _ttl=None, _s3_temp_prefix='tmp')
    mocker.patch.object(feature_data, 'connect', return_value=conn)
    org_id = str(uuid4())
    keyspace = str(uuid4())
    if key:
        store = feature_data.create_object_store(Model2, keyspace, key_type=type(key))
        store.set(org_id, key, record)
        assert store.get(org_id, key) == record
    else:
        store = feature_data.create_object_store(Model2, keyspace)
        store.set(org_id, record)
        assert store.get(org_id) == record


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize(
    'error,large_data,expected_s3_data',
    [
        (EXCEED_MAX_SIZE_ERROR, DataObject(val=1, name='test', data=['a']), {'val': 1, 'name': 'test', 'data': ['a']}),
        (
            EXCEED_NESTING_LEVELS_ERROR,
            DataObject(val=2, name='nested', data=['b', 'c']),
            {'val': 2, 'name': 'nested', 'data': ['b', 'c']},
        ),
        (HTTP_413_ERROR, DataObject(val=3, name='large', data=['x']), {'val': 3, 'name': 'large', 'data': ['x']}),
    ],
)
def test_object_store_set_should_offload_to_s3_on_size_errors(context, key, error, large_data, expected_s3_data):
    context.mock_dyn.put_item.side_effect = [error, {}]
    context.mock_s3.put_json_object.return_value = None

    store = _create_store(DataObject, context.keyspace, key)
    _set(store, context.org_id, large_data, key)

    assert context.mock_dyn.put_item.call_count == 2
    context.mock_s3.put_json_object.assert_called_once_with(
        context.conn._s3_bucket, MATCH(lambda path: path.startswith(S3_OBJECT_STORE_ROOT)), expected_s3_data
    )


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_get_should_return_s3_data_for_ref(context, key):
    s3_data = dict(val=42, name='from_s3', data=['a', 'b'])
    context.table.get_item.return_value = {'Item': {'data': 'REF:path/to/data'}}
    context.mock_s3.get_json_object.return_value = s3_data

    store = _create_store(DataObject, context.keyspace, key)
    result = _get(store, context.org_id, key)

    assert result == DataObject(**s3_data)
    context.mock_s3.get_json_object.assert_called_once_with(context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/path/to/data')


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_set_should_delete_prior_s3_data(context, key):
    context.table.get_item.return_value = {}
    context.mock_dyn.put_item.return_value = {'Attributes': {'data': 'REF:old/path'}}

    new_data = DataObject(val=99, name='new', data=['updated'])
    store = _create_store(DataObject, context.keyspace, key)
    _set(store, context.org_id, new_data, key)

    context.mock_aws_service.s3_client.delete_object.assert_called_once_with(
        Bucket=context.conn._s3_bucket, Key=f'{S3_OBJECT_STORE_ROOT}/old/path'
    )


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@freezegun.freeze_time('2025-01-15 12:00:00')
def test_object_store_set_with_ttl_should_write_ttl_attribute(context, key, mocker):
    ttl = timedelta(days=7)
    conn_with_ttl = FeatureDataConnection(
        _table=context.table, _s3_bucket=context.conn._s3_bucket, _ttl=ttl, _s3_temp_prefix='tmp'
    )
    mocker.patch.object(feature_data, 'connect', return_value=conn_with_ttl)
    context.mock_dyn.put_item.return_value = {}

    data = DataObject(val=1, name='test', data=['x'])
    store = _create_store(DataObject, context.keyspace, key, ttl=ttl)
    _set(store, context.org_id, data, key)

    def validate_item(item):
        expected_expiration = int((datetime(2025, 1, 15, 12, 0, 0) + ttl).timestamp())
        return 'ttl' in item and item['ttl'] == expected_expiration

    context.mock_dyn.put_item.assert_called_once_with(context.table, ReturnValues='ALL_OLD', Item=MATCH(validate_item))


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_get_should_return_data_from_temp_ref(context, key):
    temp_data = dict(val=5, name='temp', data=['temp'])
    context.table.get_item.return_value = {'Item': {'data': 'TEMP_REF:tmp/feature_data_object_store/path'}}
    context.mock_s3.get_json_object.return_value = temp_data

    store = _create_store(DataObject, context.keyspace, key)
    result = _get(store, context.org_id, key)

    assert result == DataObject(**temp_data)
    context.mock_s3.get_json_object.assert_called_once_with(context.conn._s3_bucket, 'tmp/feature_data_object_store/path')


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_get_should_return_none_when_temp_ref_is_missing(context, key):
    context.table.get_item.return_value = {'Item': {'data': 'TEMP_REF:tmp/feature_data_object_store/expired'}}
    context.mock_s3.get_json_object.side_effect = NOT_FOUND_ERROR

    store = _create_store(DataObject, context.keyspace, key)
    result = _get(store, context.org_id, key)

    assert result is None


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_update_should_retry_on_condition_check_failure(context, key):
    initial_data = dict(val=10, name='initial', data=['a'])
    updated_data = dict(val=20, name='updated', data=['b'])

    context.table.get_item.return_value = {'Item': {'data': initial_data}}
    context.mock_dyn.put_item.side_effect = [CONDITIONAL_CHECK_FAILED_ERROR, {}]

    store = _create_store(DataObject, context.keyspace, key)
    result = _update(store, context.org_id, lambda _: DataObject(**updated_data), key)

    assert context.table.get_item.call_count == 2
    assert context.mock_dyn.put_item.call_count == 2
    assert result == DataObject(**updated_data)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_delete_should_delete_prior_s3_data(context, key):
    context.table.delete_item.return_value = {'Attributes': {'data': 'REF:path/to/delete'}}

    store = _create_store(DataObject, context.keyspace, key)
    _delete(store, context.org_id, key)

    context.mock_aws_service.s3_client.delete_object.assert_called_once_with(
        Bucket=context.conn._s3_bucket, Key=f'{S3_OBJECT_STORE_ROOT}/path/to/delete'
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'key_args',
    [
        {'region': 'us-east-1'},
        {'region': 'us-east-1', 'account': '123'},
        {'region': 'us-east-1', 'account': '123', 'service': 's3'},
    ],
)
def test_hierarchical_key_allows_partial_key_if_fields_are_prefix(key_args):
    assert Key1(**key_args) is not None


@pytest.mark.unit
@pytest.mark.parametrize(
    'key_args',
    [
        {},
        {'account': '123'},
        {'service': 's3'},
        {'region': 'us-east-1', 'service': 's3'},
        {'account': '123', 'service': 's3'},
    ],
)
def test_hierarchical_key_disallows_partial_key_if_fields_are_not_prefix(key_args):
    with pytest.raises(ValueError, match='Fields provided must be a prefix'):
        Key1(**key_args)


@pytest.mark.unit
def test_hierarchical_key_requires_hkpart_fields(context):
    with pytest.raises(TypeError, match='Field invalid_field in InvalidKey must be of type HKPart'):

        class InvalidKey(HierarchicalKey):
            something: HKPart
            invalid_field: str


@pytest.mark.unit
@pytest.mark.parametrize('operation', ['get', 'set', 'update', 'delete'])
def test_object_store_operations_should_require_complete_key(context, operation):
    context.table.get_item.return_value = {'Item': {'data': dict(val=1, name='test', data=['a'])}}
    context.mock_dyn.put_item.return_value = {}
    context.table.delete_item.return_value = {}

    store = feature_data.create_object_store(DataObject, context.keyspace, key_type=Key1)
    partial_key = Key1(region='us-east-1')

    operation_dict = {
        'get': lambda: store.get(context.org_id, partial_key),
        'set': lambda: store.set(context.org_id, partial_key, DataObject(val=1, name='test', data=['a'])),
        'update': lambda: store.update(context.org_id, partial_key, lambda x: x),
        'delete': lambda: store.delete(context.org_id, partial_key),
    }

    with pytest.raises(ValueError, match='Complete key required for this operation'):
        operation_dict[operation]()


@pytest.mark.unit
@pytest.mark.parametrize(
    'partial_key,expected_sk_prefix',
    [
        (Key1(region='us-east-1'), 'record#us-east-1'),
        (Key1(region='us-east-1', account='123'), 'record#us-east-1#123'),
        (Key1(region='us-east-1', account='123', service='s3'), 'record#us-east-1#123#s3'),
        (Key2(env='prod'), 'record#prod'),
        (Key2(env='prod', some_enum=Enum2.Y), 'record#prod#y'),
    ],
)
def test_object_store_list_should_query_dynamo_with_correct_key(context, partial_key, expected_sk_prefix):
    context.mock_dyn.query.return_value = []

    store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(partial_key))
    list(store.list(context.org_id, partial_key))

    expected_pk = _expected_pk(context.org_id, context.keyspace)
    expected_condition = Key('PK').eq(expected_pk) & Key('SK').begins_with(expected_sk_prefix)
    context.mock_dyn.query.assert_called_once_with(
        context.conn._table, dict(KeyConditionExpression=expected_condition, ConsistentRead=True)
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'partial_key,dynamo_data,expected_result',
    [
        (
            Key1(region='us-east-1'),
            [
                {'PK': 'org#123', 'SK': 'record#us-east-1#123#s3', 'data': dict(val=1, name='foo', data=['a'])},
                {'PK': 'org#123', 'SK': 'record#us-east-1#456#ec2', 'data': dict(val=2, name='bar', data=['b'])},
            ],
            [DataObject(val=1, name='foo', data=['a']), DataObject(val=2, name='bar', data=['b'])],
        ),
        (Key1(region='us-west-2'), [], []),
        (
            Key2(env='prod'),
            [{'PK': 'org#456', 'SK': 'record#prod#X#item1', 'data': dict(val=99, name='test', data=['x', 'y'])}],
            [DataObject(val=99, name='test', data=['x', 'y'])],
        ),
    ],
)
def test_object_store_list_should_coerce_result(context, partial_key, dynamo_data, expected_result):
    context.mock_dyn.query.return_value = dynamo_data

    store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(partial_key))
    result = list(store.list(context.org_id, partial_key))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'partial_key,keyspace,expected_sk_prefix',
    [
        (Key1(region='us-east-1'), 'keyspace1', 'record#us-east-1'),
        (Key1(region='us-east-1', account='123'), 'keyspace2', 'record#us-east-1#123'),
    ],
)
def test_object_store_list_all_should_scan_dynamo_with_correct_key(context, partial_key, keyspace, expected_sk_prefix):
    context.mock_dyn.scan.return_value = []

    store = feature_data.create_object_store(DataObject, keyspace, key_type=type(partial_key))
    list(store.list_all(partial_key))

    expected_condition = (
        Attr('PK').begins_with('organization#')
        & Attr('PK').contains(f'#{keyspace}')
        & Attr('SK').begins_with(expected_sk_prefix)
    )
    context.mock_dyn.scan.assert_called_once_with(context.conn._table, dict(FilterExpression=expected_condition))


@pytest.mark.unit
@pytest.mark.parametrize(
    'partial_key,keyspace,dynamo_data,expected_result',
    [
        (
            Key1(region='us-east-1'),
            'keyspace1',
            [
                {
                    'PK': 'organization#org1#keyspace1',
                    'SK': 'record#us-east-1#123#s3',
                    'data': dict(val=1, name='foo', data=['a']),
                },
                {
                    'PK': 'organization#org2#keyspace1',
                    'SK': 'record#us-east-1#456#ec2',
                    'data': dict(val=2, name='bar', data=['b']),
                },
            ],
            [
                ListRecord(organization_id='org1', data=DataObject(val=1, name='foo', data=['a'])),
                ListRecord(organization_id='org2', data=DataObject(val=2, name='bar', data=['b'])),
            ],
        ),
        (Key1(region='us-west-2'), 'keyspace2', [], []),
        (
            Key2(env='prod', some_enum=Enum2.Y),
            'keyspace3',
            [
                {
                    'PK': 'organization#org5#keyspace3',
                    'SK': 'record#prod#Y#item1',
                    'data': dict(val=42, name='single', data=['test']),
                },
            ],
            [ListRecord(organization_id='org5', data=DataObject(val=42, name='single', data=['test']))],
        ),
    ],
)
def test_object_store_list_all_should_coerce_result(context, partial_key, keyspace, dynamo_data, expected_result):
    context.mock_dyn.scan.return_value = dynamo_data

    store = feature_data.create_object_store(DataObject, keyspace, key_type=type(partial_key))
    result = list(store.list_all(partial_key))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'partial_key,dynamo_data,mock_s3_data,expected_result',
    [
        (
            Key1(region='us-east-1'),
            [
                {'PK': 'org#123', 'SK': 'record#us-east-1#123#s3', 'data': 'REF:path/to/data1'},
                {'PK': 'org#123', 'SK': 'record#us-east-1#456#ec2', 'data': dict(val=2, name='bar', data=['b'])},
            ],
            {'path/to/data1': dict(val=1, name='foo', data=['a'])},
            [DataObject(val=1, name='foo', data=['a']), DataObject(val=2, name='bar', data=['b'])],
        ),
        (Key1(region='us-east-1'), [], {}, []),
    ],
)
def test_object_store_list_should_handle_s3_refs(context, partial_key, dynamo_data, mock_s3_data, expected_result):
    MockEx(context.mock_dyn.query).if_called_with_anything().return_value(dynamo_data)
    for path, data in mock_s3_data.items():
        MockEx(context.mock_s3.get_json_object).if_called_with(
            context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
        ).return_value(data)

    store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(partial_key))
    result = list(store.list(context.org_id, partial_key))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'partial_key,keyspace,dynamo_data,mock_s3_data,expected_result',
    [
        (
            Key1(region='us-east-1'),
            'keyspace1',
            [
                {'PK': 'organization#org1#keyspace1', 'SK': 'record#us-east-1#123#s3', 'data': 'REF:path/to/data1'},
                {
                    'PK': 'organization#org2#keyspace1',
                    'SK': 'record#us-east-1#456#ec2',
                    'data': dict(val=2, name='bar', data=['b']),
                },
            ],
            {'path/to/data1': dict(val=1, name='foo', data=['a'])},
            [
                ListRecord(organization_id='org1', data=DataObject(val=1, name='foo', data=['a'])),
                ListRecord(organization_id='org2', data=DataObject(val=2, name='bar', data=['b'])),
            ],
        ),
        (
            Key2(env='prod'),
            'keyspace2',
            [
                {'PK': 'organization#org3#keyspace2', 'SK': 'record#prod#x#item1', 'data': 'REF:s3/data2'},
                {'PK': 'organization#org4#keyspace2', 'SK': 'record#prod#y#item2', 'data': 'REF:s3/data3'},
            ],
            {
                's3/data2': dict(val=10, name='first', data=['x']),
                's3/data3': dict(val=20, name='second', data=['y']),
            },
            [
                ListRecord(organization_id='org3', data=DataObject(val=10, name='first', data=['x'])),
                ListRecord(organization_id='org4', data=DataObject(val=20, name='second', data=['y'])),
            ],
        ),
    ],
)
def test_object_store_list_all_should_handle_s3_refs(
    context, partial_key, keyspace, dynamo_data, mock_s3_data, expected_result
):
    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(dynamo_data)
    for path, data in mock_s3_data.items():
        MockEx(context.mock_s3.get_json_object).if_called_with(
            context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
        ).return_value(data)

    store = feature_data.create_object_store(DataObject, keyspace, key_type=type(partial_key))
    result = list(store.list_all(partial_key))

    assert result == expected_result


@pytest.mark.unit
def test_object_store_list_all_should_filter_out_wrong_keyspace(context):
    keyspace = 'correct_keyspace'
    dynamo_data = [
        {
            'PK': f'organization#org1#{keyspace}',
            'SK': 'record#us-east-1#123#s3',
            'data': dict(val=1, name='foo', data=['a']),
        },
        {
            'PK': 'organization#org2#wrong_keyspace',
            'SK': 'record#us-east-1#456#ec2',
            'data': dict(val=2, name='bar', data=['b']),
        },
    ]
    expected_result = [ListRecord(organization_id='org1', data=DataObject(val=1, name='foo', data=['a']))]

    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(dynamo_data)

    store = feature_data.create_object_store(DataObject, keyspace, key_type=Key1)
    result = list(store.list_all(Key1(region='us-east-1')))

    assert result == expected_result


# Migration Tests


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
def test_object_store_get_should_deserialize_json_string(context, key, serializable_data, expected_object_data):
    json_string = json.dumps(serializable_data)
    context.table.get_item.return_value = {'Item': {'data': json_string}}

    store = _create_store(DataObject, context.keyspace, key)
    result = _get(store, context.org_id, key)

    assert result == expected_object_data


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
@pytest.mark.parametrize('migration_mode', [MigrationMode.from_feature_data, MigrationMode.from_billing])
def test_object_store_get_with_migration_should_fallback_to_legacy(
    context, key, serializable_data, expected_object_data, migration_mode
):
    # New format doesn't exist, legacy format has data
    context.table.get_item.side_effect = [
        {},  # object_store format
        {'Item': {'data': serializable_data}},  # legacy format
    ]

    store = _create_store(DataObject, context.keyspace, key, migration_mode=migration_mode)
    result = _get(store, context.org_id, key)

    assert result == expected_object_data

    # Verify calls were made with correct keys in correct order
    legacy_key = (
        _expected_legacy_feature_data_key(context.org_id, context.keyspace, key)
        if migration_mode == MigrationMode.from_feature_data
        else _expected_legacy_billing_key(context.org_id, context.keyspace, key)
    )
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
        call(Key=legacy_key, ConsistentRead=True),
    ]


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
@pytest.mark.parametrize('migration_mode', [MigrationMode.from_feature_data, MigrationMode.from_billing])
def test_object_store_get_with_migration_prefers_new_format(
    context, key, serializable_data, expected_object_data, migration_mode
):
    legacy_data = dict(val=999, name='old', data=['legacy'])
    # New format exists, legacy also exists but shouldn't be used
    context.table.get_item.side_effect = [
        {'Item': {'data': serializable_data}},  # object_store format
        {'Item': {'data': legacy_data}},  # legacy format (shouldn't be accessed)
    ]

    store = _create_store(DataObject, context.keyspace, key, migration_mode=migration_mode)
    result = _get(store, context.org_id, key)

    assert result == expected_object_data
    # Should only call once since new format was found
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
    ]


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_update_with_migration_should_read_from_legacy_write_to_new(context, key):
    initial_data = dict(val=10, name='initial', data=['a'])
    updated_data = DataObject(val=20, name='updated', data=['b'])

    # New format doesn't exist, legacy does
    context.table.get_item.side_effect = [
        {},  # object_store format (for primary check)
        {'Item': {'data': initial_data}},  # legacy_feature_data format (for fallback)
    ]
    context.mock_dyn.put_item.return_value = {}

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.from_feature_data)
    result = _update(store, context.org_id, lambda _: updated_data, key)

    assert result == updated_data

    # Verify reads from correct keys in correct order
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
        call(Key=_expected_legacy_feature_data_key(context.org_id, context.keyspace, key), ConsistentRead=True),
    ]

    # Verify write to new format only
    context.mock_dyn.put_item.assert_called_once_with(context.table, ReturnValues='ALL_OLD', ConditionExpression=ANY, Item=ANY)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_update_with_migration_should_use_new_format_when_exists(context, key):
    new_data = dict(val=10, name='new', data=['a'])
    legacy_data = dict(val=999, name='legacy', data=['old'])
    updated_data = DataObject(val=20, name='updated', data=['b'])

    # New format exists
    context.table.get_item.side_effect = [
        {'Item': {'data': new_data}},  # object_store format (for primary check)
        {'Item': {'data': legacy_data}},  # legacy_feature_data format (should not be used)
    ]
    context.mock_dyn.put_item.return_value = {}

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.from_feature_data)
    result = _update(store, context.org_id, lambda x: updated_data, key)

    assert result == updated_data

    # Verify only primary key was checked (fallback not needed since primary had data)
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
    ]

    # Verify write to new format only
    context.mock_dyn.put_item.assert_called_once_with(context.table, ReturnValues='ALL_OLD', ConditionExpression=ANY, Item=ANY)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_delete_with_migration_should_delete_from_all_locations(context, key):
    context.table.delete_item.return_value = {}

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.from_feature_data)
    _delete(store, context.org_id, key)

    # Should delete from both locations in reverse order (legacy first, then new)
    assert context.table.delete_item.call_args_list == [
        call(Key=_expected_legacy_feature_data_key(context.org_id, context.keyspace, key), ReturnValues='ALL_OLD'),
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ReturnValues='ALL_OLD'),
    ]


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,object_data', TEST_DATA)
def test_object_store_set_always_writes_to_new_format(context, key, serializable_data, object_data):
    context.mock_dyn.put_item.return_value = {}

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.from_feature_data)
    _set(store, context.org_id, object_data, key)

    # Verify write happens only once to new format with correct data
    def validate_item(item):
        expected_pk = _expected_pk(context.org_id, context.keyspace)
        return item['PK'] == expected_pk and item['data'] == serializable_data

    context.mock_dyn.put_item.assert_called_once_with(context.table, ReturnValues='ALL_OLD', Item=MATCH(validate_item))


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
def test_object_store_get_should_deserialize_json_string_from_legacy_location(
    context, key, serializable_data, expected_object_data
):
    json_string = json.dumps(serializable_data)
    # New format doesn't exist, legacy has JSON string
    context.table.get_item.side_effect = [
        {},  # object_store format
        {'Item': {'data': json_string}},  # legacy_feature_data format with JSON string
    ]

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.from_feature_data)
    result = _get(store, context.org_id, key)

    assert result == expected_object_data
    # Verify both keys were checked in correct order
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
        call(Key=_expected_legacy_feature_data_key(context.org_id, context.keyspace, key), ConsistentRead=True),
    ]


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
def test_object_store_with_none_migration_mode_only_checks_new_format(context, key, serializable_data, expected_object_data):
    # Legacy data exists but should not be checked
    context.table.get_item.return_value = {}

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.none)
    result = _get(store, context.org_id, key)

    assert result is None
    # Should only check new format (no legacy fallback)
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
    ]


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_object_store_update_with_migration_handles_empty_fallback(context, key):
    updated_data = DataObject(val=1, name='new', data=['created'])

    # No data in any location
    context.table.get_item.return_value = {}
    context.mock_dyn.put_item.return_value = {}

    update_handler = Mock(return_value=updated_data)

    store = _create_store(DataObject, context.keyspace, key, migration_mode=MigrationMode.from_feature_data)
    result = _update(store, context.org_id, update_handler, key)

    # Handler should be called with None
    update_handler.assert_called_once_with(None)
    assert result == updated_data

    # Verify reads from correct keys in correct order
    assert context.table.get_item.call_args_list == [
        call(Key=_expected_object_store_key(context.org_id, context.keyspace, key), ConsistentRead=True),
        call(Key=_expected_legacy_feature_data_key(context.org_id, context.keyspace, key), ConsistentRead=True),
    ]

    # Verify write to new format
    context.mock_dyn.put_item.assert_called_once_with(context.table, ReturnValues='ALL_OLD', ConditionExpression=ANY, Item=ANY)


# Snapshot Tests


@pytest.mark.unit
@pytest.mark.parametrize(
    'snapshot_id,keyspaces', [('S-valid123', {'keyspace1', 'keyspace2'}), ('S-abcXYZ789', {'keyspaceA', 'keyspaceB'})]
)
def test_create_object_store_snapshot_should_copy_records_and_create_registry(context, snapshot_id, keyspaces):
    test_data_ks1 = [
        {'PK': f'organization#{context.org_id}#keyspace1', 'SK': 'record#key1', 'data': {'val': 1}},
        {'PK': f'organization#{context.org_id}#keyspace1', 'SK': 'record#key2', 'data': {'val': 2}},
    ]
    test_data_ks2 = [
        {'PK': f'organization#{context.org_id}#keyspace2', 'SK': 'record#keyA', 'data': {'val': 10}},
    ]

    context.mock_dyn.query.side_effect = [test_data_ks1, test_data_ks2]
    context.mock_dyn.put_item.return_value = {}

    feature_data.create_object_store_snapshot(context.org_id, snapshot_id, keyspaces)

    assert context.mock_dyn.put_item.call_count == 4

    def validate_snapshot_record(call_args):
        item = call_args[1]['Item']
        return item['PK'].startswith(snapshot_id) and SNAPSHOT_REGISTRY_KEY not in item['SK']

    snapshot_calls = [c for c in context.mock_dyn.put_item.call_args_list if validate_snapshot_record(c)]
    assert len(snapshot_calls) == 3

    def validate_registry_record(call_args):
        item = call_args[1]['Item']
        return SNAPSHOT_REGISTRY_KEY in item['SK']

    registry_calls = [c for c in context.mock_dyn.put_item.call_args_list if validate_registry_record(c)]
    assert len(registry_calls) == 1


@pytest.mark.unit
def test_create_object_store_snapshot_should_raise_on_empty_keyspaces(context):
    with pytest.raises(ValueError, match='At least one keyspace is required to create a snapshot'):
        feature_data.create_object_store_snapshot(context.org_id, 'S-test', set())


@pytest.mark.unit
@pytest.mark.parametrize('snapshot_id', ['S-valid123', 'S-abcXYZ789'])
def test_delete_object_store_snapshot_should_remove_all_records_and_registry(context, snapshot_id):
    registry_data = {'snapshot_id': snapshot_id, 'keyspaces': ['keyspace1']}

    context.table.get_item.return_value = {'Item': {'data': registry_data}}
    context.mock_dyn.query.return_value = [
        {'PK': f'{snapshot_id}#organization#{context.org_id}#keyspace1', 'SK': 'record#key1'},
        {'PK': f'{snapshot_id}#organization#{context.org_id}#keyspace1', 'SK': 'record#key2'},
    ]
    context.table.delete_item.return_value = {}

    feature_data.delete_object_store_snapshot(context.org_id, snapshot_id)

    assert context.table.delete_item.call_count == 3

    def validate_snapshot_delete(key):
        return snapshot_id in key['PK']

    snapshot_deletes = [c for c in context.table.delete_item.call_args_list if validate_snapshot_delete(c[1]['Key'])]
    assert len(snapshot_deletes) == 2

    def validate_registry_delete(key):
        return 'snapshot_registry' in key['SK']

    registry_deletes = [c for c in context.table.delete_item.call_args_list if validate_registry_delete(c[1]['Key'])]
    assert len(registry_deletes) == 1


@pytest.mark.unit
def test_delete_object_store_snapshot_should_raise_error_on_invalid_snapshot(context):
    context.table.get_item.return_value = {}

    with pytest.raises(feature_data.InvalidSnapshotIdError):
        feature_data.delete_object_store_snapshot(context.org_id, 'S-nonexistent')


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize('snapshot_id,keyspace', [('S-test123', 'keyspace1'), ('S-abcXYZ789', 'keyspace2')])
def test_snapshot_object_store_get_should_validate_and_read(context, key, snapshot_id, keyspace):
    registry_data = {'snapshot_id': snapshot_id, 'keyspaces': [keyspace]}
    serializable_data = {'val': 42, 'name': 'snapshot_data', 'data': ['a']}

    context.table.get_item.side_effect = [
        {'Item': {'data': registry_data}},
        {'Item': {'data': serializable_data}},
    ]

    store = _create_store(DataObject, keyspace, key, snapshot_id=snapshot_id)
    result = _get(store, context.org_id, key)

    assert result == DataObject(**serializable_data)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_snapshot_object_store_get_should_raise_error_on_invalid_snapshot(context, key):
    snapshot_id = 'S-invalid'
    context.table.get_item.return_value = {}

    store = _create_store(DataObject, context.keyspace, key, snapshot_id=snapshot_id)

    with pytest.raises(feature_data.InvalidSnapshotIdError):
        _get(store, context.org_id, key)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
def test_snapshot_object_store_get_should_raise_error_on_wrong_keyspace(context, key):
    snapshot_id = 'S-test123'
    registry_data = {'snapshot_id': snapshot_id, 'keyspaces': ['other_keyspace']}

    context.table.get_item.return_value = {'Item': {'data': registry_data}}

    store = _create_store(DataObject, context.keyspace, key, snapshot_id=snapshot_id)

    with pytest.raises(feature_data.InvalidSnapshotIdError, match='does not include keyspace'):
        _get(store, context.org_id, key)


@pytest.mark.unit
@pytest.mark.parametrize('key', [_key1(), None])
@pytest.mark.parametrize(
    'operation',
    [
        'set',
        'update',
        'delete',
    ],
)
def test_snapshot_object_store_write_operations_should_raise_readonly_error(context, key, operation):
    snapshot_id = 'S-test123'
    store = _create_store(DataObject, context.keyspace, key, snapshot_id=snapshot_id)

    operation_dict = {
        'set': lambda: _set(store, context.org_id, DataObject(val=1, name='test', data=['a']), key),
        'update': lambda: _update(store, context.org_id, lambda x: x, key),
        'delete': lambda: _delete(store, context.org_id, key),
    }

    with pytest.raises(feature_data.SnapshotReadOnlyError, match='Snapshots are read-only and cannot be modified'):
        operation_dict[operation]()
