import huggingface_hub

# 核心修補：Monkey Patch，解決新版 huggingface_hub 缺失 cached_download 的問題
if not hasattr(huggingface_hub, 'cached_download'):
    from huggingface_hub import hf_hub_download
    huggingface_hub.cached_download = hf_hub_download

import os
import sys
import subprocess
import threading
import time
import shutil
import gradio as gr
from huggingface_hub import snapshot_download, HfApi
from omegaconf import OmegaConf
import torch

cuda_available = torch.cuda.is_available()
print(f"Is CUDA available: {cuda_available}")

CHAMP_MASTER_DIR = "champ-master"
# 🌟 隔離區宣告：前處理專用房間的路徑
VENV_DIR = os.path.abspath("preproc_env")
VENV_PYTHON = os.path.join(VENV_DIR, "bin", "python")

# =========================================================================
# 🏗️ [Runtime 環境防禦與隔離區架設]
# =========================================================================
print("\n🏗️ [環境建置] 開始架設前處理隔離特區...")

# 1. 從模型庫同步組員的前處理專案
checkpoint_script = os.path.join(CHAMP_MASTER_DIR, "process_single_video.py")
if not os.path.exists(checkpoint_script):
    print(f"⏳ 偵測到 Space 本地缺乏前處理，正在從模型庫同步中...")
    try:
        os.makedirs(CHAMP_MASTER_DIR, exist_ok=True)
        snapshot_download(
            repo_id="yoyozs11/champ_preprocess",  
            local_dir=CHAMP_MASTER_DIR,       
            repo_type="model"
        )
        print("✅ 前處理模組同步成功！")
    except Exception as e:
        print(f"❌ 同步前處理模組失敗: {e}")

# 2. 🌟【核心關鍵：蓋隔離房間】自動建立前處理專用的獨立虛擬環境
if not os.path.exists(VENV_DIR):
    print(f"📦 正在為組員程式碼建立專屬隔離環境 (venv) 於: {VENV_DIR}...")
    try:
        # 使用系統 Python 創建一個完全乾淨隔離的子資料夾房間
        subprocess.run([sys.executable, "-m", "venv", VENV_DIR], check=True)
        
        print("⏳ 正在隔離環境內安裝舊版 NumPy 與 4D-Humans 連鎖依賴套件...")
        # 🎯 在隔離房間（preproc_env）內安裝前處理需要的所有套件
        # 這裡精確對齊了你提供的前處理 requirements.txt 清單，徹底解決 cv2 等缺失問題
        preprocess_deps = [
            "pip", "install", "--upgrade", "pip", "setuptools", "wheel",
            "numpy==1.23.5",               # 強制鎖定 numpy 1.x 版，修復 numpy._core 錯誤
            "opencv-python-headless",      # 解決 cv2 找不到的問題
            "torch==2.2.2", "torchvision==0.17.2", # 補上探測器（Detector）與核心張量庫
            "pytorch-lightning==2.1.0", "gdown", "webdataset", 
            "pandas", "scikit-image", "trimesh", "smplx==0.1.28", "chumpy",
            "timm", "einops", "yacs", "onnxruntime-gpu", "pillow", "tqdm", "pyrender"
        ]
        subprocess.run([VENV_PYTHON, "-m"] + preprocess_deps, check=True)
        print("✅ 前處理專屬隔離區套件配置大成功！")
    except Exception as e:
        print(f"❌ 隔離環境建立或套件安裝失敗: {e}")

# 3. 對齊 Blender 軟連結 (對應 Dockerfile)
blender_target_path = os.path.join(CHAMP_MASTER_DIR, "blender")
if os.path.exists("blender") and not os.path.exists(blender_target_path):
    try:
        os.symlink(os.path.abspath("blender"), blender_target_path)
        print("✅ Blender 軟連結架設成功，完美對齊隔離前處理路徑！")
    except Exception as e:
        print(f"⚠️ Blender 軟連結對齊失敗: {e}")

