"""DynamoDB adapter for db-test-helpers."""

from __future__ import annotations

from decimal import Decimal
from typing import Any
from uuid import uuid4

_ROOT_PK = "__root__"


class DynamodbAdapter:
    """DatabaseAdapter for Amazon DynamoDB via boto3.

    Args:
        resource: A ``boto3.resource('dynamodb')`` instance.
        pk_attr: Name of the partition key attribute (default ``"_pk"``).
        sk_attr: Name of the sort key attribute (default ``"_sk"``).
    """

    def __init__(self, resource: Any, *, pk_attr: str = "_pk", sk_attr: str = "_sk") -> None:
        self._resource = resource
        self._pk_attr = pk_attr
        self._sk_attr = sk_attr

    def _pk_value(self, parent_path: str) -> str:
        return parent_path if parent_path else _ROOT_PK

    def clear_group(self, group: str, parent_path: str) -> None:
        table = self._resource.Table(group)
        pk_value = self._pk_value(parent_path)

        kwargs: dict[str, Any] = {
            "KeyConditionExpression": "#pk = :pk",
            "ExpressionAttributeNames": {"#pk": self._pk_attr},
            "ExpressionAttributeValues": {":pk": pk_value},
            "ProjectionExpression": "#pk, #sk",
        }
        kwargs["ExpressionAttributeNames"]["#sk"] = self._sk_attr

        while True:
            response = table.query(**kwargs)
            with table.batch_writer() as batch:
                for item in response["Items"]:
                    batch.delete_item(Key={self._pk_attr: item[self._pk_attr], self._sk_attr: item[self._sk_attr]})
            if "LastEvaluatedKey" not in response:
                break
            kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]

    def write_record(self, group: str, parent_path: str, data: dict[str, Any]) -> dict[str, Any]:
        data = dict(data)
        doc_id = data.pop("__document_id__", None)
        if doc_id is None:
            doc_id = str(uuid4())

        item = _to_dynamodb_types(data)
        item[self._pk_attr] = self._pk_value(parent_path)
        item[self._sk_attr] = doc_id

        table = self._resource.Table(group)
        table.put_item(Item=item)

        item.pop(self._pk_attr)
        item.pop(self._sk_attr)
        return {"__document_id__": doc_id, **_from_dynamodb_types(item)}

    def read_records(self, group: str, parent_path: str) -> list[dict[str, Any]]:
        table = self._resource.Table(group)
        pk_value = self._pk_value(parent_path)

        items: list[dict[str, Any]] = []
        kwargs: dict[str, Any] = {
            "KeyConditionExpression": "#pk = :pk",
            "ExpressionAttributeNames": {"#pk": self._pk_attr},
            "ExpressionAttributeValues": {":pk": pk_value},
        }

        while True:
            response = table.query(**kwargs)
            items.extend(response["Items"])
            if "LastEvaluatedKey" not in response:
                break
            kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]

        result = []
        for item in items:
            doc_id = item.pop(self._sk_attr)
            item.pop(self._pk_attr)
            result.append({"__document_id__": doc_id, **_from_dynamodb_types(item)})
        return result

    def build_child_path(self, parent_path: str, group: str, record: dict[str, Any]) -> str:
        col_path = f"{parent_path}/{group}" if parent_path else group
        return f"{col_path}/{record['__document_id__']}"


def _to_dynamodb_types(data: Any) -> Any:
    if isinstance(data, dict):
        return {k: _to_dynamodb_types(v) for k, v in data.items()}
    if isinstance(data, list):
        return [_to_dynamodb_types(item) for item in data]
    if isinstance(data, float):
        return Decimal(str(data))
    return data


def _from_dynamodb_types(data: Any) -> Any:
    if isinstance(data, dict):
        return {k: _from_dynamodb_types(v) for k, v in data.items()}
    if isinstance(data, list):
        return [_from_dynamodb_types(item) for item in data]
    if isinstance(data, Decimal):
        if data == int(data):
            return int(data)
        return float(data)
    return data
