"""Utility Functions for DynamoDB."""

import base64
import json
import logging
import uuid
from asyncio import sleep
from types import ModuleType
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
from botocore.config import Config
from ddtrace.trace import tracer
from pydantic import BaseModel

from pdp import config
from pdp.constants import dynamo as constants
from pdp.utils import backoff

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from mypy_boto3_dynamodb import DynamoDBClient


def get_opts(config: ModuleType) -> Dict[Any, Any]:
    """Prepare options for dynamodb."""
    print("config type", type(config))
    dynamodb_config = {"config": get_config(config=config)}

    endpoint_url = config.DYNAMODB_ENDPOINT_URL
    if endpoint_url:
        dynamodb_config.update({"endpoint_url": endpoint_url})

    return dynamodb_config


def get_config(config: ModuleType) -> Config:
    """Get botocore configuration."""
    logger.info(f"[botoconfig] tcp_keepalive: {config.DYNAMODB_TCP_KEEP_ALIVE}")

    # PP-351: Add retries config here.
    return Config(tcp_keepalive=config.DYNAMODB_TCP_KEEP_ALIVE)


def encode_pagination_cursor(dynamo_key: Dict[Any, Any]) -> str:
    """Convert a dynamo LastEvaluatedKey dict into a string.

    Args:
        dynamo_key (dict): The primary key of the item where a Dynamodb
            operation stopped, inclusive of the previous result set.
            Use this value to start a new operation,
            excluding this value in the new request.
    Returns:
        str: A string representation of the pagination cursor
    """
    return base64.b64encode(json.dumps(dynamo_key).encode("utf-8")).decode("utf-8")


def encode_pagination_shorthand(
    dynamo_key: Dict[Any, Any], hash_key: str, range_key: Optional[str] = None
) -> Union[str, Any]:
    """Convert a dynamo LastEvaluatedKey dict into a human-readable shorthand string."""
    deserialized_key = deserialize_dynamo_item_to_dict(dynamo_key)
    shorthand = deserialized_key[hash_key]
    if range_key:
        shorthand = f"{shorthand}#{deserialized_key[range_key]}"
    return shorthand


def decode_pagination_cursor(encoded_dynamo_key: str) -> Any:
    """Convert an encoded dynamo key string into a dict representing ExclusiveStartKey.

    Args:
        encoded_dynamo_key (str): The string-encoded primary key of the first item
            that an Dynamodb operation will evaluate. Use the value that was
            returned for LastEvaluatedKey in the previous operation.
    """
    return json.loads(base64.b64decode(encoded_dynamo_key).decode("utf-8"))


def deserialize_dynamo_item_to_dict(dynamo_item: Dict[Any, Any]) -> Dict[Any, Any]:
    """Deserialize an object with DynamoDB types, to Python types.

    Args:
        dynamo_item: An object with DynamoDB types

            Example:
            {
                "tenant_type": {"S": "account"},
                "tenant_uuid": {"S": "b87b9586-03dc-47be-a3b0-53ca84aa4145"},
                "version": {"S": "1"},
            }
    Returns:
        result: The pythonic translation of the object

            Example:
            {
                "tenant_type": "account",
                "tenant_uuid": "p87b9590-03dc-47be-a3c9-53ca84aa4155",
                "version": 1.0,
            }
    """
    result = {}
    deserializer = TypeDeserializer()
    for key, val in dynamo_item.items():
        deserialized_val = deserializer.deserialize(val)
        result[key] = deserialized_val
    return result


def deserialize_dynamo_item_array(dynamo_items: List[Dict[Any, Any]]) -> List[Any]:
    """Deserialize a list of DynamoDB data types to Python types.

    Args:
        dynamo_items: A list of DynamoDB items

    Returns:
        deserialized_items: A list of deserialized, pythonic dicts
    """
    deserialized_items = []
    for item in dynamo_items:
        deserialized_items.append(deserialize_dynamo_item_to_dict(item))

    return deserialized_items


def serialize_to_dynamo_item(python_item: Any) -> Any:
    """Serialize a python type into a  DynamoDB data type.

    Args:
        python_item: Some python type

    Returns:
        dynamo_items: A DynamoDB item
    """
    serializer = TypeSerializer()
    if isinstance(python_item, uuid.UUID):
        python_item = str(python_item)
    return serializer.serialize(python_item)


def serialize_dict_to_dynamo_item(row: Dict[str, Any]) -> Dict[str, Any]:
    """Convert a dict to a DynamoDB item.

    Args:
        row (dict): Row to convert
    Returns:
        dict: DynamoDB item
    """
    return {k: serialize_to_dynamo_item(v) for k, v in row.items()}


@tracer.wrap()
def count_batch_write_item_request_items(request_items: Dict[str, Any]) -> int:
    """Return the number of items in BatchWriteRequest RequestItems obj."""
    total = 0
    for table_request_items in request_items.values():
        total += len(table_request_items)
    return total


@tracer.wrap()
async def batch_write_item(
    client: "DynamoDBClient",
    request_items: Dict[str, Any],
) -> Dict[str, Any]:
    """Calls client.batch_write_item for a batch size of 25.

    batch_write_item will assert that the batch size is less than or equal
    to the allowed DYNAMODB_BATCH_WRITE_ITEM_SIZE, an AWS API constraint.

    Call client.batch_write_item for the batch until there are no UnprocessedItems,
    or the DYNAMODB_BATCH_UNPROCESSED_ITEMS_MAX_ATTEMPTS has been reached.

    Returns any unprocessed items.
    """
    total_items = count_batch_write_item_request_items(request_items)
    assert total_items > 0, "At least one request item must be provided."
    assert total_items <= constants.DYNAMODB_BATCH_WRITE_ITEM_SIZE, (
        f"No more than {constants.DYNAMODB_BATCH_WRITE_ITEM_SIZE} request items can be submitted."  # noqa: E501
    )

    unprocessed_items_attempts = 0
    response = client.batch_write_item(RequestItems=request_items)
    while (
        response[constants.DYNAMODB_BATCH_WRITE_ITEM_RESPONSE_KEY_UNPROCESSED_ITEMS]
        and unprocessed_items_attempts
        < config.DYNAMODB_BATCH_UNPROCESSED_ITEMS_MAX_ATTEMPTS
    ):
        unprocessed_items = response[
            constants.DYNAMODB_BATCH_WRITE_ITEM_RESPONSE_KEY_UNPROCESSED_ITEMS
        ]
        unprocessed_items_attempts += 1
        logger.info(
            "There are %d UnprocessedItems. Attempt #%d of %d to batch_write_item",
            len(unprocessed_items),
            unprocessed_items_attempts,
            config.DYNAMODB_BATCH_UNPROCESSED_ITEMS_MAX_ATTEMPTS,
        )
        await sleep(
            backoff.get_backoff_with_full_jitter(1, 3, unprocessed_items_attempts)
        )
        response = client.batch_write_item(RequestItems=unprocessed_items)

    return response[constants.DYNAMODB_BATCH_WRITE_ITEM_RESPONSE_KEY_UNPROCESSED_ITEMS]


class DynamoDbBatchWriteItemKey(BaseModel):
    """Representation of keys to use to create a request to BatchWriteItem API."""

    request_type: str
    subelement: str


def get_batch_write_item_keys(is_delete: bool) -> DynamoDbBatchWriteItemKey:
    """Return request payload keys for BatchWriteItem API depending on operation."""
    if is_delete:
        return DynamoDbBatchWriteItemKey(
            request_type="DeleteRequest",
            subelement="Key",
        )
    return DynamoDbBatchWriteItemKey(
        request_type="PutRequest",
        subelement="Item",
    )