# 4. 自動下載 Champ 推論核心權重 (外面環境維持原本的最新狀態)
PRETRAINED_MODELS_DIR = "pretrained_models"
if not os.path.exists(PRETRAINED_MODELS_DIR) or len(os.listdir(PRETRAINED_MODELS_DIR)) < 4:
    print("⏳ 正在從 Hugging Face 官方庫下載 Champ 推論核心權重...")
    try:
        snapshot_download(
            repo_id="fudan-generative-ai/champ",
            local_dir=PRETRAINED_MODELS_DIR,
            repo_type="model",
            ignore_patterns=[".git*", "README.md"]
        )
        print("✅ Champ 核心推論權重下載成功。")
    except Exception as e:
        print(f"❌ 核心權重下載失敗: {e}")

print("🚀 [初始化完畢] 隔離區與外圍系統全線解耦串聯就緒，Gradio 即將點亮！\n")


# =========================================================================
# 🔒 自毀防禦計時器
# =========================================================================
def delayed_pause_space(repo_id, delay_seconds=180):
    print(f"🤖 [背景防禦線] 自毀計畫已啟動，倒數 {delay_seconds} 秒後將自動關機...")
    time.sleep(delay_seconds)
    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        try:
            api = HfApi(token=hf_token)
            api.pause_space(repo_id=repo_id)
            print("💤 [背景防禦線] Space 已進入安全休眠狀態。")
        except Exception: pass


