"""Tests for the DynamoDB connector."""

from typing import Any
from unittest.mock import Mock

import boto3
from moto import mock_aws
from mypy_boto3_dynamodb.client import DynamoDBClient

from vectororder import config
from vectororder.connectors import dynamodb

TEST_TABLE_NAME = "test-table"
TABLE_PKEY = "order_id"


def test_get_dynamodb_resource(mocker: Mock) -> None:
    """Test DynamoDB resource retrieval."""
    mock_session = mocker.Mock()
    mock_dynamodb_resource = mocker.Mock()
    mock_session.resource.return_value = mock_dynamodb_resource

    mock_boto3 = mocker.patch("vectororder.connectors.dynamodb.boto3")
    mock_boto3.Session.return_value = mock_session

    assert dynamodb.get_dynamodb_resource() == mock_dynamodb_resource

    assert mock_boto3.Session.call_count == 1
    mock_boto3.Session.assert_called_with(region_name=config.AWS_REGION)
    assert mock_session.resource.call_count == 1
    mock_session.resource.assert_called_with("dynamodb")


def get_mock_orders(num_items: int = 1) -> list[dict[str, str]]:
    """Return fake order dynamo data."""
    return [
        {
            "order_id": str(i),
            "encoder_id": f"value_{i}",
        }
        for i in range(1, num_items + 1)
    ]


def mock_ddb_table(client: DynamoDBClient, items: list[dict[str, Any]]) -> None:
    """Mock DDB Table."""
    client.create_table(
        TableName=TEST_TABLE_NAME,
        KeySchema=[{"AttributeName": TABLE_PKEY, "KeyType": "HASH"}],
        AttributeDefinitions=[{"AttributeName": TABLE_PKEY, "AttributeType": "S"}],
        ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
    )
    for item in items:
        client.put_item(
            TableName=TEST_TABLE_NAME,
            Item={
                "order_id": {"S": item[TABLE_PKEY]},
                "encoder_id": {"S": item["encoder_id"]},
            },
        )


def test_get_order_from_dynamo() -> None:
    """Test retrieving orders from DynamoDB."""
    with mock_aws():
        client = boto3.client("dynamodb")
        mock_items = get_mock_orders()
        mock_ddb_table(client, mock_items)
        order_response = dynamodb.ddb_connector.get_item(
            TEST_TABLE_NAME, TABLE_PKEY, "1"
        )
        assert order_response == mock_items[0]
        order_response = dynamodb.ddb_connector.get_item(
            TEST_TABLE_NAME, TABLE_PKEY, "2"
        )
        assert order_response is None


def test_get_dynamodb_items(mocker: Mock) -> None:
    """Test retrieving DynamoDB items in batches."""
    mock_items = get_mock_orders(3)
    keys = [str(v[TABLE_PKEY]) for v in mock_items]
    mocked_config = mocker.patch("vectororder.models.orders.config")
    mocked_config.DDB_BATCH_SIZE = 2
    with mock_aws():
        mocked_sleep = mocker.patch("vectororder.connectors.dynamodb.time.sleep")
        mocked_get_unprocessed_keys = mocker.patch(
            "vectororder.connectors.dynamodb.DDBClient.get_unprocessed_keys"
        )
        # Pretend we've received non-empty UnprocessedKeys on the first call.
        mocked_get_unprocessed_keys.side_effect = [["3"], []]

        client = boto3.client("dynamodb")
        mock_ddb_table(client, mock_items)
        result = dynamodb.ddb_connector.batch_get_items(
            TEST_TABLE_NAME, TABLE_PKEY, keys, "order_id,encoder_id"
        )
    assert result == [
        {"encoder_id": "value_1", TABLE_PKEY: "1"},
        {"encoder_id": "value_2", TABLE_PKEY: "2"},
        {"encoder_id": "value_3", TABLE_PKEY: "3"},
        {"encoder_id": "value_3", TABLE_PKEY: "3"},
    ]
    mocked_sleep.call_count = 1
    mocked_sleep.assert_called_with(0.3)


def test_get_unprocessed_keys() -> None:
    """Test get_unprocessed_keys."""
    actual = dynamodb.DDBClient.get_unprocessed_keys(
        {
            "UnprocessedKeys": {},
            "Responses": {},
            "ConsumedCapacity": [],
            "ResponseMetadata": {
                "RequestId": "",
                "HTTPStatusCode": 200,
                "HTTPHeaders": {},
                "RetryAttempts": 0,
                "HostId": "",
            },
        },
        "foo",
        "bar",
    )
    assert actual is None
    actual = dynamodb.DDBClient.get_unprocessed_keys(
        {
            "UnprocessedKeys": {"foo": {"Keys": []}},
            "Responses": {},
            "ConsumedCapacity": [],
            "ResponseMetadata": {
                "RequestId": "",
                "HTTPStatusCode": 200,
                "HTTPHeaders": {},
                "RetryAttempts": 0,
                "HostId": "",
            },
        },
        "foo",
        "bar",
    )
    assert actual is None
    actual = dynamodb.DDBClient.get_unprocessed_keys(
        {
            "UnprocessedKeys": {
                "foo": {
                    "Keys": [
                        {"bar": {"S": "1"}},
                        {"bar": {"S": "2"}},
                        {"bar": {"S": "3"}},
                    ]
                }
            },
            "Responses": {},
            "ConsumedCapacity": [],
            "ResponseMetadata": {
                "RequestId": "",
                "HTTPStatusCode": 200,
                "HTTPHeaders": {},
                "RetryAttempts": 0,
                "HostId": "",
            },
        },
        "foo",
        "bar",
    )
    assert actual == ["1", "2", "3"]