import os
import sys
import subprocess
from huggingface_hub import hf_hub_download, hf_hub_list
from pathlib import Path

MODEL_REPO = "tencent/HunyuanVideo-Avatar"
BASE_DIR = Path.cwd()
WEIGHTS_DIR = BASE_DIR / "weights"
CKPTS_DIR = WEIGHTS_DIR / "ckpts"

# The root folder in the repo you want to download files from
REPO_ROOT_FOLDER = "ckpts"

# List all files under the folder 'ckpts' in the repo (recursively)
def list_ckpt_files():
    print(f"🔍 Listing files under '{REPO_ROOT_FOLDER}' in repo {MODEL_REPO}...")
    files = hf_hub_list(repo_id=MODEL_REPO, repo_type="model", revision="main", allow_regex=f"^{REPO_ROOT_FOLDER}/.*")
    return files

# Download all files in the list into weights/ckpts preserving folder structure
def download_ckpts(files):
    print(f"⬇️ Downloading {len(files)} files under '{REPO_ROOT_FOLDER}' ...")
    for file in files:
        # Target local path preserving folder structure after 'ckpts/'
        relative_path = Path(file).relative_to(REPO_ROOT_FOLDER)
        local_path = CKPTS_DIR / relative_path
        
        # Make sure directory exists
        local_path.parent.mkdir(parents=True, exist_ok=True)

        # Download the file if not present
        if not local_path.is_file():
            print(f"Downloading {file} -> {local_path}")
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=file,
                repo_type="model",
                cache_dir=str(WEIGHTS_DIR),
                local_dir=str(WEIGHTS_DIR),
                local_dir_use_symlinks=False,
                force_filename=str(local_path.name),
                revision="main",
            )
        else:
            print(f"File {local_path} already exists, skipping download.")
    print("✅ All checkpoint files downloaded.")

def run_sample_gpu_poor():
    checkpoint_fp8 = CKPTS_DIR / "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
    if not checkpoint_fp8.is_file():
        print(f"❌ Checkpoint file {checkpoint_fp8} not found!")
        sys.exit(1)

    OUTPUT_BASEPATH = BASE_DIR / "results-poor"
    OUTPUT_BASEPATH.mkdir(parents=True, exist_ok=True)

    cmd = [
        "python3", "hymm_sp/sample_gpu_poor.py",
        "--input", "assets/test.csv",
        "--ckpt", str(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", str(OUTPUT_BASEPATH),
        "--use-fp8",
        "--cpu-offload",
        "--infer-min"
    ]

    env = os.environ.copy()
    env["PYTHONPATH"] = "./"
    env["MODEL_BASE"] = str(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():
    files = list_ckpt_files()
    if not files:
        print("❌ No checkpoint files found in repo under ckpts folder.")
        sys.exit(1)

    download_ckpts(files)
    run_sample_gpu_poor()

if __name__ == "__main__":
    main()