"""Phase 2 component test: Wan2.2 pipeline + LoRA stacking.

Verifies:
- ``Wan22Pipeline`` loads successfully (exercises the real LightX2V
  set_config -> init_runner flow).
- ``load_loras`` / ``unload_loras`` survive with the two user LoRAs at
  ``/cache/loras/wan22-[HL]-e8.safetensors``.

Supports both fp8 and GGUF DIT quantisation.  Set the ``DIT_QUANT``
environment variable to switch (default: ``fp8-sgl``).

    DIT_QUANT=gguf-Q4_K_M docker compose exec voice-chat \
        python -m tests.component.test_02_wan22_loras

Requires GPU and a first-run download of both HF repos (base support files
~12 GB, DIT size depends on quant — fp8 ~30 GB, GGUF Q4_K_M ~19 GB).
If LightX2V isn't installed the test is skipped.

Run (default fp8):
    docker compose exec voice-chat python -m tests.component.test_02_wan22_loras
"""
from __future__ import annotations

import os
import sys

from tests.component._common import get_logger

log = get_logger("test_02")

# --- Quant-dependent defaults ------------------------------------------------

DIT_QUANT = os.environ.get("DIT_QUANT", "fp8-sgl")

if DIT_QUANT.startswith("gguf-"):
    CONFIG_JSON = "/app/configs/lightx2v/wan22_i2v_gguf_distill.json"
    DIT_REPO = "QuantStack/Wan2.2-I2V-A14B-GGUF"
else:
    CONFIG_JSON = "/app/configs/lightx2v/wan22_i2v_fp8_distill.json"
    DIT_REPO = "lightx2v/Wan2.2-Distill-Models"

LORA_HIGH = "/cache/loras/wan22-H-e8.safetensors"
LORA_LOW = "/cache/loras/wan22-L-e8.safetensors"


def run():
    try:
        from server.video_models.wan22 import Wan22Pipeline
    except ImportError as e:
        log.error("Wan22Pipeline import failed: %s", e)
        log.warning("SKIP: phase 2 deps not installed")
        sys.exit(0)

    from server.video import LoRASpec

    log.info("[case 1] Instantiate Wan22Pipeline "
             "(quant=%s, dit_repo=%s)...", DIT_QUANT, DIT_REPO)
    try:
        pipe = Wan22Pipeline(
            base_repo="Wan-AI/Wan2.2-I2V-A14B",
            dit_repo=DIT_REPO,
            config_json=CONFIG_JSON,
            model_cls="wan2.2_moe_distill",
            resolution=480,
            fps=16,
            dit_quant_scheme=DIT_QUANT,
        )
    except Exception as e:
        log.error("FAIL: Wan22Pipeline construction raised: %s", e)
        log.error("Check: LightX2V install, HF cache at /cache/huggingface, "
                  "VRAM headroom, and that %s exists inside the container.",
                  CONFIG_JSON)
        sys.exit(2)
    log.info("  PASS: pipeline constructed")

    # --- LoRAs ---
    log.info("[case 2] load_loras with empty list -> no-op")
    pipe.load_loras([])
    log.info("  PASS")

    if not (os.path.isfile(LORA_HIGH) and os.path.isfile(LORA_LOW)):
        log.warning("SKIP: expected LoRA files not found at %s / %s",
                    LORA_HIGH, LORA_LOW)
        log.info("ALL PASSED (partial — LoRA cases skipped)")
        return

    log.info("[case 3] load_loras with the two MoE distill LoRAs")
    specs = [
        LoRASpec(
            path=LORA_HIGH,
            weight=1.0,
            target="high_noise",
            name="wan22-H-e8",
        ),
        LoRASpec(
            path=LORA_LOW,
            weight=1.0,
            target="low_noise",
            name="wan22-L-e8",
        ),
    ]
    try:
        pipe.load_loras(specs)
    except Exception as e:
        log.error("FAIL: load_loras raised: %s", e)
        log.error("Check: switch_lora support for wan2.2_moe_distill in the "
                  "installed LightX2V build. If it errors there, pre-declare "
                  "LoRAs in the config_json 'lora_configs' field instead.")
        sys.exit(3)
    log.info("  PASS: LoRAs applied")

    log.info("[case 4] unload_loras")
    try:
        pipe.unload_loras()
    except Exception as e:
        log.error("FAIL: unload_loras raised: %s", e)
        sys.exit(4)
    log.info("  PASS")

    log.info("ALL PASSED")


if __name__ == "__main__":
    run()