from lyra_worker.claim import claim_next, reclaim_stuck_jobs
from tests.conftest import insert_request


def _set_state(conn, job_id, state, stage="download", claimed=True, progress=0.5):
    with conn.cursor() as cur:
        cur.execute(
            'UPDATE "Job" SET state = %s, "currentStage" = %s, '
            '"claimedAt" = ' + ("now()" if claimed else "NULL") + ', '
            '"downloadProgress" = %s, error = %s WHERE id = %s',
            (state, stage, progress, "boom" if state == "needs_attention" else None, job_id),
        )
    conn.commit()


def _job(conn, job_id):
    with conn.cursor() as cur:
        cur.execute('SELECT state, "currentStage", "claimedAt", "downloadProgress", error '
                    'FROM "Job" WHERE id = %s', (job_id,))
        return cur.fetchone()


def test_reclaims_in_flight_states(conn):
    ids = {}
    for state in ("matching", "matched", "downloading", "tagging"):
        jid = insert_request(conn, album=state)
        _set_state(conn, jid, state)
        ids[state] = jid

    assert reclaim_stuck_jobs(conn) == 4

    for state, jid in ids.items():
        st, stage, claimed_at, progress, err = _job(conn, jid)
        assert st == "requested", state
        assert stage == "intake"
        assert claimed_at is None
        assert progress == 0
        assert err is None
    # a reclaimed job is immediately claimable again
    assert claim_next(conn) is not None


def test_leaves_terminal_and_queued_jobs_untouched(conn):
    imported = insert_request(conn, album="imported")
    _set_state(conn, imported, "imported", stage="import")
    attention = insert_request(conn, album="attention")
    _set_state(conn, attention, "needs_attention", stage="download")
    queued = insert_request(conn, album="queued")  # stays 'requested'/'intake' from insert

    assert reclaim_stuck_jobs(conn) == 0

    assert _job(conn, imported)[0] == "imported"
    assert _job(conn, attention)[0] == "needs_attention"
    assert _job(conn, attention)[4] == "boom"  # its error is preserved
    assert _job(conn, queued)[0] == "requested"