import os
import sys
import subprocess
import gradio as gr
from huggingface_hub import list_repo_files, hf_hub_download

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")

# Specific file to include from transformers/
ALLOWED_TRANSFORMER_FILE = "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"

def list_ckpt_files():
    all_files = list_repo_files(repo_id=MODEL_REPO, repo_type="model", revision="main")

    filtered_files = []
    for f in all_files:
        # Only allow this one file from transformers/
        if f.startswith("ckpts/hunyuan-video-t2v-720p/transformers/"):
            if f == ALLOWED_TRANSFORMER_FILE:
                filtered_files.append(f)
        else:
            filtered_files.append(f)
    return filtered_files

def download_ckpts():
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    logs = []

    for filepath in list_ckpt_files():
        relative_path = os.path.relpath(filepath, "ckpts")
        target_path = os.path.join(WEIGHTS_DIR, "ckpts", relative_path)

        if os.path.exists(target_path):
            logs.append(f"✅ Exists: {target_path}")
            continue

        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        logs.append(f"⬇️ Downloading: {filepath}")
        try:
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=filepath,
                local_dir=os.path.dirname(target_path),
                local_dir_use_symlinks=False
            )
        except Exception as e:
            logs.append(f"❌ Failed to download {filepath}: {e}")
    return "\n".join(logs)

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):
        return f"❌ Missing checkpoint: {checkpoint_fp8}"

    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"

    result = subprocess.run(cmd, env=env, capture_output=True, text=True)
    if result.returncode != 0:
        return f"❌ sample_gpu_poor.py failed:\n{result.stderr}"
    return f"✅ sample_gpu_poor.py finished:\n{result.stdout}"

def download_and_run():
    log1 = download_ckpts()
    log2 = run_sample_gpu_poor()
    return f"{log1}\n\n---\n\n{log2}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📦 Download All Checkpoints (Except Only One File from Transformers)")
    output = gr.Textbox(lines=30, label="Logs")
    button = gr.Button("🚀 Download + Run")
    button.click(fn=download_and_run, outputs=output)

if __name__ == "__main__":
    demo.launch(share=False)