#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Test runner for Wan 2.2 5B on Hugging Face ZeroGPU for KieuStory cinema shots.
"""

import sys
import os
import json
import time
from pathlib import Path

# Ensure UTF-8 output on Windows
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

BASE_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BASE_DIR))

from gradio_version.hf_client import HuggingFaceVideoClient

def test_shot_wan22(
    shot_id: str = "ep01_scene02_shot04",
    image_path: str = "04_Assets/characters/02_Vuong_Family_And_Fate/vuong_ong_55yo/views/vuong_ong_55yo_portrait.png",
    prompt: str = "Trung cận Vương Ông (55 tuổi) đạo mạo trong trang phục áo dài lụa màu sẫm, nét mặt đôn hậu, ngồi bên bàn trà gỗ trắc chậm rãi nâng chén trà sứ hoa lam, phong thái Nho gia điềm tĩnh, 24fps cinematic."
):
    print("=" * 70)
    print(f"🎬 [TEST WAN 2.2 5B] Bắt đầu render shot: {shot_id}")
    print(f"   • Input Image: {image_path}")
    print(f"   • Prompt: {prompt[:80]}...")
    print("=" * 70)

    resolved_img = BASE_DIR / image_path
    if not resolved_img.exists():
        print(f"[!] Lỗi: Không tìm thấy ảnh đầu vào tại {resolved_img}")
        return False

    client = HuggingFaceVideoClient()
    output_filename = f"{shot_id}_wan22_5b.mp4"
    
    t0 = time.time()
    ok, video_path = client.generate_i2v(
        image_path=str(resolved_img),
        prompt=prompt,
        model_key="wan22",
        width=1280,
        height=704,
        output_filename=output_filename
    )
    elapsed = time.time() - t0

    if ok and video_path and Path(video_path).exists():
        print(f"\n[✓] Render Wan 2.2 5B thành công trong {elapsed:.1f}s!")
        print(f"    File: {video_path}")
        
        # Trích xuất tail frame
        tail_dir = BASE_DIR / "04_Assets" / "keyframes" / shot_id
        tail_dir.mkdir(parents=True, exist_ok=True)
        tail_file = tail_dir / "clean_frame_239.jpg"
        HuggingFaceVideoClient.extract_tail_frame(str(video_path), str(tail_file))
        print(f"    Tail Frame: {tail_file}")
        return True
    else:
        print(f"\n[!] Thất bại khi sinh video với Wan 2.2 5B.")
        return False

if __name__ == "__main__":
    shot = sys.argv[1] if len(sys.argv) > 1 else "ep01_scene02_shot04"
    img = sys.argv[2] if len(sys.argv) > 2 else "04_Assets/characters/02_Vuong_Family_And_Fate/vuong_ong_55yo/views/vuong_ong_55yo_portrait.png"
    p = sys.argv[3] if len(sys.argv) > 3 else "Trung cận Vương Ông (55 tuổi) đạo mạo trong trang phục áo dài lụa màu sẫm, nét mặt đôn hậu, ngồi bên bàn trà gỗ trắc chậm rãi nâng chén trà sứ hoa lam, phong thái Nho gia điềm tĩnh, 24fps cinematic."
    test_shot_wan22(shot, img, p)