"""Lambda handlers — thin event-shaped wrappers around task functions."""

import datetime
import itertools
import logging
import uuid
from collections.abc import Sequence
from dataclasses import dataclass, field

from fansifter_common.utils import timezone

from resonance_engine.adapters import aws_sqs
from resonance_engine.adapters.db import db
from resonance_engine.config import settings
from resonance_engine.dsp.enums import DSPId
from resonance_engine.dsp.models import DSPClient
from resonance_engine.fandata import collector
from resonance_engine.fandata.collector import FanBatchOutcome
from resonance_engine.fandata.models import FanCollectionState, FanConnectionFilter
from resonance_engine.tasks import planner
from resonance_engine.tasks.enums import FanoutSource, TaskFinishedReason, TaskStatus
from resonance_engine.tasks.models import CollectTask, FanoutTask
from resonance_engine.tasks.types import FanCollectBatch
from resonance_engine.tasks.utils import queue_url_for

STAMP_BATCH = 500

logger = logging.getLogger(__name__)


@dataclass(kw_only=True, frozen=True)
class RunFanoutRequest:
    source: FanoutSource
    filters: FanConnectionFilter = field(default_factory=planner.get_default_filters)
    force: bool = False


@dataclass(kw_only=True, frozen=True)
class RunFanoutResponse:
    results: list[planner.FanoutPlanResult]
    task: FanoutTask | None


def run_fanout(request: RunFanoutRequest) -> RunFanoutResponse:
    # Decorated so the write retries on DSQL serialization conflicts (OC000).
    @db.transaction
    def _close_stale() -> int:
        return FanoutTask.query.close_stale(
            dispatch_stale_s=settings.fan_fanout_dispatch_stale_s,
        )

    stale = _close_stale()
    if stale:
        logger.warning("run_fanout: closed %d stale fanout task(s)", stale)

    # Read-only planning runs outside a write transaction: its per-client COUNT
    # and fan-selection queries must not accumulate against DSQL's 300s
    # transaction time limit (which otherwise aborts the whole fanout tick).
    with db.autocommit():
        results = planner.plan_fanout(filters=request.filters, source=request.source)

    # Nothing to dispatch
    if not any(result.fans for result in results):
        logger.info("run_fanout: no eligible fans — nothing to dispatch")
        return RunFanoutResponse(results=[], task=None)

    now = timezone.now()
    with db.transaction():
        task = FanoutTask(
            started_at=now,
            source=request.source,
            status=TaskStatus.running,
            filters=request.filters,
        )
        task.save()

    pairs = [pair for result in results for pair in result.dispatch_pairs]
    for batch in itertools.batched(pairs, STAMP_BATCH, strict=False):
        _stamp_dispatched(batch, now=now)

    if settings.fan_fanout_async_enabled:
        _run_fanout_async(task.id, results, force=request.force)
    else:
        _run_fanout_sync(task.id, results, force=request.force)

    return RunFanoutResponse(results=results, task=task)


def _make_collect_batches(
    task_id: uuid.UUID, result: planner.FanoutPlanResult, *, force: bool
) -> list[FanCollectBatch]:
    return [
        FanCollectBatch(
            dsp_client_name=result.client_name,
            fanout_task_id=task_id,
            fans=list(chunk),
            force=force,
        )
        for chunk in itertools.batched(result.fans, result.batch_size, strict=False)
    ]


@db.transaction
def _stamp_dispatched(
    pairs: Sequence[tuple[str, DSPId]], now: datetime.datetime
) -> None:
    FanCollectionState.query.stamp_dispatched(pairs, now=now)


def _run_fanout_async(
    task_id: uuid.UUID, results: list[planner.FanoutPlanResult], *, force: bool
) -> None:
    fans_dispatched = 0
    messages_sent = 0

    for result in results:
        if not result.fans:
            continue

        queue_url = queue_url_for(result.client_name)
        if queue_url is None:
            logger.warning(
                "Fanout %s — %d fans planned but no SQS queue configured",
                result.client_name,
                len(result.fans),
                extra={
                    "dsp_client_name": result.client_name,
                    "fanout_task_id": task_id,
                },
            )
            continue

        batches = _make_collect_batches(task_id, result, force=force)

        aws_sqs.send_messages_batch(
            queue_url=queue_url,
            bodies=(batch.model_dump_json() for batch in batches),
        )

        fans_dispatched += len(result.fans)
        messages_sent += len(batches)

        logger.info(
            "Fanout %s — dispatched %d fans in %d batches",
            result.client_name,
            len(result.fans),
            len(batches),
            extra={
                "dsp_client_name": result.client_name,
                "fanout_task_id": task_id,
            },
        )

    with db.autocommit():
        FanoutTask.query.update_dispatch(
            task_id,
            fans_dispatched=fans_dispatched,
            messages_sent=messages_sent,
        )


def _run_fanout_sync(
    task_id: uuid.UUID, results: list[planner.FanoutPlanResult], *, force: bool
) -> None:
    for result in results:
        if not result.fans:
            continue

        batches = _make_collect_batches(task_id, result, force=force)

        for batch in batches:
            try:
                run_collect(RunCollectRequest(batch=batch))
            except Exception:
                logger.exception(
                    "Fanout %s — inline collect failed for batch",
                    result.client_name,
                    extra={
                        "dsp_client_name": result.client_name,
                        "fanout_task_id": task_id,
                        "collect_task_id": batch.collect_task_id,
                    },
                )

        # Increment after each DSP client so progress is visible mid-run.
        with db.autocommit():
            FanoutTask.query.increment_dispatch(
                task_id,
                fans=len(result.fans),
                messages=len(batches),
            )


