"""Run a destructive-free PostgreSQL smoke test for job lease fencing.

The command creates one synthetic queue row, races two claimers, verifies that
only one wins, proves that a stale token is rejected, completes the winning
lease, and removes the synthetic row before exiting.
"""

from __future__ import annotations

import json
import sys
import uuid
from concurrent.futures import ThreadPoolExecutor

import psycopg

from app.core.config import Settings
from app.persistence.job_queue import (
    BackgroundJob,
    JobLease,
    LeaseLostError,
    PsycopgJobQueue,
)


def _dsn(settings: Settings) -> str:
    url = settings.database_url().set(drivername="postgresql")
    return url.render_as_string(hide_password=False)


def _insert_job(dsn: str, job_id: uuid.UUID, capability: str, idempotency_key: str) -> None:
    with psycopg.connect(dsn, connect_timeout=5) as connection:
        connection.execute(
            """
            INSERT INTO rag.background_jobs (
              id, job_type, required_capability, resource_type, resource_id,
              idempotency_key, payload, stage, max_attempts
            ) VALUES (%s, 'WORKER_SMOKE', %s, 'synthetic_smoke', %s, %s, '{}'::jsonb,
                      'VERIFYING_FENCE', 2)
            """,
            (job_id, capability, job_id, idempotency_key),
        )


def _delete_job(dsn: str, job_id: uuid.UUID) -> None:
    with psycopg.connect(dsn, connect_timeout=5) as connection:
        connection.execute("DELETE FROM rag.background_jobs WHERE id = %s", (job_id,))


def _expire_lease(dsn: str, job_id: uuid.UUID) -> None:
    """Move only the synthetic smoke lease into the past for recovery verification."""

    with psycopg.connect(dsn, connect_timeout=5) as connection:
        connection.execute(
            """
            UPDATE rag.background_jobs
               SET lease_until = now() - interval '1 second'
             WHERE id = %s AND status = 'RUNNING'
            """,
            (job_id,),
        )


def _race_claim(
    queues: tuple[PsycopgJobQueue, PsycopgJobQueue],
    capability: str,
    *,
    worker_prefix: str,
) -> tuple[BackgroundJob | None, BackgroundJob | None]:
    with ThreadPoolExecutor(max_workers=2) as executor:
        return tuple(
            executor.map(
                lambda item: item[0].claim(
                    worker_id=item[1],
                    worker_capabilities=(capability,),
                    lease_seconds=30,
                ),
                (
                    (queues[0], f"{worker_prefix}-a"),
                    (queues[1], f"{worker_prefix}-b"),
                ),
            )
        )


def run_smoke(settings: Settings) -> dict[str, object]:
    dsn = _dsn(settings)
    job_id = uuid.uuid4()
    nonce = uuid.uuid4().hex
    capability = f"worker-smoke-{nonce}"
    idempotency_key = f"worker-smoke:{nonce}"
    _insert_job(dsn, job_id, capability, idempotency_key)
    try:
        queues = (PsycopgJobQueue(dsn), PsycopgJobQueue(dsn))
        claims = _race_claim(queues, capability, worker_prefix="smoke-worker")
        winners = tuple(claim for claim in claims if claim is not None)
        if len(winners) != 1:
            raise RuntimeError("concurrent claim contract failed")
        winner = winners[0]
        stale = JobLease(
            job_id=winner.lease.job_id,
            worker_id=winner.lease.worker_id,
            lease_token=uuid.uuid4(),
        )
        try:
            queues[0].heartbeat(stale, lease_seconds=30)
        except LeaseLostError:
            fence_rejected = True
        else:
            fence_rejected = False
        if not fence_rejected:
            raise RuntimeError("stale lease fence was accepted")
        queues[0].heartbeat(winner.lease, lease_seconds=30)
        _expire_lease(dsn, job_id)
        try:
            queues[0].heartbeat(winner.lease, lease_seconds=30)
        except LeaseLostError:
            expired_lease_rejected = True
        else:
            expired_lease_rejected = False
        if not expired_lease_rejected:
            raise RuntimeError("expired lease was renewed")
        recovery_claims = _race_claim(queues, capability, worker_prefix="recovery-worker")
        recovery_winners = tuple(claim for claim in recovery_claims if claim is not None)
        if len(recovery_winners) != 1:
            raise RuntimeError("expired lease recovery contract failed")
        recovered = recovery_winners[0]
        if recovered.lease.lease_token == winner.lease.lease_token:
            raise RuntimeError("recovered lease did not rotate its fence token")
        terminal = queues[0].complete(recovered.lease)
        if terminal.status != "SUCCEEDED":
            raise RuntimeError("winning lease did not complete")
        return {
            "status": "ok",
            "claim_winners": 1,
            "stale_fence_rejected": True,
            "expired_lease_rejected": True,
            "recovery_claim_winners": 1,
            "recovery_token_rotated": True,
            "terminal_status": terminal.status,
        }
    finally:
        _delete_job(dsn, job_id)


def main() -> None:
    try:
        result = run_smoke(Settings())
        exit_code = 0
    except Exception:
        result = {"status": "failed", "error_kind": "worker_smoke_failed"}
        exit_code = 1
    sys.stdout.write(json.dumps(result, sort_keys=True) + "\n")
    raise SystemExit(exit_code)


if __name__ == "__main__":
    main()