"""
Persistence helpers for audit runs and findings.

The append-only ``AppendOnlyManager`` on ``AuditRun`` /
``AuditFinding`` blocks ``update()`` and ``bulk_update()`` at the
manager and QuerySet layers; per-finding state changes go through
``instance.save(update_fields=[...])`` exclusively.

These helpers wrap the common write patterns:

- ``bulk_create_findings``: insert PENDING findings in batches.
- ``flip_run_state``: transition an ``AuditRun`` between lifecycle
  states with the correct ``finished_at`` and ``error_message``
  semantics.
- ``persist_first_run_coercion``: writes the in-place coerced
  policy + ``dry_run=True`` + ``first_run_lock_was_active=True``
  + ``summary['first_run_lock']=True``.
- ``persist_finding_applied`` /``_skipped`` / ``_failed``: per-finding
  state writes with the correct allow-listed fields.
"""

from __future__ import annotations

import logging
from collections.abc import Callable, Collection, Mapping
from typing import TYPE_CHECKING, Final, TypeAlias

from django.db.models import Count, QuerySet
from django.utils import timezone

from .models import (
    AuditFinding,
    AuditFindingState,
    AuditRun,
    AuditRunState,
    FindingSkipReason,
    OrphanAction,
)

logger = logging.getLogger(f"extensions.{__name__}")

if TYPE_CHECKING:
    from .audit import OrphanFinding

# Resolver from finding to the action that should be persisted on the
# finding row. Lives outside persistence: plan/CLI/policy decide the
# mapping; this module just applies the resolver to each finding.
# The resolver receives the full ``OrphanFinding`` so callers can
# route by category, by aa_user_id, or both — the per-group / per-
# state policy in ``CategoryPolicy`` needs the linked AA user.
ActionResolver: TypeAlias = Callable[["OrphanFinding"], OrphanAction]

# Discord member pages are 1000 each; finding insert batches stay
# small enough to bound commit cost (50k findings -> 100 commits,
# not 50k).
_BULK_CREATE_BATCH_SIZE: Final[int] = 500

# Terminal run states for ``flip_run_state`` ↔ stamps ``finished_at``.
# Frozenset so the membership check is O(1) without rebuilding it on
# each call. ``INTERRUPTED`` is intentionally NOT terminal — it is
# resumable via the ``--resume`` subcommand, which transitions
# ``INTERRUPTED → RUNNING`` (see ``_VALID_RUN_TRANSITIONS`` below).
_TERMINAL_RUN_STATES: Final[frozenset[AuditRunState]] = frozenset(
    {
        AuditRunState.SUCCEEDED,
        AuditRunState.FAILED,
        AuditRunState.BUDGET_EXCEEDED,
        AuditRunState.ABANDONED,
    }
)

# Explicit FSM transition matrix for ``AuditRun.state``. Only legal
# transitions appear here; ``flip_run_state`` validates the requested
# ``(expected_state, new_state)`` pair against the matrix BEFORE the
# CAS-update so an illegal transition raises at the call site instead
# of silently no-op'ing through ``rowcount=0``. The unconditional
# ``flip_run_state`` path used by ``--abandon`` bypasses this matrix
# by design — abandon is the operator's force-quit hatch.
#
# Terminal states (``SUCCEEDED``, ``FAILED``, ``BUDGET_EXCEEDED``,
# ``ABANDONED``) have no outgoing transitions: a SUCCEEDED run is
# done; reopening it for further work would require a fresh row.
# ``INTERRUPTED`` is the only non-terminal "rest" state — it can be
# resumed (``→ RUNNING``) or formally abandoned (``→ ABANDONED``).
_VALID_RUN_TRANSITIONS: Final[
    Mapping[AuditRunState, frozenset[AuditRunState]]
] = {
    AuditRunState.PENDING: frozenset(
        {
            AuditRunState.RUNNING,
            AuditRunState.ABANDONED,
            AuditRunState.INTERRUPTED,
        }
    ),
    AuditRunState.RUNNING: frozenset(
        {
            AuditRunState.SUCCEEDED,
            AuditRunState.FAILED,
            AuditRunState.BUDGET_EXCEEDED,
            AuditRunState.ABANDONED,
            AuditRunState.INTERRUPTED,
        }
    ),
    AuditRunState.INTERRUPTED: frozenset(
        {AuditRunState.RUNNING, AuditRunState.ABANDONED}
    ),
    AuditRunState.SUCCEEDED: frozenset(),
    AuditRunState.FAILED: frozenset(),
    AuditRunState.BUDGET_EXCEEDED: frozenset(),
    AuditRunState.ABANDONED: frozenset(),
}


