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

from __future__ import annotations

import base64
import copy
import io
import logging
import os
import tempfile
from pathlib import Path
from typing import Any

import pytest
import requests
import yaml
from PIL import Image

from tests.e2e.accuracy.helpers import (
    SimilarityThresholds,
    assert_images_pixel_close,
    assert_similarity,
    model_output_dir,
    resolve_device_profile,
    resolve_similarity_thresholds,
)
from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServer

pytestmark = [pytest.mark.diffusion]

logger = logging.getLogger(__name__)

MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
SEED = 42
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 2.5
HEIGHT = 1024
WIDTH = 1024
PROMPT = "A brown and white dog is running on the grass."
MEAN_THRESHOLD = 3e-2
P99_THRESHOLD = 3e-1
PSNR_THRESHOLD_NPU = 26.0

# Per-device SSIM/PSNR for online/offline vs baseline. Unlisted devices use ``default``.
SIMILARITY_THRESHOLDS_BY_DEVICE: dict[str, SimilarityThresholds] = {
    "B200": SimilarityThresholds(ssim=0.96, psnr=28.0),
    "default": SimilarityThresholds(ssim=0.97, psnr=30.0),
}


def _psnr_threshold(thresholds: SimilarityThresholds | None = None) -> float:
    from vllm_omni.platforms import current_omni_platform

    if current_omni_platform.is_npu():
        return PSNR_THRESHOLD_NPU
    if thresholds is None:
        thresholds = resolve_similarity_thresholds(SIMILARITY_THRESHOLDS_BY_DEVICE)
    return thresholds.psnr


_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent


def _baseline_path() -> Path:
    from vllm_omni.platforms import current_omni_platform

    suffix = "_npu" if current_omni_platform.is_npu() else ""
    path = get_asset_path(f"hunyuan/hunyuan_baseline{suffix}.png")
    assert path.exists(), f"Baseline image not found at {path}"
    return path


_OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"

# DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
_DEPLOY_CONFIG: dict[str, Any] = {
    "pipeline": "hunyuan_image3_dit",
    "async_chunk": False,
    "trust_remote_code": True,
    "stages": [
        {
            "stage_id": 0,
            "max_num_seqs": 1,
            "gpu_memory_utilization": 0.9,
            "enforce_eager": True,
            "trust_remote_code": True,
            "devices": "0,1,2,3",  # set dynamically by _write_deploy_config
            "vae_use_slicing": False,
            "moe_backend": "flashinfer_cutlass",
            "vae_use_tiling": False,
            "parallel_config": {
                "pipeline_parallel_size": 1,
                "data_parallel_size": 1,
                "tensor_parallel_size": 1,
                "enable_expert_parallel": True,
                "sequence_parallel_size": 1,
                "ulysses_degree": 1,
                "ring_degree": 1,
                "allgather_degree": 4,
                "cfg_parallel_size": 1,
                "vae_patch_parallel_size": 1,
                "use_hsdp": False,
                "hsdp_shard_size": -1,
                "hsdp_replicate_size": 1,
            },
            "default_sampling_params": {
                "seed": SEED,
            },
        },
    ],
    "platforms": {
        "npu": {
            "stages": [
                {
                    "stage_id": 0,
                    "gpu_memory_utilization": 0.65,
                    "devices": "0,1,2,3",
                    "moe_backend": "auto",
                    "max_num_batched_tokens": 32768,
                    "parallel_config": {
                        "tensor_parallel_size": 4,
                        "enable_expert_parallel": True,
                    },
                },
            ],
        },
    },
}


def _model_name() -> str:
    return os.environ.get("HUNYUAN_IMAGE3_MODEL", MODEL_NAME)


def _devices() -> str:
    return os.environ.get("HUNYUAN_IMAGE3_DEVICES", "0,1,2,3")


def _write_deploy_config(path: Path) -> None:
    config: dict[str, Any] = copy.deepcopy(_DEPLOY_CONFIG)
    devices = _devices()
    num_devices = len(devices.split(","))
    config["stages"][0]["devices"] = devices
    parallel_config = config["stages"][0]["parallel_config"]
    if parallel_config.get("allgather_degree", 1) > 1:
        parallel_config["allgather_degree"] = num_devices
    elif parallel_config.get("tensor_parallel_size", 1) > 1:
        parallel_config["tensor_parallel_size"] = num_devices

    sequence_parallel_size = (
        parallel_config.get("allgather_degree", 1)
        if parallel_config.get("allgather_degree", 1) > 1
        else parallel_config.get("ulysses_degree", 1) * parallel_config.get("ring_degree", 1)
    )
    configured_world_size = (
        parallel_config.get("pipeline_parallel_size", 1)
        * parallel_config.get("data_parallel_size", 1)
        * parallel_config.get("tensor_parallel_size", 1)
        * sequence_parallel_size
        * parallel_config.get("cfg_parallel_size", 1)
    )
    assert configured_world_size == num_devices, (
        f"Configured parallel world size ({configured_world_size}) does not match "
        f"the number of devices ({num_devices}): {parallel_config}"
    )

    parallel_config["sequence_parallel_size"] = sequence_parallel_size
    config_text = yaml.dump(config, default_flow_style=False, sort_keys=False)
    path.write_text(config_text)
    logger.info(
        "HunyuanImage3 pixel accuracy deploy config:\n"
        "  path=%s\n"
        "  CUDA_VISIBLE_DEVICES=%s\n"
        "  devices=%s\n"
        "  tensor_parallel_size=%s\n"
        "  allgather_degree=%s\n%s",
        path,
        os.environ.get("CUDA_VISIBLE_DEVICES", "<unset>"),
        devices,
        parallel_config.get("tensor_parallel_size", 1),
        parallel_config.get("allgather_degree", 1),
        config_text,
    )


