"""Quick smoke test: generate a video clip with the GGUF pipeline.

Calls Wan22Pipeline.generate_i2v directly (no MuseTalk, no VideoEngine)
and writes the result to tests/component/_out/phase9_gguf.mp4.

Run:
    docker compose exec -e DIT_QUANT=gguf-Q4_K_M voice-chat \
        python -m tests.component.test_09_gguf_generate
"""
from __future__ import annotations

import os
import sys

from tests.component._common import ensure_sample_avatar, get_logger, write_bytes

log = get_logger("test_09")

DIT_QUANT = os.environ.get("DIT_QUANT", "gguf-Q4_K_M")

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"


def run():
    try:
        from server.video_models.wan22 import Wan22Pipeline
    except ImportError as e:
        log.error("Import failed: %s", e)
        sys.exit(0)

    avatar = ensure_sample_avatar()
    log.info("Avatar: %s", avatar)

    log.info("Building pipeline (quant=%s)...", DIT_QUANT)
    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,
        t5_quantized=True,
    )
    log.info("Pipeline ready.")

    # Debug: verify DTYPE is set correctly for GGUF
    from lightx2v.utils.envs import GET_DTYPE
    log.info("GET_DTYPE() = %s (DTYPE env = %s)", GET_DTYPE(), os.environ.get("DTYPE"))

    log.info("Generating 3-second i2v clip...")
    frames = pipe.generate_i2v(
        image_path=avatar,
        prompt="a person looking at the camera, natural lighting, soft focus background",
        seconds=3,
        seed=42,
    )
    log.info("Got frames: shape=%s dtype=%s", frames.shape, frames.dtype)

    # Encode to MP4
    import imageio.v3 as iio
    import tempfile

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tf:
        tmp = tf.name
    try:
        iio.imwrite(tmp, frames, fps=16, codec="libx264")
        with open(tmp, "rb") as f:
            mp4_bytes = f.read()
    finally:
        os.remove(tmp)

    out = write_bytes("phase9_gguf.mp4", mp4_bytes)
    log.info("PASS: video written to %s (%d bytes, %d frames)", out, len(mp4_bytes), frames.shape[0])


if __name__ == "__main__":
    run()