def _set_attrs(obj: object, fields: Mapping[str, object]) -> None:
    """Mirror ``name → value`` mapping onto an instance via ``setattr``."""
    # Used by ``flip_run_state`` (both branches) and
    # ``_transition_finding`` to keep the in-memory instance in sync
    # with the SQL write that just succeeded. Centralises the loop so
    # a future change (logging, accept-list) lands in one place.
    for name, value in fields.items():
        setattr(obj, name, value)


def bulk_create_findings(
    run: AuditRun,
    findings: Collection[OrphanFinding],
    *,
    action_resolver: ActionResolver,
) -> int:
    """
    Insert PENDING ``AuditFinding`` rows for each ``OrphanFinding``.

    ``action_resolver`` is a callable taking
    ``(category) -> OrphanAction`` so the resolver lives outside
    persistence (the plan resolves actions per-finding via the
    frozen policy + CLI override).
    """
    rows = [
        AuditFinding(
            run=run,
            member_id=f.member_id,
            member_username=f.member_username,
            category=f.category.value,
            aa_user_id=f.aa_user_id,
            state=AuditFindingState.PENDING.value,
            action=action_resolver(f).value,
        )
        for f in findings
    ]
    AuditFinding.objects.bulk_create(rows, batch_size=_BULK_CREATE_BATCH_SIZE)
    return len(rows)


def flip_run_state(
    run: AuditRun,
    state: AuditRunState,
    *,
    error_message: str = "",
    summary: dict[str, object] | None = None,
    expected_state: AuditRunState | tuple[AuditRunState, ...] | None = None,
) -> bool:
    """
    Transition ``run`` to the new state with the right side fields.

    When ``expected_state`` is supplied, the transition is gated by
    a SQL compare-and-set: if the DB row's current state no longer
    matches (a concurrent worker already advanced the FSM), the
    function is a no-op and returns ``False``. The in-memory ``run``
    instance is NOT mutated when the transition is skipped, so the
    caller sees the same view as the DB.

    ``expected_state`` accepts either a single ``AuditRunState`` or a
    tuple thereof. Tuple form is used by ``--resume``, which accepts
    both ``PENDING`` (run never started) and ``INTERRUPTED`` (signal
    cut a prior run short) as resumable predecessors of ``RUNNING``.

    The gated form additionally validates the requested transition
    against ``_VALID_RUN_TRANSITIONS`` and raises ``ValueError`` for
    an illegal pair (e.g. trying to flip ``SUCCEEDED → RUNNING``).
    Catching the illegal transition at the call site beats silently
    returning ``False`` because of a CAS rowcount of 0 — the latter
    looks identical to a concurrency race.

    Without ``expected_state`` the legacy unconditional ``save`` path
    is used (retained for ``--abandon``, the operator's force-quit
    hatch). New call sites should prefer the gated form.

    Returns ``True`` if the row was transitioned, ``False`` otherwise.
    """
    extra_fields: dict[str, object] = {}
    if state in _TERMINAL_RUN_STATES:
        extra_fields["finished_at"] = timezone.now()
    if error_message:
        extra_fields["error_message"] = error_message
    if summary is not None:
        extra_fields["summary"] = summary

    if expected_state is not None:
        expected_tuple: tuple[AuditRunState, ...] = (
            (expected_state,)
            if isinstance(expected_state, AuditRunState)
            else expected_state
        )
        for predecessor in expected_tuple:
            if state not in _VALID_RUN_TRANSITIONS.get(
                predecessor, frozenset()
            ):
                raise ValueError(
                    f"illegal AuditRun transition "
                    f"{predecessor.value} → {state.value}; "
                    f"see _VALID_RUN_TRANSITIONS"
                )
        rowcount = AuditRun.objects.filter(pk=run.pk).transition_state(
            expected_state=tuple(s.value for s in expected_tuple),
            new_state=state.value,
            **extra_fields,
        )
        if rowcount == 0:
            return False
        run.state = state.value
        _set_attrs(run, extra_fields)
        return True

    run.state = state.value
    _set_attrs(run, extra_fields)
    run.save(update_fields=["state", *extra_fields.keys()])
    return True


def persist_first_run_coercion(
    run: AuditRun, coerced_policy: dict[str, object]
) -> None:
    """
    Write the first-run-lock coercion via the allow-listed fields.

    All four fields touched here (``dry_run``,
    ``first_run_lock_was_active``, ``policy_snapshot``,
    ``summary``) are in the ``AuditRun.save()`` allow-list.
    """
    run.dry_run = True
    run.first_run_lock_was_active = True
    run.policy_snapshot = coerced_policy
    run.summary = {**(run.summary or {}), "first_run_lock": True}
    run.save(
        update_fields=[
            "dry_run",
            "first_run_lock_was_active",
            "policy_snapshot",
            "summary",
        ]
    )


