"""Renew Manager module."""

import asyncio
import time
import uuid
from collections import defaultdict
from typing import Any

from lambdacommon import util
from vector_utils.datadog.metrics import call_datadog_with_metric

from src import config, queries
from src.logger import get_current_logger
from src.sqs_helpers import fan_out_to_sqs, get_queue_name_for_delivery_job
from src.task_protection import with_task_protection


def run_renew(log: Any) -> None:
    """Promote renewable jobs to 'ready_to_encode' and fan them out to SQS."""
    asyncio.run(_run_renew_async(log))


async def _run_renew_async(log: Any) -> None:
    """Async body of run_renew; uses aioboto3 for the SQS fan-out."""
    with util.dd_connection(config.DD_MYSQL_CONN_INFO) as conn:
        with conn.cursor() as cursor:
            cursor.execute(queries.SQL_GET_STORE_CAPS, (config.MAX_JOBS_PER_RUN,))
            caps = cursor.fetchall()
        log.info(f"Store caps fetched: {len(caps)} stores")

        if not caps:
            return

        dms_ids = tuple(int(row["dms_master_master_id"]) for row in caps)

        _emit_per_dms_metric("renew.attempt", dms_ids)

        case_params: tuple[int, ...] = tuple(
            v
            for row in caps
            for v in (int(row["dms_master_master_id"]), int(row["jobs_to_queue"]))
        )
        ranked_jobs_sql = queries.SQL_GET_RANKED_JOBS.format(
            case_body=" ".join(["WHEN %s THEN %s"] * len(caps))
        )

        with conn.cursor() as cursor:
            cursor.execute(ranked_jobs_sql, (dms_ids, *case_params))
            jobs = cursor.fetchall()
        log.info(f"Ranked jobs fetched: {len(jobs)} jobs")

        if not jobs:
            _emit_per_dms_metric("renew.run", dms_ids)
            return

        # Jobs can span many (store, encoder, jp.priority) combinations and
        # therefore many SQS queues — group by destination queue before batching.
        messages_by_queue: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for job in jobs:
            cmm_priority = int(job["store_priority"])
            delivery_job = {
                "upc": job["upc"],
                "cd": job["cd"],
                "track_id": job["track_id"],
                "clip_number": job["clip_number"],
                "order_type": job["encoding_order_type"],
                "dms_master_master_id": job["dms_master_master_id"],
                "encoding_queue_detail_id": job["encoding_queue_detail_id"],
                "encoder_id": job["encoder_id"],
                "meta_update": job["meta_update"] == "Y",
                "encoding_order_priority": job["encoding_order_priority"],
                "priority": cmm_priority,
            }
            queue_name = get_queue_name_for_delivery_job(
                encoder_id=int(job["encoder_id"]),
                dms_id=int(job["dms_master_master_id"]),
                priority=cmm_priority,
                encoding_order_priority=int(job["encoding_order_priority"]),
            )
            messages_by_queue[queue_name].append(delivery_job)

        log.info(f"grouped {len(jobs)} jobs into {len(messages_by_queue)} queues")

        # Phase 1: promote all selected jobs to 'ready_to_encode' in chunks.
        eqd_ids_to_promote = [int(j["encoding_queue_detail_id"]) for j in jobs]
        total_db_batches = (
            len(eqd_ids_to_promote) + config.DB_BATCH_SIZE - 1
        ) // config.DB_BATCH_SIZE
        for db_batch_idx, batch_start in enumerate(
            range(0, len(eqd_ids_to_promote), config.DB_BATCH_SIZE), start=1
        ):
            batch_ids = tuple(
                eqd_ids_to_promote[batch_start : batch_start + config.DB_BATCH_SIZE]
            )
            with conn.cursor() as cursor:
                cursor.execute(
                    queries.SQL_SET_JOBS_STATUS,
                    ("ready_to_encode", batch_ids, "new"),
                )
            conn.commit()
            log.debug(
                f"db batch {db_batch_idx}/{total_db_batches} ({len(batch_ids)} jobs)"
            )
        log.info(f"DB phase done: {len(eqd_ids_to_promote)} rows promoted")

        # Phase 2: SQS fan-out. All failure tracking is internal to
        # fan_out_to_sqs; this returns the authoritative list of ids to reset.
        failed_ids = await fan_out_to_sqs(messages_by_queue, eqd_ids_to_promote, log)
        sent_count = len(eqd_ids_to_promote) - len(failed_ids)
        log.info(f"SQS phase done; sent: {sent_count}, failed: {len(failed_ids)}")

        # Phase 3: reset any rows whose SQS send failed back to 'new' so the
        # next iteration picks them up.
        if failed_ids:
            log.warning(f"resetting {len(failed_ids)} failed ids back to 'new'")
            for reset_start in range(0, len(failed_ids), config.DB_BATCH_SIZE):
                batch_ids = tuple(
                    failed_ids[reset_start : reset_start + config.DB_BATCH_SIZE]
                )
                with conn.cursor() as cursor:
                    cursor.execute(
                        queries.SQL_SET_JOBS_STATUS,
                        ("new", batch_ids, "ready_to_encode"),
                    )
                conn.commit()

    _emit_per_dms_metric("renew.run", dms_ids)


def _emit_per_dms_metric(metric: str, dms_ids: tuple[int, ...]) -> None:
    """Emit a Datadog metric once per DMS, tagged with dms_id and environment."""
    for dms_id in dms_ids:
        call_datadog_with_metric(
            metric,
            [f"dms_id:{dms_id}", f"environment:{config.ENVIRONMENT}"],
            api_key=config.DATADOG_API_KEY,
            app_key=config.DATADOG_APP_KEY,
        )


def main() -> None:
    """Run the renew daemon, throttled to one run per interval."""
    protected_run_renew = with_task_protection(run_renew)

    for i in range(config.ITERATIONS):
        start = time.monotonic()
        log = get_current_logger(uuid.uuid4())
        log.info(
            f'Running "{config.SCRIPT_NAME}" iteration {i + 1}/{config.ITERATIONS}'
        )

        if protected_run_renew(log):
            call_datadog_with_metric(
                "renew_manager.run",
                [f"environment:{config.ENVIRONMENT}"],
                api_key=config.DATADOG_API_KEY,
                app_key=config.DATADOG_APP_KEY,
            )
        elapsed = time.monotonic() - start
        log.info(
            f'Finished "{config.SCRIPT_NAME}" iteration {i + 1}/{config.ITERATIONS} '
            f"in {elapsed:.2f}s"
        )

        if i < config.ITERATIONS - 1:
            time.sleep(max(0, config.INTERVAL_SECONDS - elapsed))


if __name__ == "__main__":
    main()