# =========================================================================
# 5. 全新絕對沙盒防禦對齊版 Gradio Trigger Function (環境路徑劫持版)
# =========================================================================
def gradio_predict(input_image, input_video, start_frame, end_frame):
    import shutil # 引入實體檔案搬運工
    
    if not input_image or not input_video or not isinstance(input_video, str):
        return "❌ 錯誤：請確保正確上傳了參考人像與動作影片！", None, None
    
    # -------------------------------------------------------------------------
    # 🛡️ 實體防線：既然組員腳本硬要去讀 /app/champ-master/images，我們就直接蓋在它要的地方！
    # -------------------------------------------------------------------------
    custom_ref_root = os.path.abspath(CHAMP_MASTER_DIR)
    custom_ref_images = os.path.join(custom_ref_root, "images")
    
    if os.path.exists(custom_ref_images):
        shutil.rmtree(custom_ref_images)
    os.makedirs(custom_ref_images, exist_ok=True)
    
    target_image_path = os.path.join(custom_ref_images, "0001.png")
    try:
        shutil.copy(input_image, target_image_path)
        print(f"🎯 [物理降維打擊成功] 參考圖已直接入駐組員大本營: {target_image_path}")
    except Exception as e:
        return f"❌ 參考圖絕對沙盒配置失敗: {e}", None, None
    # -------------------------------------------------------------------------

    print("\n" + "="*50)
    print("🎬 [階段 1] 啟動隔離房間內的前處理執行組...")
    print("="*50)
    try:
        script_name = "process_single_video.py"
        
        # 🌟 核心破關線：複製當前的環境變數，並進行「路徑劫持」
        env = os.environ.copy()
        
        # 1. 把隔離房間的 bin 夾（裡面裝有我們剛裝好的 python）強行插到 Linux PATH 的最前面！
        # 這樣組員腳本在裡面不論呼叫幾百次 python，全部都會被強行導流進隔離房間！
        venv_bin_dir = os.path.join(VENV_DIR, "bin")
        env["PATH"] = f"{venv_bin_dir}:{env.get('PATH', '')}"
        
        # 2. 同步強注原始碼白名單
        master_path = os.path.abspath(CHAMP_MASTER_DIR)
        fourd_path = os.path.abspath(os.path.join(CHAMP_MASTER_DIR, "4D-Humans"))
        env["PYTHONPATH"] = f"{master_path}:{fourd_path}:{env.get('PYTHONPATH', '')}"
        
        # 3. 呼叫隔離房間的 Python 啟動前處理主程式
        preprocess_cmd = [
            VENV_PYTHON, script_name,
            "--video", os.path.abspath(input_video),
            "--ref", "./" 
        ]
        
        subprocess.run(preprocess_cmd, check=True, cwd=CHAMP_MASTER_DIR, env=env)
        print("✅ 隔離區前處理連鎖通車成功！")
    except Exception as e:
        return f"❌ 前處理隔離階段發生錯誤: {e}", None, None

    # 動態計算前處理生成的資料夾路徑
    video_filename = os.path.basename(input_video)
    video_name, _ = os.path.splitext(video_filename)
    computed_guidance_folder = os.path.join(CHAMP_MASTER_DIR, "transferd_result", video_name)
    
    if not os.path.exists(computed_guidance_folder):
        backup_folder = os.path.join(CHAMP_MASTER_DIR, "transferd_result")
        if os.path.exists(backup_folder) and len(os.listdir(backup_folder)) > 0:
            computed_guidance_folder = backup_folder

    print(f"🔍 最終定位條件影像資料夾: {computed_guidance_folder}")
    if not os.path.exists(computed_guidance_folder):
        return f"❌ 錯誤：找不到前處理生成目錄 {computed_guidance_folder}", None, None

    # 外圍主要大腦（Champ 推理層）路徑對齊
    current_dir = os.path.dirname(os.path.abspath(__file__))
    if current_dir not in sys.path: sys.path.insert(0, current_dir)
    if master_abs not in sys.path: sys.path.insert(0, master_abs)

    try:
        from inference import main as run_champ_inference
    except Exception as e:
        return f"❌ 載入 inference.py 失敗: {e}", None, None
    
    config = OmegaConf.load(os.path.join("configs", "inference", "inference.yaml"))
    config.data.ref_image_path = os.path.abspath(target_image_path)        
    config.data.guidance_data_folder = os.path.abspath(computed_guidance_folder)  
    config.data.frame_range = [int(start_frame), int(end_frame)]
    
    try:
        print("🚀 [階段 2] 外圍核心啟動擴散模型推理...")
        run_champ_inference(config)
        
        video_animation, video_grid, video_grid_wguidance = None, None, None
        results_dir = "results"
        if os.path.exists(results_dir):
            subdirs = [os.path.join(results_dir, d) for d in os.listdir(results_dir) if os.path.isdir(os.path.join(results_dir, d))]
            if subdirs:
                latest_dir = max(subdirs, key=os.path.getmtime)
                if os.path.exists(os.path.join(latest_dir, "animation.mp4")): video_animation = os.path.join(latest_dir, "animation.mp4")
                if os.path.exists(os.path.join(latest_dir, "grid.mp4")): video_grid = os.path.join(latest_dir, "grid.mp4")
                if os.path.exists(os.path.join(latest_dir, "grid_wguidance.mp4")): video_grid_wguidance = os.path.join(latest_dir, "grid_wguidance.mp4")

        threading.Thread(target=delayed_pause_space, args=("yoyozs11/champ_demo", 180), daemon=True).start()
        return video_animation, video_grid, video_grid_wguidance
        
    except Exception as e:
        return f"❌ 推論階段發生崩潰，錯誤訊息: {e}", None, None


# ==========================================
# 6. Gradio Web UI 佈局
# ==========================================
with gr.Blocks() as demo:
    gr.Markdown("# Champ 3D 動作生成系統 (高階獨立虛擬環境隔離版)")
    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(type="filepath", label="1. 上傳人像圖片")
            input_vid = gr.Video(label="2. 上傳自訂動作影片 (MP4)", format="mp4")
            with gr.Row():
                start_f = gr.Number(value=0, label="動作起始幀", precision=0)
                end_f = gr.Number(value=20, label="動作結束幀", precision=0)
            btn = gr.Button("開始前處理與影片生成", variant="primary")
        with gr.Column(scale=1):
            output_video_anim = gr.Video(label="3-1. 純動作生成結果")
            output_video_grid = gr.Video(label="3-2. 人像與結果對照組")
            output_video_grid_wguid = gr.Video(label="3-3. 完整骨架引導對照組")
            
    btn.click(
        fn=gradio_predict, 
        inputs=[input_img, input_vid, start_f, end_f], 
        outputs=[output_video_anim, output_video_grid, output_video_grid_wguid]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Default())