"""PR-Worker-Pool-Registry-1B — /internal/worker/* 5 endpoint 정상 동작 검증.

5 항목:
  1. /register UPSERT — 같은 worker_id 두 번 호출하면 단일 row 유지 + last_registered_at 갱신.
  2. /heartbeat — INSERT row 추가 + status='available' 저장 확인.
  3. /claim — queue empty 시 204 + body 0. pending row 있으면 200 + payload + status='processing'.
  4. /result — completed 시 status='completed' + result JSONB 저장.
  5. /drain — heartbeat row status='draining' INSERT.

DB integration. GPU docker compose exec fastapi pytest 환경 가정.
"""

from __future__ import annotations

import os
import sys
import uuid

import pytest
import pytest_asyncio

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "app"))

from httpx import ASGITransport, AsyncClient
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from _worker_pool_helpers import (
    cleanup_worker_capabilities,
    cleanup_worker_jobs,
    ensure_user,
    get_database_url,
    mint_access_token,
)


@pytest_asyncio.fixture
async def db_session(monkeypatch):
    monkeypatch.setenv("LAPTOP_WORKER_BOT_USERNAME", "laptop-worker-bot")
    engine = create_async_engine(get_database_url())
    sm = async_sessionmaker(engine, expire_on_commit=False)
    async with sm() as session:
        yield session
    await engine.dispose()


@pytest_asyncio.fixture
async def client():
    from main import app

    async with AsyncClient(
        transport=ASGITransport(app=app),
        base_url="http://test",
    ) as ac:
        yield ac


@pytest_asyncio.fixture
async def worker_token(db_session):
    await ensure_user(db_session, "laptop-worker-bot")
    return mint_access_token("laptop-worker-bot")


@pytest_asyncio.fixture
async def owner_id(db_session):
    return await ensure_user(db_session, "test-owner-user-1b")


@pytest_asyncio.fixture
async def worker_id_unique():
    wid = f"test-mbp-1b-{uuid.uuid4().hex[:8]}"
    yield wid


@pytest_asyncio.fixture
async def cleanup_after(db_session, worker_id_unique):
    yield
    await cleanup_worker_jobs(db_session, "test-rec-1b")
    await cleanup_worker_capabilities(db_session, worker_id_unique)


def _auth(token: str) -> dict:
    return {"Authorization": f"Bearer {token}"}


@pytest.mark.asyncio
async def test_register_upsert(client, db_session, worker_token, worker_id_unique, cleanup_after):
    body = {
        "worker_id": worker_id_unique,
        "device_label": "MacBook Pro M3 Pro",
        "worker_class": "laptop",
        "tier": "laptop_small",
        "capabilities": ["short_summary"],
        "models_loaded": ["gemma-3-4b-it"],
    }
    r1 = await client.post("/internal/worker/register", json=body, headers=_auth(worker_token))
    assert r1.status_code == 200, r1.text
    # 두번째 호출 = upsert 갱신
    body["device_label"] = "MacBook Pro M3 Pro (renamed)"
    r2 = await client.post("/internal/worker/register", json=body, headers=_auth(worker_token))
    assert r2.status_code == 200
    res = await db_session.execute(
        text(
            "SELECT device_label, last_registered_at FROM worker_capabilities WHERE worker_id = :w"
        ),
        {"w": worker_id_unique},
    )
    rows = res.all()
    assert len(rows) == 1
    assert rows[0][0] == "MacBook Pro M3 Pro (renamed)"


@pytest.mark.asyncio
async def test_heartbeat_insert(
    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
):
    # 사전: worker_capabilities row 필요 (FK)
    await db_session.execute(
        text(
            "INSERT INTO worker_capabilities (worker_id, user_id, device_label, "
            "worker_class, tier) VALUES (:w, :u, 'lbl', 'laptop', 'laptop_small') "
            "ON CONFLICT (worker_id) DO NOTHING"
        ),
        {"w": worker_id_unique, "u": owner_id},
    )
    await db_session.commit()

    r = await client.post(
        "/internal/worker/heartbeat",
        json={"worker_id": worker_id_unique, "status": "available"},
        headers=_auth(worker_token),
    )
    assert r.status_code == 200
    res = await db_session.execute(
        text(
            "SELECT status FROM worker_heartbeats WHERE worker_id = :w ORDER BY heartbeat_at DESC LIMIT 1"
        ),
        {"w": worker_id_unique},
    )
    assert res.scalar_one() == "available"


