# 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

from datetime import datetime, timedelta, timezone
from enum import Enum, StrEnum
from typing import NamedTuple
from unittest.mock import Mock
from uuid import uuid4

import freezegun
import pytest
from botocore.exceptions import ClientError
from pydantic import BaseModel

from czc import feature_data
from czc.feature_data import S3_OBJECT_STORE_ROOT, FeatureDataConnection, HierarchicalKey
from czc.type_util import deserializable
from czt.unittest import MATCH, Context

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()
    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,
    )
    with Context(mocker, feature_data, ['dyn', 's3', 'uuid', 'connect'], extra=extra) as context:
        context.mock_connect.return_value = context.conn
        context.mock_dyn.prefix.side_effect = lambda p, x: f'{p}#{x}'
        context.mock_dyn.compose_key.side_effect = lambda *x: '#'.join(x)
        context.mock_dyn.type_for_dynamo.side_effect = lambda x: x
        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: str = 'us-east-1'
    account: str = '123456789'

    def expected_sk(self) -> str:
        return f'record#{self.region}#{self.account}'


class Key2(HierarchicalKey):
    env: str = 'prod'
    some_enum: Enum2 = Enum2.X

    def expected_sk(self) -> str:
        return f'record#{self.env}#{self.some_enum.value}'


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}'


@pytest.mark.unit
@pytest.mark.parametrize('key', [Key1(), Key2(), 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 {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        result = store.get(context.org_id, key)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        result = store.get(context.org_id)

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('key', [Key1(), None])
def test_object_store_get_should_call_get_item_with_correct_key(context, key):
    context.table.get_item.return_value = {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        store.get(context.org_id, key)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        store.get(context.org_id)

    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 = {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        store.set(context.org_id, key, object_data)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        store.set(context.org_id, object_data)

    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(), Key2(), 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 = {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)

    new_val = 99
    new_object_data = expected_object_data._replace(val=new_val)
    update_handler = Mock(return_value=new_object_data)
    if key:
        result = store.update(context.org_id, key, update_handler)
    else:
        result = store.update(context.org_id, update_handler)

    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 = {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)

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

    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 = {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)

    update_handler = Mock(return_value=handler_return)
    if key:
        result = store.update(context.org_id, key, update_handler)
    else:
        result = store.update(context.org_id, update_handler)

    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 = {}
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)

    if key:
        store.delete(context.org_id, key)
    else:
        store.delete(context.org_id)

    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', [EXCEED_MAX_SIZE_ERROR, EXCEED_NESTING_LEVELS_ERROR, HTTP_413_ERROR])
def test_object_store_set_should_offload_to_s3_on_size_errors(context, key, error):
    large_data = DataObject(val=1, name='large', data=['x'] * 1000)
    context.mock_dyn.put_item.side_effect = [error, {}]
    context.mock_s3.put_json_object.return_value = None

    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        store.set(context.org_id, key, large_data)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        store.set(context.org_id, large_data)

    assert context.mock_dyn.put_item.call_count == 2
    context.mock_s3.put_json_object.assert_called_once()
    s3_call_args = context.mock_s3.put_json_object.call_args
    assert s3_call_args[0][0] == context.conn._s3_bucket
    assert s3_call_args[0][1].startswith(S3_OBJECT_STORE_ROOT)


@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

    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        result = store.get(context.org_id, key)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        result = store.get(context.org_id)

    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, mocker):
    mock_aws_service = mocker.patch.object(feature_data, 'aws_service', autospec=False)
    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'])
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        store.set(context.org_id, key, new_data)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        store.set(context.org_id, new_data)

    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'])
    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key), ttl=ttl)
        store.set(context.org_id, key, data)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace, ttl=ttl)
        store.set(context.org_id, data)

    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()
    call_args = context.mock_dyn.put_item.call_args
    assert validate_item(call_args[1]['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

    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        result = store.get(context.org_id, key)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        result = store.get(context.org_id)

    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

    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        result = store.get(context.org_id, key)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        result = store.get(context.org_id)

    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, {}]

    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        result = store.update(context.org_id, key, lambda _: DataObject(**updated_data))
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        result = store.update(context.org_id, lambda _: DataObject(**updated_data))

    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, mocker):
    mock_aws_service = mocker.patch.object(feature_data, 'aws_service', autospec=False)
    context.table.delete_item.return_value = {'Attributes': {'data': 'REF:path/to/delete'}}

    if key:
        store = feature_data.create_object_store(DataObject, context.keyspace, key_type=type(key))
        store.delete(context.org_id, key)
    else:
        store = feature_data.create_object_store(DataObject, context.keyspace)
        store.delete(context.org_id)

    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'
    )