# ─── Collect ──────────────────────────────────────────────────────────────────


@dataclass(kw_only=True, frozen=True)
class RunCollectRequest:
    batch: FanCollectBatch


def run_collect(request: RunCollectRequest) -> FanBatchOutcome | None:
    fans = request.batch.fans
    started_at = timezone.now()

    # Transaction 1: look up the DSP client and get-or-create the CollectTask.
    # Committed immediately so status=running is visible during collection.
    # Get-or-create prevents duplicate rows when SQS retries the message.
    with db.transaction():
        dsp_client = DSPClient.query.where(
            DSPClient.name == request.batch.dsp_client_name
        ).one_or_none()

        if dsp_client is None:
            logger.error(
                "Collect %s — DSP client not found, aborting batch",
                request.batch.dsp_client_name,
                extra={
                    "dsp_client_name": request.batch.dsp_client_name,
                    "fanout_task_id": request.batch.fanout_task_id,
                },
            )
            return None

        # Get-or-create using the collect_task_id baked into the SQS message at dispatch
        # time. On retry, SQS re-delivers the same body → same ID → no duplicate rows.
        task = CollectTask.query.where(
            CollectTask.id == request.batch.collect_task_id
        ).one_or_none()

        if task is not None and task.status == TaskStatus.done:
            logger.info(
                "Collect %s — already completed, skipping (idempotent retry)",
                request.batch.dsp_client_name,
                extra={
                    "dsp_client_name": request.batch.dsp_client_name,
                    "fanout_task_id": request.batch.fanout_task_id,
                    "collect_task_id": task.id,
                },
            )
            return None

        if task is not None:
            # Running task from a failed previous attempt — re-run collection below
            # using the same record so the outcome update lands on it.
            logger.warning(
                "Collect %s — retrying failed attempt (existing task reused)",
                request.batch.dsp_client_name,
                extra={
                    "dsp_client_name": request.batch.dsp_client_name,
                    "fanout_task_id": request.batch.fanout_task_id,
                    "collect_task_id": task.id,
                },
            )
        else:
            task = CollectTask(
                started_at=started_at,
                fanout_task_id=request.batch.fanout_task_id,
                dsp_client_name=request.batch.dsp_client_name,
                status=TaskStatus.running,
                fans_total=len(fans),
            )
            task.id = request.batch.collect_task_id
            task.save()

    outcome = collector.collect_fans(
        fans=fans, client=dsp_client, force=request.batch.force
    )

    # Record the outcome, then close the FanoutTask if this was the last running
    # sibling. Decorated to retry the OCC conflict siblings hit on the shared row.
    @db.transaction
    def _finalize() -> bool:
        CollectTask.query.update_outcome(
            task.id,
            fans_processed=outcome.fans_processed,
            fans_errors=outcome.fans_errors,
            fans_stale_tokens=outcome.fans_stale_tokens,
            requests_rate_limited=outcome.requests_rate_limited,
            fans_skipped=outcome.fans_skipped,
            requests_skipped=outcome.requests_skipped,
            requests=outcome.requests,
        )
        return FanoutTask.query.close_if_all_collected(task.fanout_task_id)

    if _finalize():
        logger.info(
            "Collect %s — fanout complete",
            task.dsp_client_name,
            extra={
                "dsp_client_name": task.dsp_client_name,
                "fanout_task_id": task.fanout_task_id,
            },
        )

    return outcome


# ─── Collect DLQ ──────────────────────────────────────────────────────────────


@dataclass(kw_only=True, frozen=True)
class RunCollectDLQRequest:
    batch: FanCollectBatch


def run_collect_dlq(request: RunCollectDLQRequest) -> None:
    """Handle a collect SQS message that landed in the DLQ (all retries exhausted)."""
    batch = request.batch
    now = timezone.now()

    with db.transaction():
        task = CollectTask.query.get(batch.collect_task_id)

        if task is None:
            # Lambda never started (e.g. throttled immediately to DLQ).
            # Create a stale record so close_if_all_collected can settle the fanout.
            task = CollectTask(
                started_at=now,
                fanout_task_id=batch.fanout_task_id,
                dsp_client_name=batch.dsp_client_name,
                status=TaskStatus.stale,
                fans_total=len(batch.fans),
                finished_at=now,
                finished_reason=TaskFinishedReason.timed_out,
            )
            task.id = batch.collect_task_id
            task.save()
            changed = True
        else:
            changed = CollectTask.query.mark_stale_one(batch.collect_task_id)

    if not changed:
        logger.info(
            "Collect DLQ %s — task already done, skipping",
            batch.dsp_client_name,
            extra={
                "dsp_client_name": batch.dsp_client_name,
                "fanout_task_id": batch.fanout_task_id,
                "collect_task_id": batch.collect_task_id,
            },
        )
        return

    logger.warning(
        "Collect DLQ %s — marked stale, %d fans lost (SQS retries exhausted)",
        batch.dsp_client_name,
        len(batch.fans),
        extra={
            "dsp_client_name": batch.dsp_client_name,
            "fanout_task_id": batch.fanout_task_id,
            "collect_task_id": batch.collect_task_id,
            "fans_total": len(batch.fans),
        },
    )

    # Decorated so the shared-row close retries on serialization conflicts.
    @db.transaction
    def _close_fanout() -> bool:
        return FanoutTask.query.close_if_all_collected(batch.fanout_task_id)

    if _close_fanout():
        logger.warning(
            "Collect DLQ %s — fanout closed with stale tasks",
            batch.dsp_client_name,
            extra={
                "dsp_client_name": batch.dsp_client_name,
                "fanout_task_id": batch.fanout_task_id,
            },
        )