def _run_vllm_omni_hunyuan_image3_online(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:
    server_args = [
        "--deploy-config",
        deploy_config,
        "--stage-init-timeout",
        "300",
        "--init-timeout",
        "900",
        "--enforce-eager",
        "--trust-remote-code",
    ]
    with OmniServer(model, server_args, use_omni=True) as omni_server:
        response = requests.post(
            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
            json={
                "model": omni_server.model,
                "prompt": PROMPT,
                "size": f"{WIDTH}x{HEIGHT}",
                "n": 1,
                "response_format": "b64_json",
                "num_inference_steps": NUM_INFERENCE_STEPS,
                "guidance_scale": GUIDANCE_SCALE,
                "seed": SEED,
                "bot_task": "none",
                "use_system_prompt": "en_unified",
            },
            timeout=600,
        )
        response.raise_for_status()
        payload = response.json()
        assert len(payload["data"]) == 1
        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        image.load()
        image.save(output_path)
        return image


def _run_vllm_omni_hunyuan_image3_offline(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:
    import subprocess
    import sys

    output_dir = str(output_path.parent)
    subprocess.run(
        [
            sys.executable,
            str(_OFFLINE_SCRIPT),
            "--modality",
            "text2img",
            "--deploy-config",
            deploy_config,
            "--prompts",
            PROMPT,
            "--output",
            output_dir,
            "--steps",
            str(NUM_INFERENCE_STEPS),
            "--guidance-scale",
            str(GUIDANCE_SCALE),
            "--seed",
            str(SEED),
            "--height",
            str(HEIGHT),
            "--width",
            str(WIDTH),
            "--bot-task",
            "none",
            "--sys-type",
            "en_unified",
            "--model",
            model,
            "--enforce-eager",
        ],
        check=True,
    )
    images = sorted(Path(output_dir).glob("output_*.png"))
    assert images, f"No output image found in {output_dir}"
    image = Image.open(images[0]).convert("RGB")
    image.load()
    image.save(output_path)
    return image


def _assert_against_baseline(image: Image.Image, label: str) -> None:
    baseline_path = _baseline_path()
    baseline_image = Image.open(baseline_path).convert("RGB")
    device_profile = resolve_device_profile(profiles=SIMILARITY_THRESHOLDS_BY_DEVICE)
    thresholds = resolve_similarity_thresholds(SIMILARITY_THRESHOLDS_BY_DEVICE, device_profile)
    psnr_threshold = _psnr_threshold(thresholds)
    logger.info(
        "HunyuanImage3 pixel accuracy thresholds: profile=%s ssim>=%s psnr>=%s",
        device_profile,
        thresholds.ssim,
        psnr_threshold,
    )

    assert_images_pixel_close(
        model_name=f"{MODEL_NAME} ({label} vs baseline)",
        vllm_image=image,
        diffusers_image=baseline_image,
        mean_threshold=MEAN_THRESHOLD,
        p99_threshold=P99_THRESHOLD,
    )
    assert_similarity(
        model_name=f"{MODEL_NAME} ({label} vs baseline)",
        vllm_image=image,
        diffusers_image=baseline_image,
        ssim_threshold=thresholds.ssim,
        psnr_threshold=psnr_threshold,
    )


@pytest.mark.full_model
@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=4)
def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None:
    model = _model_name()
    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)

    with tempfile.TemporaryDirectory() as tmpdir:
        deploy_config_path = Path(tmpdir) / "deploy.yaml"
        _write_deploy_config(deploy_config_path)
        image = _run_vllm_omni_hunyuan_image3_online(
            model=model, deploy_config=str(deploy_config_path), output_path=output_dir / "vllm_omni_online.png"
        )
    _assert_against_baseline(image, "online")


@pytest.mark.full_model
@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=4)
def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None:
    model = _model_name()
    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)

    with tempfile.TemporaryDirectory() as tmpdir:
        deploy_config_path = Path(tmpdir) / "deploy.yaml"
        _write_deploy_config(deploy_config_path)
        image = _run_vllm_omni_hunyuan_image3_offline(
            model=model, deploy_config=str(deploy_config_path), output_path=output_dir / "vllm_omni_offline.png"
        )
    _assert_against_baseline(image, "offline")