"""Gemini Batch API. List and bogus-name probes always run; creating a batch is gated by RUN_BATCH_TESTS=true
(marker run_batch_tests) and xfails on FAILED_PRECONDITION (Batch API is not available on the free tier — observed 2026-09-18)."""
from __future__ import annotations

import time
import pytest

TERMINAL = {"BATCH_STATE_SUCCEEDED", "BATCH_STATE_FAILED", "BATCH_STATE_CANCELLED", "BATCH_STATE_EXPIRED"}


def test_list_batches(gemini):
    st, body, _ = gemini("GET", "/v1beta/batches?pageSize=5", note="test_batch list")
    assert st == 200 and isinstance(body, dict)
    for op in body.get("operations", []):
        assert op["name"].startswith("batches/") and "done" in op


def test_get_bogus_batch(gemini):
    st, body, _ = gemini("GET", "/v1beta/batches/bogus-batch", note="test_batch bogus")
    assert st in (400, 403, 404) and "error" in body


@pytest.mark.run_batch_tests
def test_inline_batch_lifecycle(gemini, models):
    req = {"request": {"contents": [{"parts": [{"text": "Reply with OK."}]}], "generationConfig": {"maxOutputTokens": 8}}, "metadata": {"key": "r1"}}
    body = {"batch": {"display_name": "atlas-test", "input_config": {"requests": {"requests": [req]}}}}
    st, op, _ = gemini("POST", f"/v1beta/models/{models['gemini']}:batchGenerateContent", body, est_cost_usd=0.00005, note="test_batch create")
    if st == 400 and op["error"]["status"] == "FAILED_PRECONDITION":
        pytest.xfail("Batch API not enabled for this key (no free tier): " + op["error"]["message"])
    assert st == 200 and op["name"].startswith("batches/"), op
    assert op["metadata"]["state"] == "BATCH_STATE_PENDING"
    name = op["name"]
    try:
        t0 = time.time()
        state = op["metadata"]["state"]
        while not op.get("done") and time.time() - t0 < 300:
            time.sleep(20)
            st, op, _ = gemini("GET", f"/v1beta/{name}", note="test_batch poll")
            assert st == 200
            state = op["metadata"]["state"]
        if not op.get("done"):
            st, _, _ = gemini("POST", f"/v1beta/{name}:cancel", note="test_batch cancel")
            assert st == 200
        else:
            assert state in TERMINAL
            if state == "BATCH_STATE_SUCCEEDED":
                resp = op["response"]
                assert "inlinedResponses" in resp or "responsesFile" in resp
                if "inlinedResponses" in resp:
                    r0 = resp["inlinedResponses"]["inlinedResponses"][0]
                    assert r0["metadata"]["key"] == "r1" and ("response" in r0 or "error" in r0)
    finally:
        st, _, _ = gemini("DELETE", f"/v1beta/{name}", note="test_batch delete")
        assert st == 200