import os
import sys
from huggingface_hub import hf_hub_download, list_repo_files
import subprocess

MODEL_REPO = "tencent/HunyuanVideo-Avatar"
BASE_DIR = os.getcwd()
WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")


# Only download this specific file from transformers
ESSENTIAL_PATHS = [
    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
]

# Download everything from these folders
FULL_DIRS = [
    "hunyuan-video-t2v-720p/vae",
    #"llava_llama_image",
    "text_encoder_2",
    "whisper-tiny",
    "det_align",
]


def list_ckpt_files():
    """Return a list of files to download: specific files + all from some folders."""
    try:
        all_files = list_repo_files(MODEL_REPO)
    except Exception as e:
        print(f"❌ Failed to list files from {MODEL_REPO}: {e}")
        return []

    files_to_download = ESSENTIAL_PATHS.copy()
    for path in all_files:
        if any(path.startswith(folder + "/") for folder in FULL_DIRS):
            files_to_download.append(path)

    return files_to_download


def download_ckpts(files):
    """Download selected files to local directory, preserving structure."""
    for file_path in files:
        local_path = os.path.join(WEIGHTS_DIR, file_path)
        if os.path.exists(local_path):
            print(f"✅ Already exists: {local_path}")
            continue

        print(f"⬇️  Downloading: {file_path}")
        try:
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=file_path,
                local_dir=WEIGHTS_DIR,
                local_dir_use_symlinks=False,
            )
        except Exception as e:
            print(f"❌ Failed to download {file_path}: {e}")


def run_sample_gpu_poor():
    ckpt_fp8 = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states_fp8.pt")

    if not os.path.isfile(ckpt_fp8):
        print(f"❌ Missing checkpoint: {ckpt_fp8}")
        return

    cmd = [
        "python3", "hymm_sp/sample_gpu_poor.py",
        "--input", "assets/test.csv",
        "--ckpt", ckpt_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...")
    result = subprocess.run(cmd, env=env, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"❌ sample_gpu_poor.py failed:\n{result.stderr}")
    else:
        print("✅ sample_gpu_poor.py ran 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()