import time

import psycopg

# (state, currentStage) pairs the job moves through, in order, after claim.
_STAGES = [
    ("matched", "rank"),
    ("downloading", "download"),
    ("tagging", "tag"),
    ("imported", "import"),
]


def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            'UPDATE "Job" SET state = %s, "currentStage" = %s, "updatedAt" = now() WHERE id = %s',
            (state, stage, job_id),
        )
    conn.commit()


def run_pipeline(conn: psycopg.Connection, job_id: str, stage_delay: float = 0.0) -> None:
    """Fake acquisition: walk the claimed job through to 'imported'."""
    for state, stage in _STAGES:
        if stage_delay:
            time.sleep(stage_delay)
        _set_state(conn, job_id, state, stage)

    with conn.cursor() as cur:
        cur.execute(
            'UPDATE "Request" SET status = \'completed\' '
            'WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
            (job_id,),
        )
    conn.commit()