import os
import sys
import subprocess
from huggingface_hub import list_repo_files

MODEL_REPO = "tencent/HunyuanVideo-Avatar"
REPO_ROOT_FOLDER = "ckpts"  # folder inside repo to download

BASE_DIR = os.getcwd()
WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")

def list_ckpt_files():
    print(f"🔍 Listing files under '{REPO_ROOT_FOLDER}' in repo {MODEL_REPO}...")
    all_files = list_repo_files(repo_id=MODEL_REPO, repo_type="model", revision="main")
    files = [f for f in all_files if f.startswith(REPO_ROOT_FOLDER + "/")]
    return files

def download_ckpts():
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    for filepath in list_ckpt_files():
        # Target path relative to weights directory
        relative_path = os.path.relpath(filepath, REPO_ROOT_FOLDER)
        target_path = os.path.join(WEIGHTS_DIR, REPO_ROOT_FOLDER, relative_path)

        # Skip if already exists
        if os.path.exists(target_path):
            print(f"✅ File exists: {target_path}")
            continue

        # Create dirs if needed
        os.makedirs(os.path.dirname(target_path), exist_ok=True)

        # Download the file blob from HF repo using huggingface_hub api (requires v0.14+)
        print(f"⬇️ Downloading {filepath} to {target_path} ...")
        try:
            from huggingface_hub import hf_hub_download
            hf_hub_download(repo_id=MODEL_REPO, filename=filepath, local_dir=os.path.dirname(target_path), local_dir_use_symlinks=False)
        except ImportError:
            print("⚠️ hf_hub_download not available, please upgrade huggingface_hub")
            sys.exit(1)

def run_sample_gpu_poor():
    checkpoint_fp8 = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states_fp8.pt")
    if not os.path.isfile(checkpoint_fp8):
        print(f"❌ Checkpoint file not found at {checkpoint_fp8}. Cannot run sampling.")
        sys.exit(1)

    cmd = [
        "python3", "hymm_sp/sample_gpu_poor.py",
        "--input", "assets/test.csv",
        "--ckpt", checkpoint_fp8,
        "--sample-n-frames", "129",
        "--seed", "128",
        "--image-size", "704",
        "--cfg-scale", "7.5",
        "--infer-steps", "50",
        "--use-deepcache", "1",
        "--flow-shift-eval-video", "5.0",
        "--save-path", OUTPUT_BASEPATH,
        "--use-fp8",
        "--cpu-offload",
        "--infer-min"
    ]

    env = os.environ.copy()
    env["PYTHONPATH"] = "./"
    env["MODEL_BASE"] = WEIGHTS_DIR
    env["CPU_OFFLOAD"] = "1"
    env["CUDA_VISIBLE_DEVICES"] = "0"

    print("🎬 Running sample_gpu_poor.py ...")
    proc = subprocess.run(cmd, env=env)
    if proc.returncode != 0:
        print("❌ sample_gpu_poor.py failed.")
        sys.exit(1)
    print("✅ sample_gpu_poor.py completed successfully.")

def main():
    download_ckpts()
    run_sample_gpu_poor()

if __name__ == "__main__":
    main()