"""Opt-in real PostgreSQL queue races, using an explicitly disposable database."""
import asyncio
import os
from pathlib import Path
from uuid import uuid4

import pytest
from fastapi import HTTPException
from psycopg import AsyncConnection
from psycopg.conninfo import conninfo_to_dict
from psycopg.rows import dict_row

from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository


def test_concurrent_submission_claims_fencing_and_project_isolation():
    dsn = os.getenv("CALCULATION_QUEUE_TEST_DSN")
    if not dsn:
        pytest.skip("Set CALCULATION_QUEUE_TEST_DSN to a disposable PostgreSQL database")
    assert conninfo_to_dict(dsn)["dbname"].startswith("tjwater_calc_test_")

    async def run():
        project, other, owner = uuid4(), uuid4(), uuid4()
        async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
            await conn.execute((Path(__file__).parents[2] / "resources/sql/026_calculation_jobs.sql").read_text())

        async def submit(key, parameters=None, selected_project=project):
            async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
                return await CalculationJobRepository(conn, selected_project).enqueue(
                    owner_user_id=owner, request_key=key, tool="run_burst_location",
                    parameters=parameters or {"x": 1}, binding={"native_session_id": "session"}, timeout_seconds=600)

        # Four independent API connections, same logical submission.
        submissions = await asyncio.gather(*(submit("same") for _ in range(4)))
        assert len({str(row["job_id"]) for row in submissions}) == 1
        with pytest.raises(HTTPException) as conflict:
            await submit("same", {"x": 2})
        assert conflict.value.status_code == 409
        second = await submit("second")
        foreign = await submit("foreign", selected_project=other)

        async def claim(worker):
            async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
                return await CalculationJobRepository(conn, project).claim(worker, 60)

        claimed = await asyncio.gather(*(claim(f"worker-{index}") for index in range(4)))
        rows = [row for row in claimed if row]
        assert len(rows) == 2
        assert len({str(row["job_id"]) for row in rows}) == 2
        assert str(foreign["job_id"]) not in {str(row["job_id"]) for row in rows}

        async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
            repo = CalculationJobRepository(conn, project)
            assert await repo.get(foreign["job_id"], owner) is None
            assert await repo.get(second["job_id"], uuid4()) is None
            row = rows[0]
            stale = {**row, "lease_generation": row["lease_generation"] - 1}
            assert await repo.heartbeat(stale, row["lease_owner"], 60) is None
            assert not await repo.finish(stale, row["lease_owner"], "succeeded", result={})
            cancelled = await repo.cancel(row["job_id"], owner)
            assert cancelled["status"] == "cancel_requested"
            assert not await repo.finish(row, row["lease_owner"], "succeeded", result={})
            assert await repo.finish(row, row["lease_owner"], "cancelled")
            expired = rows[1]
            await conn.execute("UPDATE workbench.calculation_jobs SET lease_expires_at=now()-interval '1 second' WHERE job_id=%s", (expired["job_id"],))
            await repo.recover()
            assert (await repo.get(expired["job_id"], owner))["status"] == "interrupted"
            assert not await repo.finish(expired, expired["lease_owner"], "succeeded", result={})
            assert await repo.claim("replacement", 60) is None
    asyncio.run(run())