#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
=============================================================================
THỬ NGHIỆM SINH VIDEO I2V BẰNG HUGGING FACE SPACE (ZERO-GPU FREE API)
=============================================================================
Ví dụ sử dụng:
  # 1. Thử nghiệm nhanh với LTX-Video (Mặc định, tốc độ ~20s):
  python gradio_version/render_test.py --image "04_Assets/characters/thuy_kieu_maiden_16yo/front.png" --prompt "cinematic slow portrait shot, ancient Vietnamese dress, gentle breeze blowing hair, soft natural lighting, 24fps"

  # 2. Thử nghiệm với Wan 2.1 (Alibaba Wan-Video I2V):
  python gradio_version/render_test.py --image "04_Assets/characters/thuy_kieu_maiden_16yo/front.png" --prompt "gentle breeze, historical cinema style" --model wan21

  # 3. Thử nghiệm theo Shot ID trong phim (Tự động nạp Prompt và Start Frame):
  python gradio_version/render_test.py --shot-id ep01_scene01_shot01 --model ltx
=============================================================================
"""

import sys
import json
import argparse
from pathlib import Path

# Đảm bảo UTF-8 Console
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

# Thêm thư mục gradio_version vào sys.path
CURRENT_DIR = Path(__file__).resolve().parent
PROJECT_DIR = CURRENT_DIR.parent
if str(PROJECT_DIR) not in sys.path:
    sys.path.insert(0, str(PROJECT_DIR))

from gradio_version.hf_config import (
    BASE_DIR,
    SUPPORTED_SPACES,
    DEFAULT_MODEL,
    OUTPUTS_DIR,
    ASSETS_VIDEOS_DIR,
    ASSETS_KEYFRAMES_DIR
)
from gradio_version.hf_client import HuggingFaceVideoClient


def load_shot_metadata(shot_id: str):
    """Tìm thông tin prompt và start frame từ danh mục dự án."""
    prompts_file = BASE_DIR / "02_AI_Prompts" / "muse_ai_video_prompts.json"
    prompt_text = "cinematic slow motion, ancient historical drama, traditional Vietnamese costume, atmospheric lighting, 24fps"
    image_path = None

    if prompts_file.exists():
        try:
            with open(prompts_file, "r", encoding="utf-8") as f:
                data = json.load(f)
                shots = data.get("shots", {})
                if shot_id in shots:
                    shot_info = shots[shot_id]
                    prompt_text = shot_info.get("prompt", prompt_text)
        except Exception as e:
            print(f"[!] Không đọc được file prompt: {e}")

    # Tìm Start Frame tiềm năng
    kf_candidates = [
        ASSETS_KEYFRAMES_DIR / f"{shot_id}_start.png",
        ASSETS_KEYFRAMES_DIR / f"{shot_id}_start.jpg",
        ASSETS_KEYFRAMES_DIR / shot_id / "clean_frame_0.jpg",
        ASSETS_KEYFRAMES_DIR / shot_id / "start_frame.png",
        BASE_DIR / "04_Assets" / "characters" / "thuy_kieu_maiden_16yo" / "front.png"
    ]
    for c in kf_candidates:
        if c.exists():
            image_path = str(c)
            break

    return prompt_text, image_path


def main():
    parser = argparse.ArgumentParser(description="HF ZeroGPU I2V Test Runner")
    parser.add_argument("--image", type=str, help="Đường dẫn tới ảnh đầu vào (Start Frame)")
    parser.add_argument("--prompt", type=str, help="Prompt mô tả chuyển động")
    parser.add_argument("--model", type=str, default=DEFAULT_MODEL, choices=list(SUPPORTED_SPACES.keys()),
                        help="Lựa chọn mô hình: ltx (mặc định), wan21, cogvideox")
    parser.add_argument("--shot-id", type=str, help="Mã Shot ID (ep01_scene01_shot01...) để tự lấy dữ liệu từ kịch bản")
    parser.add_argument("--duration", type=float, default=4.0, help="Thời lượng video (giây)")
    parser.add_argument("--tail-frame", action="store_true", default=True, help="Tự động trích xuất Tail Frame sau khi sinh")
    parser.add_argument("--token", type=str, help="Hugging Face User Access Token (tùy chọn)")

    args = parser.parse_args()

    image_path = args.image
    prompt_text = args.prompt

    if args.shot_id:
        p, img = load_shot_metadata(args.shot_id)
        if not prompt_text:
            prompt_text = p
        if not image_path:
            image_path = img

    if not image_path:
        # Fallback thử tìm 1 ảnh chân dung nhân vật trong repo
        sample_img = BASE_DIR / "04_Assets" / "characters" / "thuy_kieu_maiden_16yo" / "front.png"
        if sample_img.exists():
            image_path = str(sample_img)
            print(f"[*] Sử dụng ảnh mẫu mặc định: {sample_img}")
        else:
            print("[!] Lỗi: Bạn cần chỉ định ảnh bằng --image hoặc --shot-id.")
            sys.exit(1)

    if not prompt_text:
        prompt_text = "cinematic camera movement, soft atmospheric wind, detailed clothing embroidery, high quality 24fps"

    client = HuggingFaceVideoClient(token=args.token)
    out_name = f"{args.model}_{args.shot_id}.mp4" if args.shot_id else None

    success, video_file = client.generate_i2v(
        image_path=image_path,
        prompt=prompt_text,
        model_key=args.model,
        duration=args.duration,
        output_filename=out_name
    )

    if success and video_file:
        print(f"\n[🎉 THÀNH CÔNG] Video đã được tạo tại: {video_file}")
        if args.tail_frame:
            tail = HuggingFaceVideoClient.extract_tail_frame(video_file)
            if tail:
                print(f"   • Tail Frame kế thừa: {tail}")
    else:
        print("\n[❌ THẤT BẠI] Quá trình sinh video không hoàn tất.")
        sys.exit(1)


if __name__ == "__main__":
    main()