import asyncio
import operator
from collections.abc import Sequence
from functools import reduce
from typing import Any

from aiodynamo.client import Client
from aiodynamo.errors import ItemNotFound
from aiodynamo.expressions import F
from aiodynamo.models import BatchGetRequest, BatchWriteRequest
from aiodynamo.types import Item

from url_shortener.exceptions import DynamodbBatchWriteError, PathDoesNotExistsError


class PathsRepository:
    def __init__(
        self,
        dynamodb_client: Client,
        dynamodb_tablename: str,
        semaphore: asyncio.Semaphore,
        dynamodb_max_write_retry_count: int,
    ) -> None:
        self.dynamodb_client = dynamodb_client
        self.dynamodb_tablename = dynamodb_tablename
        self.semaphore = semaphore
        self.dynamodb_max_write_retry_count = dynamodb_max_write_retry_count

    async def put_url(
        self,
        url: str,
        path: str,
        is_personalized: bool,
        created_at: float,
        additional_attributes: dict[str, Any] | None = None,
    ) -> None:
        await self.dynamodb_client.put_item(
            self.dynamodb_tablename,
            {
                "path": path,
                "is_personalized": is_personalized,
                "url": url,
                "created_at": created_at,
                "additional_attributes": additional_attributes,
            },
        )

    async def put_urls(
        self,
        url: str,
        created_at: float,
        paths: Sequence[str],
        additional_attributes: list[dict[str, Any]] | None = None,
    ) -> None:
        if additional_attributes:
            items_to_put = [
                {
                    "path": path,
                    "url": url,
                    "additional_attributes": url_additional_attributes,
                    "created_at": created_at,
                    "is_personalized": True,
                }
                for path, url_additional_attributes in zip(
                    paths, additional_attributes, strict=True
                )
            ]
        else:
            items_to_put = [
                {
                    "path": path,
                    "url": url,
                    "created_at": created_at,
                    "additional_attributes": None,
                    "is_personalized": True,
                }
                for path in paths
            ]

        await self._write_data(items_to_put=items_to_put)

    async def get_url(self, path: str) -> Item | None:
        try:
            response = await self.dynamodb_client.get_item(
                self.dynamodb_tablename, {"path": path}
            )
            return response
        except ItemNotFound:
            return None

    async def get_urls(self, paths: list[str]) -> list[Item]:
        response = await self.dynamodb_client.batch_get(
            {
                self.dynamodb_tablename: BatchGetRequest(
                    [{"path": path} for path in paths]
                )
            }
        )
        return response.items[self.dynamodb_tablename]

    async def delete_paths(self, paths: list[str]) -> None:
        await self._write_data(keys_to_delete=[{"path": path} for path in paths])

    async def update_path(self, path: str, update_attributes: dict[str, Any]) -> None:
        if not update_attributes:
            return

        if not await self.get_url(path):
            raise PathDoesNotExistsError

        update_expr = reduce(
            operator.and_,
            (F(attr).set(value) for attr, value in update_attributes.items()),
        )

        await self.dynamodb_client.update_item(
            table=self.dynamodb_tablename,
            key={"path": path},
            update_expression=update_expr,
        )

    async def is_path_available(self, path: str) -> bool:
        return not await self.get_url(path)

    async def _write_data(
        self,
        items_to_put: list[Item] | None = None,
        keys_to_delete: list[Item] | None = None,
        retry_number=0,
        backoff_delay=1,
    ) -> None:
        if retry_number >= self.dynamodb_max_write_retry_count:
            raise DynamodbBatchWriteError

        async with self.semaphore:
            result = await self.dynamodb_client.batch_write(
                {
                    self.dynamodb_tablename: BatchWriteRequest(
                        items_to_put=items_to_put, keys_to_delete=keys_to_delete
                    )
                }
            )

        unput_items = None
        undeleted_keys = None

        if result.get(self.dynamodb_tablename):
            unput_items = result[self.dynamodb_tablename].unput_items or None
            undeleted_keys = result[self.dynamodb_tablename].undeleted_keys or None

        if unput_items or undeleted_keys:
            await asyncio.sleep(backoff_delay)
            await self._write_data(
                items_to_put=unput_items,
                keys_to_delete=undeleted_keys,
                retry_number=retry_number + 1,
                backoff_delay=backoff_delay + 1,
            )