"""GPU-job queue engine (#114): enqueue dedupe + lease/heartbeat/complete/fail."""
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy import select

from backend.app.models import GpuJob, ImageRecord
from backend.app.services.ml.gpu_jobs import GpuJobService

pytestmark = pytest.mark.integration


async def _img(db, sha) -> ImageRecord:
    img = ImageRecord(
        path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
        width=1, height=1, origin="imported_filesystem", integrity_status="unknown",
    )
    db.add(img)
    await db.flush()
    return img


@pytest.mark.asyncio
async def test_enqueue_dedupes_same_pair(db):
    img = await _img(db, "a" * 64)
    svc = GpuJobService(db)
    first = await svc.enqueue(img.id, "ccip")
    dup = await svc.enqueue(img.id, "ccip")
    other = await svc.enqueue(img.id, "siglip_region")
    await db.commit()
    assert first is not None
    assert dup is None              # same (image, task) already queued
    assert other is not None        # different task is fine


@pytest.mark.asyncio
async def test_lease_claims_then_skips_when_held(db):
    img = await _img(db, "b" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()

    leased = await svc.lease("agent-1", batch_size=8)
    await db.commit()
    assert len(leased) == 1
    assert leased[0].status == "leased" and leased[0].lease_token == "agent-1"
    assert leased[0].attempts == 1

    # Already leased + not expired → a second agent gets nothing.
    again = await svc.lease("agent-2", batch_size=8)
    await db.commit()
    assert again == []


@pytest.mark.asyncio
async def test_expired_lease_is_reclaimed(db):
    img = await _img(db, "c" * 64)
    svc = GpuJobService(db)
    job = await svc.enqueue(img.id, "ccip")
    await db.commit()
    # Force the lease into the past.
    job.status = "leased"
    job.lease_token = "dead-agent"
    job.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
    await db.commit()

    leased = await svc.lease("agent-2", batch_size=8)
    await db.commit()
    assert len(leased) == 1
    assert leased[0].lease_token == "agent-2"
    assert leased[0].attempts == 1   # re-lease incremented from 0 (was set directly)


@pytest.mark.asyncio
async def test_heartbeat_extends_only_own_lease(db):
    img = await _img(db, "d" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]
    await db.commit()

    assert await svc.heartbeat("agent-1", [job.id]) == 1
    assert await svc.heartbeat("someone-else", [job.id]) == 0


@pytest.mark.asyncio
async def test_complete_closes_job(db):
    img = await _img(db, "e" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]
    await db.commit()

    assert await svc.complete("wrong-token", job.id) is False
    assert await svc.complete("agent-1", job.id) is True
    await db.commit()
    fresh = await db.get(GpuJob, job.id)
    assert fresh.status == "done" and fresh.lease_token is None


@pytest.mark.asyncio
async def test_fail_requeues_until_cap(db):
    img = await _img(db, "f" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]   # attempts -> 1
    await db.commit()
    # Under the cap → back to pending for a retry.
    assert await svc.fail("agent-1", job.id, "boom") is True
    await db.commit()
    assert (await db.get(GpuJob, job.id)).status == "pending"

    # At the attempt cap → terminal 'error'.
    j = await db.get(GpuJob, job.id)
    j.attempts = 3
    j.status = "leased"
    j.lease_token = "agent-1"
    j.lease_expires_at = datetime.now(UTC) + timedelta(minutes=5)
    await db.commit()
    assert await svc.fail("agent-1", job.id, "boom again") is True
    await db.commit()
    assert (await db.get(GpuJob, job.id)).status == "error"