def log_pending_purges(queryset: QuerySet[AuditRun]) -> None:
    """Emit one INFO line per PENDING run before retention cascade."""
    # Used by the Celery ``prune_audit_runs`` task and the management
    # command symmetrically. Annotates the finding count once via
    # ``Count('findings')`` so the preamble executes a single SELECT
    # instead of N COUNT(*) subqueries (the original N+1 path on
    # installs with many stuck PENDING runs).
    pending_runs = queryset.filter(state=AuditRunState.PENDING.value).annotate(
        _finding_count=Count("findings")
    )
    for run in pending_runs:
        # ``getattr`` rather than direct attribute access: the annotated
        # alias is dynamic and not on the model class, so neither the
        # type checker nor the ``SLF001`` private-access lint sees it.
        finding_count = getattr(run, "_finding_count")  # noqa: B009
        logger.info(
            "auto_purged pending audit_run id=%s findings=%s started_at=%s",
            run.pk,
            finding_count,
            run.started_at.isoformat(),
        )


def _transition_finding(
    finding: AuditFinding,
    *,
    new_state: AuditFindingState,
    **fields: object,
) -> bool:
    """CAS-transition a finding from PENDING + mirror fields locally."""
    # Idempotent under retry: the SQL UPDATE is gated by
    # ``state = PENDING`` so a second invocation on an already-terminal
    # row is a no-op (rowcount=0). On success the in-memory ``finding``
    # is updated to match the row write so callers reading the
    # instance see consistent values. ``**fields`` is forwarded
    # verbatim to ``transition_state``; the carve-out's allow-list
    # enforces that only ``_APPEND_ONLY_UPDATABLE_FIELDS`` keys reach
    # the SQL UPDATE.
    rowcount = AuditFinding.objects.filter(pk=finding.pk).transition_state(
        expected_state=AuditFindingState.PENDING.value,
        new_state=new_state.value,
        **fields,
    )
    if rowcount == 0:
        return False
    finding.state = new_state.value
    _set_attrs(finding, fields)
    return True


def persist_finding_applied(
    finding: AuditFinding,
    *,
    role_ids_removed: Collection[int],
    role_names_removed: Collection[str],
    kicked: bool = False,
) -> bool:
    """Mark a finding as APPLIED with the role-removal record (idempotent)."""
    return _transition_finding(
        finding,
        new_state=AuditFindingState.APPLIED,
        role_ids_removed=list(role_ids_removed),
        role_names_removed=list(role_names_removed),
        kicked=bool(kicked),
    )


def persist_finding_skipped(
    finding: AuditFinding, *, reason: FindingSkipReason
) -> bool:
    """Mark a finding as SKIPPED with the reason (idempotent)."""
    return _transition_finding(
        finding,
        new_state=AuditFindingState.SKIPPED,
        skipped_reason=reason.value,
    )


def persist_finding_failed(
    finding: AuditFinding, *, error_message: str
) -> bool:
    """Mark a finding as FAILED with the error message (idempotent)."""
    return _transition_finding(
        finding,
        new_state=AuditFindingState.FAILED,
        error_message=error_message,
    )


def increment_run_counters(
    run: AuditRun,
    *,
    members_scanned: int = 0,
    actions_applied: int = 0,
    actions_skipped: int = 0,
) -> None:
    """
    Bump the run's monotonic counters atomically.

    Uses ``AppendOnlyQuerySet.bump_counters`` so the SQL is
    ``UPDATE … SET col = col + N``; this is safe under concurrent
    writers (a future fan-out worker bumping the same run does
    not lose updates). The in-memory ``run`` instance is refreshed
    from the deltas so callers reading ``run.actions_applied``
    after this call see the bumped value (single-writer assumption
    is documented in apply_plan; multi-writer correctness comes
    from the F() expression on the DB side).
    """
    deltas: dict[str, int] = {
        name: value
        for name, value in (
            ("members_scanned", members_scanned),
            ("actions_applied", actions_applied),
            ("actions_skipped", actions_skipped),
        )
        if value
    }
    if not deltas:
        return
    AuditRun.objects.filter(pk=run.pk).bump_counters(**deltas)
    # Refresh the in-memory copy so callers reading the bumped
    # fields after this call see consistent values. Under the
    # single-writer-per-run invariant documented in apply_plan,
    # the in-memory delta matches the DB delta exactly. The counter
    # fields have ``default=0`` and are non-nullable, so direct
    # ``getattr(run, field)`` never yields ``None``.
    for field, delta in deltas.items():
        setattr(run, field, getattr(run, field) + delta)
