from aiobotocore.session import get_session
from apollo_utils.core.silent import SilentResult, silent
from apollo_utils.service.clients.aiohttp.clients.utils import retry
from botocore.exceptions import BotoCoreError
from typing import Any, Dict, List, Tuple, Type

from server import config
from server.publishers.clients.sns.utils import log
from server.publishers.parallel import parallel


async def _sns_publish_bulk(
    messages: List[Dict[str, Any]],
    chunk_size: int = config.PUBLISH_CHUNK_SIZE,
    retry_count: int = config.PUBLISH_RETRY_COUNT,
    retry_delay: int = config.PUBLISH_RETRY_DELAY,
    retried_exceptions: Tuple[Type[Exception]] = (BotoCoreError,),
) -> List[SilentResult]:
    session = get_session()
    async with session.create_client("sns", region_name=config.AWS_DEFAULT_REGION) as client:
        return await parallel(chunk_size=chunk_size, items_kwarg="PublishBatchRequestEntries")(
            silent(retry(retry_count, delay=retry_delay, excepted=retried_exceptions)(log(client.publish_batch)))
        )(TopicArn=config.MESSAGES_TOPIC_ARN, PublishBatchRequestEntries=messages)


def _parse_ok_sns_publish_bulk_result(result, id_to_error_map: Dict[int, dict], error_type: str = "publish_error"):
    error_data_keys = ("Code", "Message")
    failed = result.content.get("Failed", [])
    for response_item in failed:
        id_to_error_map[int(response_item["Id"])] = {
            "type": error_type,
            "data": {k.lower(): response_item.get(k) for k in error_data_keys},
        }


def _parse_error_sns_publish_bulk_result(result, id_to_error_map: Dict[int, dict], error_type: str = "publish_error"):
    error = {
        "type": error_type,
        "data": {
            "message": str(result.content),
        },
    }
    for request_item in result.kwargs["PublishBatchRequestEntries"]:
        id_to_error_map[int(request_item["Id"])] = error


def _parse_sns_publish_bulk_results(results: List[SilentResult], error_type: str = "publish_error") -> Dict[int, dict]:
    id_to_error_map = {}
    for result in results:
        parser = _parse_ok_sns_publish_bulk_result if result else _parse_error_sns_publish_bulk_result
        parser(result, id_to_error_map, error_type)
    return id_to_error_map


async def sns_publish_bulk(messages: List[Dict[str, Any]]) -> Dict[int, dict]:
    return _parse_sns_publish_bulk_results(await _sns_publish_bulk(messages))