"""scan_library_for_rule Celery task — iterates image_record in keyset-
paginated batches, evaluates the audit rule per image, populates
LibraryAuditRun.matched_ids. Runs on the maintenance queue with a 2h soft
time limit (plenty of margin for 100k+ image libraries at ~100ms PIL
decode + histogram per image).

State machine:
  start: status='running'
  end success: status='ready'
  end error: status='error', error=traceback
  oversize: status='error', error='matched too many images; tighten threshold'
  external cancel: scan sees status='cancelled' between batches, exits.
"""

import logging
import time
import traceback
from datetime import UTC, datetime

from celery.exceptions import SoftTimeLimitExceeded
from PIL import Image
from sqlalchemy import select, update
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..models import ImageRecord, LibraryAuditRun
from ..services.audits import single_color, transparency
from ._sync_engine import sync_session_factory as _sync_session_factory

log = logging.getLogger(__name__)

_BATCH = 500
_PROGRESS_TICK = 100
_MAX_MATCHED = 50_000
# One chunk's wall-clock budget. Was a single 2h pass that timed out on large
# libraries and held the concurrency-1 maintenance queue the whole time
# (operator-flagged 2026-06-07). Now: scan ~10 min, persist the keyset cursor +
# matches, re-enqueue to continue — so backups/vacuum/normalize chunks can
# interleave. soft/hard limits sit just above so the budget fires first.
_CHUNK_SECONDS = 600

_RULES = {
    "transparency": transparency.evaluate,
    "single_color": single_color.evaluate,
}


@celery.task(
    name="backend.app.tasks.library_audit.scan_library_for_rule",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError),
    retry_backoff=5,
    retry_backoff_max=60,
    retry_jitter=True,
    max_retries=3,
    soft_time_limit=900,
    time_limit=1000,
)
def scan_library_for_rule(self, audit_id: int) -> dict:
    """See module docstring. Time-boxed + self-resuming: one call scans a
    ~10-min chunk, persists the resume cursor + matches, and re-enqueues itself
    until the library is exhausted. Returns a small summary dict for eager-mode
    test assertions (real workers ignore the return value)."""
    SessionLocal = _sync_session_factory()
    start = time.monotonic()
    try:
        with SessionLocal() as session:
            audit = session.get(LibraryAuditRun, audit_id)
            if audit is None:
                return {"audit_id": audit_id, "status": "missing"}
            evaluate = _RULES.get(audit.rule)
            if evaluate is None:
                _mark_error(session, audit_id, f"unknown rule {audit.rule!r}")
                return {"audit_id": audit_id, "status": "error"}
            params = dict(audit.params or {})
            # Resume from the previous chunk's persisted state.
            matched: list[int] = list(audit.matched_ids or [])
            scanned = audit.scanned_count or 0
            last_id = audit.resume_after_id or 0
            while True:
                # Cancellation check between batches.
                current_status = session.execute(
                    select(LibraryAuditRun.status)
                    .where(LibraryAuditRun.id == audit_id)
                ).scalar_one()
                if current_status == "cancelled":
                    return {"audit_id": audit_id, "status": "cancelled"}
                # Time-box: persist the cursor + matches and re-enqueue so the
                # queue is freed between chunks. The next call resumes here.
                if time.monotonic() - start >= _CHUNK_SECONDS:
                    _persist_chunk(session, audit_id, scanned, matched, last_id)
                    scan_library_for_rule.delay(audit_id)
                    return {
                        "audit_id": audit_id, "status": "running",
                        "partial": True, "scanned": scanned,
                    }
                rows = session.execute(
                    select(ImageRecord.id, ImageRecord.path)
                    .where(ImageRecord.id > last_id)
                    .where(ImageRecord.mime.like("image/%"))
                    .order_by(ImageRecord.id.asc())
                    .limit(_BATCH)
                ).all()
                if not rows:
                    break
                for image_id, image_path in rows:
                    last_id = image_id
                    scanned += 1
                    try:
                        with Image.open(image_path) as im:
                            try:
                                if evaluate(im, **params):
                                    matched.append(image_id)
                            except Exception as exc:  # noqa: BLE001
                                log.warning(
                                    "audit %s: rule evaluate failed on %s: %s",
                                    audit_id, image_path, exc,
                                )
                    except FileNotFoundError:
                        log.warning(
                            "audit %s: image_record %s file missing at %s; skipping",
                            audit_id, image_id, image_path,
                        )
                    except OSError as exc:
                        log.warning(
                            "audit %s: PIL load failed for %s: %s",
                            audit_id, image_path, exc,
                        )
                    if len(matched) > _MAX_MATCHED:
                        _mark_error(
                            session, audit_id,
                            f"matched > {_MAX_MATCHED} images; "
                            "tighten threshold and re-run",
                        )
                        return {"audit_id": audit_id, "status": "error"}
                    if scanned % _PROGRESS_TICK == 0:
                        # Cheap heartbeat: scanned_count + last_progress_at so the
                        # recovery sweep sees the multi-chunk audit is alive. The
                        # cursor + matches are persisted at chunk boundaries.
                        session.execute(
                            update(LibraryAuditRun)
                            .where(LibraryAuditRun.id == audit_id)
                            .values(
                                scanned_count=scanned,
                                last_progress_at=datetime.now(UTC),
                            )
                        )
                        session.commit()
            # Final state.
            session.execute(
                update(LibraryAuditRun)
                .where(LibraryAuditRun.id == audit_id)
                .values(
                    scanned_count=scanned,
                    matched_count=len(matched),
                    matched_ids=matched,
                    resume_after_id=last_id,
                    status="ready",
                    finished_at=datetime.now(UTC),
                    last_progress_at=datetime.now(UTC),
                )
            )
            session.commit()
            return {
                "audit_id": audit_id,
                "status": "ready",
                "scanned": scanned,
                "matched": len(matched),
            }
    except SoftTimeLimitExceeded:
        # Backstop (the in-chunk budget should fire first): the audit stays
        # 'running' with its last committed cursor; re-enqueue to continue from
        # there rather than marking the whole run an error.
        log.warning(
            "audit %s: soft time limit hit — re-enqueuing to resume", audit_id,
        )
        scan_library_for_rule.delay(audit_id)
        return {"audit_id": audit_id, "status": "running", "partial": True}
    except (OperationalError, DBAPIError):
        # Retryable per the decorator; leave row in 'running' and let
        # autoretry try again. Recovery sweep catches if all retries fail.
        raise
    except Exception:  # noqa: BLE001
        tb = traceback.format_exc()
        with SessionLocal() as session:
            _mark_error(session, audit_id, tb)
        raise


def _persist_chunk(session, audit_id, scanned, matched, last_id) -> None:
    """Persist a chunk boundary: scanned count, matches so far, and the keyset
    cursor the next chunk resumes from. Keeps status='running'."""
    session.execute(
        update(LibraryAuditRun)
        .where(LibraryAuditRun.id == audit_id)
        .values(
            scanned_count=scanned,
            matched_count=len(matched),
            matched_ids=list(matched),
            resume_after_id=last_id,
            last_progress_at=datetime.now(UTC),
        )
    )
    session.commit()


def _mark_error(session, audit_id: int, error_msg: str) -> None:
    session.execute(
        update(LibraryAuditRun)
        .where(LibraryAuditRun.id == audit_id)
        .values(
            status="error",
            error=error_msg,
            finished_at=datetime.now(UTC),
        )
    )
    session.commit()
