# grantflow/tests/integration_helpers.py

import time

from fastapi.testclient import TestClient

TERMINAL_JOB_STATUSES = {"done", "error", "pending_hitl"}


def wait_for_terminal_status(client: TestClient, job_id: str, timeout_s: float = 3.0) -> dict:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        response = client.get(f"/status/{job_id}")
        assert response.status_code == 200
        body = response.json()
        if body["status"] in TERMINAL_JOB_STATUSES:
            return body
        time.sleep(0.05)
    raise AssertionError("Timed out waiting for job completion")


def drain_hitl_to_done(
    client: TestClient,
    job_id: str,
    *,
    initial_status: dict | None = None,
    max_cycles: int = 6,
) -> dict:
    status = initial_status or wait_for_terminal_status(client, job_id)
    for _ in range(max_cycles):
        if status.get("status") == "done":
            return status
        assert status.get("status") == "pending_hitl"
        checkpoint_id = str(status.get("checkpoint_id") or "").strip()
        checkpoint_stage = str(status.get("checkpoint_stage") or "").strip().lower()
        assert checkpoint_id
        assert checkpoint_stage in {"toc", "logframe"}

        approve = client.post(
            "/hitl/approve",
            json={"checkpoint_id": checkpoint_id, "approved": True, "feedback": f"Auto approve {checkpoint_stage}"},
        )
        assert approve.status_code == 200

        resume = client.post(f"/resume/{job_id}", json={})
        assert resume.status_code == 200
        expected_resume_from = "mel" if checkpoint_stage == "toc" else "critic"
        assert resume.json()["resuming_from"] == expected_resume_from
        status = wait_for_terminal_status(client, job_id)
    raise AssertionError("HITL pipeline did not reach done within max_cycles")
