"""DynamoDB connector.

Allows to configure and get a 'dynamodb' resource.
"""

import time
from typing import Any

import boto3
from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
from mypy_boto3_dynamodb.client import DynamoDBClient
from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource
from mypy_boto3_dynamodb.type_defs import BatchGetItemOutputTypeDef
from werkzeug.exceptions import InternalServerError

from vectororder import config


def get_dynamodb_resource() -> DynamoDBServiceResource:
    """Get DynamoDB resource.

    Returns:
        DynamoDBServiceResource: DynamoDB resource.
    """
    boto_session = boto3.Session(region_name=config.AWS_REGION)
    dynamodb = boto_session.resource("dynamodb")
    return dynamodb


class DDBClient:
    """Shared DynamoDB Resource."""

    _client: DynamoDBClient | None
    _region: str

    def __init__(self, region: str) -> None:
        """Init."""
        self._region = region
        self._client = None

    def _deferred_client_init(self) -> None:
        if self._client is None:
            self._client = boto3.session.Session().client(
                "dynamodb", region_name=self._region
            )

    def _get_client(self) -> DynamoDBClient:
        self._deferred_client_init()

        if not self._client:
            raise InternalServerError("DynamoDB client failed to initialize")

        return self._client

    def get_item(
        self,
        table_name: str,
        table_key: str,
        item_id: str,
        consistent_read: bool = False,
    ) -> dict[str, Any] | None:
        """Get an item by id."""
        result = self._get_client().get_item(
            TableName=table_name,
            Key=DDBClient.python_to_dynamo({table_key: item_id}),
            ConsistentRead=consistent_read,
        )
        if "Item" in result:
            return DDBClient.dynamo_to_python(result["Item"])
        return None

    def batch_get_items(
        self,
        table_name: str,
        table_key: str,
        item_ids: list[str],
        projection_expression: str,
        consistent_read: bool = False,
    ) -> list[dict[str, Any]]:
        """Get multiple items by id."""
        key_chunks = (
            item_ids[i : i + config.DDB_BATCH_SIZE]
            for i in range(0, len(item_ids), config.DDB_BATCH_SIZE)
        )
        all_items = []
        for keys_chunk in key_chunks:
            items_chunk = self._get_items(
                table_name,
                table_key,
                keys_chunk,
                projection_expression,
                consistent_read,
            )
            all_items.extend(items_chunk)
        return all_items

    def _get_items(
        self,
        table_name: str,
        table_key: str,
        item_ids: list[str],
        projection_expression: str,
        consistent_read: bool = False,
    ) -> list[dict[str, Any]]:
        items_batch = []
        processing_keys = list(item_ids) or None
        while processing_keys:
            result = self._get_client().batch_get_item(
                RequestItems={
                    table_name: {
                        "Keys": [
                            DDBClient.python_to_dynamo({table_key: item_id})
                            for item_id in processing_keys
                        ],
                        "ConsistentRead": consistent_read,
                        "ProjectionExpression": projection_expression,
                    }
                }
            )
            processing_keys = DDBClient.get_unprocessed_keys(
                result, table_name, table_key
            )
            result_items = [
                DDBClient.dynamo_to_python(item)
                for item in result["Responses"][table_name]
            ]
            items_batch.extend(result_items)
            if processing_keys:
                time.sleep(0.3)
        return items_batch

    @staticmethod
    def python_to_dynamo(python_object: dict[str, Any]) -> dict[str, Any]:
        """Convert python dict to dynamo syntax."""
        serializer = TypeSerializer()
        return {k: serializer.serialize(v) for k, v in python_object.items()}

    @staticmethod
    def dynamo_to_python(dynamo_object: dict[str, Any]) -> dict[str, Any]:
        """Convert dynamo syntax to python dict."""
        deserializer = TypeDeserializer()
        return {k: deserializer.deserialize(v) for k, v in dynamo_object.items()}

    @staticmethod
    def get_unprocessed_keys(
        result: BatchGetItemOutputTypeDef, table_name: str, table_key: str
    ) -> list[str] | None:
        """Get unprocessed keys."""
        if "UnprocessedKeys" not in result or not result["UnprocessedKeys"]:
            return None
        if (
            table_name not in result["UnprocessedKeys"]
            or not result["UnprocessedKeys"][table_name]
        ):
            return None
        if (
            "Keys" not in result["UnprocessedKeys"][table_name]
            or not result["UnprocessedKeys"][table_name]["Keys"]
        ):
            return None
        return [
            DDBClient.dynamo_to_python(item)[table_key]
            for item in result["UnprocessedKeys"][table_name]["Keys"]
        ]


ddb_connector = DDBClient(config.AWS_REGION)