"""xAI Imagine images tests. Model listing and validation-error tests are free and always run; generation is gated by RUN_IMAGE_TESTS
(~$0.02 per run with grok-imagine-image)."""
from __future__ import annotations

import base64

import pytest

IMG_MODEL = "grok-imagine-image"


def test_list_image_generation_models(xai):
    st, body, _ = xai("GET", "/v1/image-generation-models", note="test_images list models (free)")
    assert st == 200, body
    ids = {m["id"]: m for m in body["models"]}
    assert {"grok-imagine-image", "grok-imagine-image-2.0"} <= set(ids)
    m = ids["grok-imagine-image"]
    assert m["object"] == "model" and m["output_modalities"] == ["image"] and m["image_price"] == 200000000  # $0.02
    tiers = {(t["quality"], t["resolution"]) for t in ids["grok-imagine-image-2.0"]["pricing"]}
    assert ("low", "1k") in tiers and ("medium", "2k") in tiers


def test_get_image_model_by_alias(xai):
    st, body, _ = xai("GET", "/v1/image-generation-models/grok-imagine-image-2026-03-02", note="test_images get by alias (free)")
    assert st == 200, body
    assert body["id"] == "grok-imagine-image" and "grok-imagine-image-2026-03-02" in body["aliases"]


def test_invalid_n_error_shape(xai):
    st, body, _ = xai("POST", "/v1/images/generations", {"model": IMG_MODEL, "prompt": "a plain white square", "n": 0}, note="test_images n=0 (free)")
    assert st == 400, body
    assert body["code"] == "invalid-argument" and "between 1 and 10" in body["error"]


def test_invalid_response_format(xai):
    st, body, _ = xai("POST", "/v1/images/generations", {"model": IMG_MODEL, "prompt": "x", "response_format": "png"}, note="test_images bad response_format (free)")
    assert st == 400 and body["code"] == "invalid-argument" and body["error"] == "Invalid format."


def test_invalid_aspect_ratio_is_422_text(xai):
    st, body, hdrs = xai("POST", "/v1/images/generations", {"model": IMG_MODEL, "prompt": "x", "aspect_ratio": "7:7"}, note="test_images bad aspect_ratio (free)")
    assert st == 422, body
    assert hdrs.get("Content-Type", "").startswith("text/plain")
    assert b"aspect_ratio: unknown variant `7:7`" in body


def test_retired_grok_2_image_slug(xai):
    st, body, _ = xai("POST", "/v1/images/generations", {"model": "grok-2-image", "prompt": "x"}, note="test_images grok-2-image retired (free)")
    assert st == 404 and body["code"] == "not-found", body


@pytest.mark.run_image_tests
def test_generate_one_image_b64(xai):
    """$0.02: one 1k 1:1 image with the cheapest Imagine model, base64 output."""
    st, body, _ = xai("POST", "/v1/images/generations", {"model": IMG_MODEL, "prompt": "a plain white square", "n": 1, "aspect_ratio": "1:1",
                                                        "resolution": "1k", "response_format": "b64_json"},
                      est_cost_usd=0.02, note="test_images generate b64 ($0.02)")
    assert st == 200, body
    item = body["data"][0]
    raw = base64.b64decode(item["b64_json"])
    assert raw[:3] == b"\xff\xd8\xff" or raw[:4] == b"\x89PNG"
    assert item.get("mime_type", "image/jpeg").startswith("image/")
    assert body["usage"]["cost_in_usd_ticks"] == 200000000


@pytest.mark.run_image_tests
def test_edit_generated_image(xai):
    """~$0.022: edit a freshly generated image via JSON data URL (input image is billed ~$0.002 extra)."""
    st, gen, _ = xai("POST", "/v1/images/generations", {"model": IMG_MODEL, "prompt": "a plain white square", "response_format": "b64_json"},
                     est_cost_usd=0.02, note="test_images edit: source image ($0.02)")
    assert st == 200, gen
    data_url = "data:image/jpeg;base64," + gen["data"][0]["b64_json"]
    st, body, _ = xai("POST", "/v1/images/edits", {"model": IMG_MODEL, "prompt": "draw a small red circle in the center", "image": {"url": data_url},
                                                  "response_format": "b64_json"}, est_cost_usd=0.022, note="test_images edit ($0.022)")
    assert st == 200, body
    assert body["data"][0]["b64_json"] and body["usage"]["cost_in_usd_ticks"] > 200000000