# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""
Online serving tests: text-to-image.
See examples/online_serving/text_to_image/README.md

The output materials are organized in a three-level directory structure:
- Set at init: `self.output_root` for all tests (from env OUTPUT_DIR)
- Set at `self.run(...)`: `output_subfolder` for a specific example page (e.g., `example_offline_t2i`)
- Generated by `extract_readme_snippets`: `snippet.test_id` for a specific code block (matching H2 titles, e.g., `basic_usage_001`)
"""

import sys
from pathlib import Path

import pytest

from tests.examples.helpers import EXAMPLES, OUTPUT_DIR, run_cmd, write_zimage_lora
from tests.helpers.assertions import assert_image_valid
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams

pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": ["H100", "B200"]})]

T2I_ONLINE_CLIENT = EXAMPLES / "online_serving" / "text_to_image" / "openai_chat_client.py"
EXAMPLE_OUTPUT_SUBFOLDER = "example_online_t2i"


@pytest.fixture(scope="module")
def example_output_dir() -> Path:
    d = OUTPUT_DIR / EXAMPLE_OUTPUT_SUBFOLDER
    d.mkdir(parents=True, exist_ok=True)
    return d


# Ensuring two indirect parametrization receive exactly the same parameter list.
# Then, only one omni_server instance is started for all test functions using this param set.
# Two additional pre-requisites:
# - omni_server must (and indeed is) defined as a module-scoped fixture
# - test functions sharing the same param set must be adjacent to each other
qwen_image_server_params = [OmniServerParams(model="Qwen/Qwen-Image")]
z_image_turbo_server_params = [OmniServerParams(model="Tongyi-MAI/Z-Image-Turbo")]


# --- ### Method 1: Using curl ---


@pytest.mark.parametrize("omni_server", qwen_image_server_params, indirect=True)
def test_api_calls_001(omni_server: OmniServer, example_output_dir: Path):
    url = f"http://{omni_server.host}:{omni_server.port}/v1/chat/completions"
    case_dir = example_output_dir / "api_calls-001"
    case_dir.mkdir(parents=True, exist_ok=True)
    out = case_dir / "api_calls_001.png"
    json_str = """{
    "messages": [
      {"role": "user", "content": "A beautiful landscape painting"}
    ],
    "extra_body": {
      "height": 1024,
      "width": 1024,
      "num_inference_steps": 50,
      "true_cfg_scale": 4.0,
      "seed": 42
    }
}"""
    run_cmd(
        f"curl -s '{url}'"
        " -H 'Content-Type: application/json'"
        f" -d '{json_str}'"
        " | jq -r '.choices[0].message.content[0].image_url.url'"
        f" | cut -d',' -f2- | base64 -d > '{out}'",
        shell=True,
    )
    assert_image_valid(out, width=1024, height=1024)


# --- ### Method 2: Using Python Client ---


@pytest.mark.parametrize("omni_server", qwen_image_server_params, indirect=True)
def test_api_calls_002(omni_server: OmniServer, example_output_dir: Path):
    case_dir = example_output_dir / "api_calls-002"
    case_dir.mkdir(parents=True, exist_ok=True)
    out = case_dir / "api_calls_002.png"
    run_cmd(
        [
            sys.executable,
            str(T2I_ONLINE_CLIENT),
            "--prompt",
            "A beautiful landscape painting",
            "--output",
            str(out),
            "--server",
            f"http://{omni_server.host}:{omni_server.port}",
        ]
    )
    assert_image_valid(out)


@pytest.mark.skip("README section 'Method 3: Using Gradio Demo' is intentionally excluded for examples tests")
def test_api_calls_003(): ...


# --- ### Using Python Client with LoRA ---


@pytest.mark.parametrize("omni_server", z_image_turbo_server_params, indirect=True)
def test_lora_001(omni_server: OmniServer, example_output_dir: Path, tmp_path: Path):
    lora_dir = tmp_path / "zimage_lora_a"
    write_zimage_lora(lora_dir, v_scale=8.0)
    case_dir = example_output_dir / "lora-001"
    case_dir.mkdir(parents=True, exist_ok=True)
    out = case_dir / "lora_001.png"
    run_cmd(
        [
            sys.executable,
            str(T2I_ONLINE_CLIENT),
            "--prompt",
            "A piece of cheesecake",
            "--lora-path",
            str(lora_dir),
            "--lora-name",
            "a",
            "--lora-scale",
            "1.0",
            "--output",
            str(out),
            "--server",
            f"http://{omni_server.host}:{omni_server.port}",
        ]
    )
    assert_image_valid(out)


# --- ### Using curl with LoRA (Images API) ---


@pytest.mark.skip(reason="Covered by tests/diffusion/lora/test_images_generations_lora.py")
def test_lora_002(): ...