"""Helpers for SQS queue resolution and async fan-out used by run_renew."""

import asyncio
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Any

import aioboto3
from aiobotocore.config import AioConfig
from vector_utils.aws_utils import sqs

from src import config

# Cached across daemon iterations; queue URLs are stable for the process lifetime.
_queue_url_cache: dict[str, str] = {}


def get_queue_name_for_delivery_job(
    encoder_id: int,
    dms_id: int,
    priority: int,
    encoding_order_priority: int,
) -> str:
    """Get the SQS queue name for a delivery job.

    Args:
        encoder_id: The encoder ID.
        dms_id: The DMS store ID.
        priority: The DMS store priority (cmm.priority).
        encoding_order_priority: The encoding order priority (jp.priority).

    Returns:
        The name of the SQS queue for the delivery job.
    """
    if encoder_id == config.PHYSICAL_ENCODER:
        if dms_id in config.GRAS_DELIVERY_STORE_IDS:
            queue_name: str = sqs.GRAS_DELIVERY_QUEUE_NAME_PATTERN.format(
                env=config.ENVIRONMENT,
            )
            return queue_name
    else:
        if priority > config.NUM_TOP_PRIORITY_STORES:
            priority = config.NON_TOP_PRIORITY_VALUE
    queue_name = sqs.format_queue_name(
        sqs.ENCODING_QUEUE_NAME_PATTERN,
        encoder_id=encoder_id,
        dms_priority=priority,
        priority=encoding_order_priority,
        env=config.ENVIRONMENT,
    )
    return queue_name


async def _get_or_create_queue_url(sqs_client: Any, queue_name: str) -> str:
    """Return the URL for `queue_name`, creating the queue if missing."""
    if queue_name in _queue_url_cache:
        return _queue_url_cache[queue_name]
    try:
        resp = await sqs_client.get_queue_url(QueueName=queue_name)
    except sqs_client.exceptions.QueueDoesNotExist:
        resp = await sqs_client.create_queue(QueueName=queue_name)
    url = str(resp["QueueUrl"])
    _queue_url_cache[queue_name] = url
    return url


async def _drain_with_workers[T, R](
    items: Sequence[T], handler: Callable[[T], Awaitable[R]]
) -> list[R]:
    """Map `handler` over `items`, bounded to SQS_WORKER_COUNT concurrent calls."""
    work_q: asyncio.Queue[T] = asyncio.Queue()
    for item in items:
        work_q.put_nowait(item)

    async def worker() -> list[R]:
        out: list[R] = []
        while not work_q.empty():
            out.append(await handler(work_q.get_nowait()))
        return out

    grouped = await asyncio.gather(*[worker() for _ in range(config.SQS_WORKER_COUNT)])
    return [result for worker_results in grouped for result in worker_results]


async def _resolve_queue_urls(
    sqs_client: Any, queue_names: Sequence[str], log: Any
) -> tuple[dict[str, str], list[str]]:
    """Resolve queue names to URLs via the worker pool.

    Returns the resolved name→URL map and the names that could not be resolved.
    """

    async def resolve(name: str) -> tuple[str, str | None]:
        try:
            return name, await _get_or_create_queue_url(sqs_client, name)
        except Exception:
            log.exception(f"could not resolve queue url for {name}")
            return name, None

    results = await _drain_with_workers(queue_names, resolve)
    resolved = {name: url for name, url in results if url is not None}
    unresolved = [name for name, url in results if url is None]
    return resolved, unresolved


def _build_work_units(
    messages_by_queue: Mapping[str, Sequence[Mapping[str, Any]]],
    queue_urls: Mapping[str, str],
) -> list[tuple[str, list[dict[str, str]]]]:
    """Split each resolved queue's messages into (queue_url, entries) send batches."""
    work_units: list[tuple[str, list[dict[str, str]]]] = []
    for queue_name, queue_url in queue_urls.items():
        messages = messages_by_queue[queue_name]
        for batch_start in range(0, len(messages), config.SQS_BATCH_SIZE):
            sqs_batch = messages[batch_start : batch_start + config.SQS_BATCH_SIZE]
            entries = [
                {
                    "Id": str(m["encoding_queue_detail_id"]),
                    "MessageBody": json.dumps(dict(m)),
                }
                for m in sqs_batch
            ]
            work_units.append((queue_url, entries))
    return work_units


async def _send_batch(
    sqs_client: Any, queue_url: str, entries: list[dict[str, str]], log: Any
) -> list[int]:
    """Send one entries batch and return the encoding_queue_detail_ids that failed."""
    try:
        resp = await sqs_client.send_message_batch(QueueUrl=queue_url, Entries=entries)
    except Exception:
        log.exception(f"send_message_batch raised for {queue_url}")
        return [int(e["Id"]) for e in entries]
    if not resp.get("Failed"):
        return []
    failed = [int(f["Id"]) for f in resp["Failed"]]
    log.error(f"partial SQS failure ({len(failed)}/{len(entries)}): {resp['Failed']}")
    return failed


async def fan_out_to_sqs(
    messages_by_queue: Mapping[str, Sequence[Mapping[str, Any]]],
    promoted_ids: Sequence[int],
    log: Any,
) -> list[int]:
    """Send `messages_by_queue` to SQS in parallel and return the failed eqd_ids."""
    client_config = AioConfig(max_pool_connections=config.SQS_WORKER_COUNT)
    try:
        session = aioboto3.Session()
        async with session.client(
            "sqs", region_name="us-east-1", config=client_config
        ) as sqs_client:
            # 1. Resolve queue names → URLs (creating queues as needed).
            queue_urls, unresolved = await _resolve_queue_urls(
                sqs_client, list(messages_by_queue), log
            )
            # 2. Build the flat list of (queue_url, entries) batches to send.
            work_units = _build_work_units(messages_by_queue, queue_urls)
            log.debug(
                f"resolved {len(queue_urls)}/{len(messages_by_queue)} queues; "
                f"{len(work_units)} batches to send"
            )

            # 3. Send every batch via the same bounded worker pool.
            async def send(unit: tuple[str, list[dict[str, str]]]) -> list[int]:
                queue_url, entries = unit
                return await _send_batch(sqs_client, queue_url, entries, log)

            batch_failures = await _drain_with_workers(work_units, send)
    except Exception:
        log.exception("SQS fan-out raised; all promoted rows will be reset")
        return list(promoted_ids)
    # 4. Failed ids = ids in unresolved queues + ids from failed/partial sends.
    failed = [
        int(m["encoding_queue_detail_id"])
        for name in unresolved
        for m in messages_by_queue[name]
    ]
    for batch_failed in batch_failures:
        failed.extend(batch_failed)
    return failed