import psycopg

# States a job passes through while a worker is actively processing it. On startup this
# fresh worker owns none of them, so any job left here was orphaned by a crash/restart.
_IN_FLIGHT_STATES = ("matching", "matched", "downloading", "tagging")


def reclaim_stuck_jobs(conn: psycopg.Connection) -> int:
    """Reset jobs left mid-pipeline by a prior worker crash/restart back to 'requested' so
    they get re-claimed. Safe to call once at startup: nothing is processing them, staging is
    cleared beforehand, and the pipeline redoes intake→import idempotently (atomic import).
    Returns how many were reclaimed. Terminal jobs (imported/needs_attention) are untouched."""
    with conn.cursor() as cur:
        cur.execute(
            """
            UPDATE "Job"
            SET state = 'requested', "currentStage" = 'intake', "claimedAt" = NULL,
                error = NULL, "downloadProgress" = 0, "updatedAt" = now()
            WHERE state IN ('matching', 'matched', 'downloading', 'tagging')
            """
        )
        n = cur.rowcount
    conn.commit()
    return n


_MAX_ATTEMPTS = 5


def requeue_or_fail(
    conn: psycopg.Connection, job_id: str, error: str, max_attempts: int = _MAX_ATTEMPTS
) -> None:
    """Recover a job whose pipeline raised an *unhandled* exception (e.g. a dropped MusicBrainz
    TLS connection surfacing as an SSL EOF). Without this the job is left in its mid-pipeline
    'matching' state; since claim_next only picks 'requested' it is orphaned until the next
    restart's reclaim_stuck_jobs — showing "Searching sources…" on the press indefinitely.

    Such errors are usually transient, so requeue to 'requested' until the attempt cap, then fall
    to 'needs_attention' so a persistently-failing job stays visible instead of looping forever.
    `attempts` is bumped by claim_next on every claim, so it already counts tries."""
    with conn.cursor() as cur:
        cur.execute('SELECT attempts, "requestId" FROM "Job" WHERE id = %s', (job_id,))
        row = cur.fetchone()
        if row is None:
            return
        attempts, request_id = row
        if attempts >= max_attempts:
            cur.execute(
                'UPDATE "Job" SET state = \'needs_attention\', error = %s, "claimedAt" = NULL, '
                '"updatedAt" = now() WHERE id = %s',
                (error, job_id),
            )
            cur.execute('UPDATE "Request" SET status = \'needs_attention\' WHERE id = %s', (request_id,))
        else:
            cur.execute(
                'UPDATE "Job" SET state = \'requested\', "currentStage" = \'intake\', error = NULL, '
                '"claimedAt" = NULL, "downloadProgress" = 0, "updatedAt" = now() WHERE id = %s',
                (job_id,),
            )
            cur.execute('UPDATE "Request" SET status = \'pending\' WHERE id = %s', (request_id,))
    conn.commit()


def claim_next(conn: psycopg.Connection) -> str | None:
    """Atomically claim the oldest queued job. Returns its id or None."""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT id FROM "Job"
            WHERE state = 'requested' AND NOT paused
            ORDER BY "createdAt" ASC
            FOR UPDATE SKIP LOCKED
            LIMIT 1
            """
        )
        row = cur.fetchone()
        if row is None:
            conn.rollback()
            return None
        job_id = row[0]
        cur.execute(
            """
            UPDATE "Job"
            SET state = 'matching',
                "currentStage" = 'match',
                "claimedAt" = now(),
                attempts = attempts + 1,
                "updatedAt" = now()
            WHERE id = %s
            """,
            (job_id,),
        )
    conn.commit()
    return job_id