@pytest.mark.asyncio
async def test_claim_204_when_empty(client, worker_token, cleanup_after):
    r = await client.post(
        "/internal/worker/claim",
        json={"worker_id": "test-empty-1b", "job_type": "test-rec-1b-empty"},
        headers=_auth(worker_token),
    )
    assert r.status_code == 204
    # 정정 #4: body 0
    assert r.content == b""


@pytest.mark.asyncio
async def test_claim_processing_transition(
    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
):
    # seed worker_capabilities + worker_jobs
    await db_session.execute(
        text(
            "INSERT INTO worker_capabilities (worker_id, user_id, device_label, "
            "worker_class, tier) VALUES (:w, :u, 'lbl', 'laptop', 'laptop_small') "
            "ON CONFLICT (worker_id) DO NOTHING"
        ),
        {"w": worker_id_unique, "u": owner_id},
    )
    jt = "test-rec-1b-claim"
    job_id = (
        await db_session.execute(
            text(
                "INSERT INTO worker_jobs (user_id, job_type, payload) "
                "VALUES (:u, :j, '{\"week\":\"2026-W20\"}'::jsonb) RETURNING id"
            ),
            {"u": owner_id, "j": jt},
        )
    ).scalar_one()
    await db_session.commit()

    r = await client.post(
        "/internal/worker/claim",
        json={"worker_id": worker_id_unique, "job_type": jt},
        headers=_auth(worker_token),
    )
    assert r.status_code == 200
    js = r.json()
    assert js["id"] == job_id
    assert js["job_type"] == jt
    assert js["attempts"] == 1
    assert js["payload"]["week"] == "2026-W20"

    res = await db_session.execute(
        text("SELECT status, worker_id, attempts FROM worker_jobs WHERE id = :i"),
        {"i": job_id},
    )
    status_, w_id, attempts = res.first()
    assert status_ == "processing"
    assert w_id == worker_id_unique
    assert attempts == 1


@pytest.mark.asyncio
async def test_result_completed(
    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
):
    # seed
    await db_session.execute(
        text(
            "INSERT INTO worker_capabilities (worker_id, user_id, device_label, "
            "worker_class, tier) VALUES (:w, :u, 'lbl', 'laptop', 'laptop_small') "
            "ON CONFLICT (worker_id) DO NOTHING"
        ),
        {"w": worker_id_unique, "u": owner_id},
    )
    jt = "test-rec-1b-result"
    job_id = (
        await db_session.execute(
            text(
                "INSERT INTO worker_jobs (user_id, job_type, status, worker_id, "
                "attempts, claimed_at) VALUES (:u, :j, 'processing', :w, 1, NOW()) RETURNING id"
            ),
            {"u": owner_id, "j": jt, "w": worker_id_unique},
        )
    ).scalar_one()
    await db_session.commit()

    r = await client.post(
        "/internal/worker/result",
        json={
            "job_id": job_id,
            "worker_id": worker_id_unique,
            "status": "completed",
            "result": {"summary": "test"},
        },
        headers=_auth(worker_token),
    )
    assert r.status_code == 200
    res = await db_session.execute(
        text("SELECT status, result FROM worker_jobs WHERE id = :i"), {"i": job_id}
    )
    s, result_jsonb = res.first()
    assert s == "completed"
    assert result_jsonb == {"summary": "test"}


@pytest.mark.asyncio
async def test_drain_heartbeat_insert(
    client, db_session, worker_token, worker_id_unique, owner_id, cleanup_after
):
    await db_session.execute(
        text(
            "INSERT INTO worker_capabilities (worker_id, user_id, device_label, "
            "worker_class, tier) VALUES (:w, :u, 'lbl', 'laptop', 'laptop_small') "
            "ON CONFLICT (worker_id) DO NOTHING"
        ),
        {"w": worker_id_unique, "u": owner_id},
    )
    await db_session.commit()

    r = await client.post(
        "/internal/worker/drain",
        json={"worker_id": worker_id_unique, "reason": "sleep"},
        headers=_auth(worker_token),
    )
    assert r.status_code == 200
    res = await db_session.execute(
        text(
            "SELECT status, raw_payload FROM worker_heartbeats WHERE worker_id = :w "
            "ORDER BY heartbeat_at DESC LIMIT 1"
        ),
        {"w": worker_id_unique},
    )
    s, rp = res.first()
    assert s == "draining"
    assert rp == {"reason": "sleep"}