import gradio as gr
import torch
import os
import sys
import tempfile
import shutil
import subprocess
import spaces
from typing import Any, Dict, Union

# --- Configuration ---
# Path to the cloned UniRig repository directory within the Space
UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")

# Path to the Blender Python site-packages
BLENDER_PYTHON_PATH = "/opt/blender-4.2.0-linux-x64/4.2/python/lib/python3.11/site-packages"

BLENDER_PYTHON_EXEC = "/opt/blender-4.2.0-linux-x64/4.2/python/bin/python3.11"

# Path to the setup script
SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")

# Check if Blender is installed
if not os.path.exists("/usr/local/bin/blender"):
    print("Blender not found. Installing...")
    subprocess.run(["bash", SETUP_SCRIPT], check=True)
else:
    print("Blender is already installed.")

# Verify Blender Python path
if os.path.exists(BLENDER_PYTHON_PATH):
    print(f"Blender Python modules found at {BLENDER_PYTHON_PATH}")
    if os.path.exists(os.path.join(BLENDER_PYTHON_PATH, "bpy.so")):
        print("bpy.so found - Blender Python integration should work")
    else:
        print("bpy.so not found - Check Blender installation or path")
else:
    print(f"Blender Python modules not found at {BLENDER_PYTHON_PATH} - Update BLENDER_PYTHON_PATH")

if not os.path.isdir(UNIRIG_REPO_DIR):
    print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
    # Consider raising an error or displaying it in the UI if UniRig is critical for startup

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")
if DEVICE.type == 'cuda':
    print(f"CUDA Device Name: {torch.cuda.get_device_name(0)}")
    print(f"CUDA Version: {torch.version.cuda}")
else:
    print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")

def patch_asset_py():
    """Temporary patch to fix type hinting error in asset.py"""
    asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
    try:
        with open(asset_py_path, "r") as f:
            content = f.read()
        # Check if patch is already applied
        if "meta: Union[Dict[str, Any], None]" in content:
            print("Patch already applied to asset.py")
            return
        # Replace the problematic line
        content = content.replace("meta: Union[Dict[str, ...], None]=None", "meta: Union[Dict[str, Any], None]=None")
        # Ensure 'Any' is imported
        if "from typing import Any" not in content:
            content = "from typing import Any\n" + content
        with open(asset_py_path, "w") as f:
            f.write(content)
        print("Successfully patched asset.py")
    except Exception as e:
        print(f"Failed to patch asset.py: {e}")
        raise

@spaces.GPU # Decorator for ZeroGPU
def run_unirig_command(command_list, step_name):
    """Run UniRig commands using Blender's Python interpreter."""
    cmd = [BLENDER_PYTHON_EXEC] + command_list[1:]  # Use Blender's Python instead of bash!!
    
    print(f"Running {step_name}: {' '.join(cmd)}")
    
    process_env = os.environ.copy()
    
    # Determine the path to the 'src' directory within UniRig, where the 'unirig' package resides.
    unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")

    # Explicitly add BLENDER_PYTHON_PATH, UNIRIG_REPO_DIR/src, and UNIRIG_REPO_DIR to PYTHONPATH
    new_pythonpath_parts = [BLENDER_PYTHON_PATH, unirig_src_dir, UNIRIG_REPO_DIR]
    #existing_pythonpath = process_env.get('PYTHONPATH', '')
    #if existing_pythonpath:
    #    new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts))
    print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")

    try:
        result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
        print(f"{step_name} STDOUT:\n{result.stdout}")
        if result.stderr:
            print(f"{step_name} STDERR (non-fatal or warnings):\n{result.stderr}")
    except subprocess.CalledProcessError as e:
        print(f"ERROR during {step_name}:")
        print(f"Command: {' '.join(e.cmd)}")
        print(f"Return code: {e.returncode}")
        print(f"Stdout: {e.stdout}")
        print(f"Stderr: {e.stderr}")
        error_summary = e.stderr.splitlines()[-5:]
        raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
    except FileNotFoundError:
        print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}")
        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs and script paths.")
    except Exception as e_general:
        print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
        raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}")

@spaces.GPU
def rig_glb_mesh_multistep(input_glb_file_obj):
    """Rig a GLB mesh using UniRig's multi-step process."""
    patch_asset_py()

    if not os.path.isdir(UNIRIG_REPO_DIR):
        raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}.")

    if input_glb_file_obj is None:
        raise gr.Error("No input file provided. Please upload a .glb mesh.")

    input_glb_path = input_glb_file_obj
    print(f"Input GLB path received: {input_glb_path}")

    processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
    print(f"Using temporary processing directory: {processing_temp_dir}")

    try:
        base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
        abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
        abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
        abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")

        # Step 1: Skeleton Prediction
        print("Step 1: Predicting Skeleton...")
        skeleton_cmd = ["python", "launch/inference/generate_skeleton.py", "--input", input_glb_path, "--output", abs_skeleton_output_path]
        run_unirig_command(skeleton_cmd, "Skeleton Prediction")
        if not os.path.exists(abs_skeleton_output_path):
            raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")

        # Step 2: Skinning Weight Prediction
        print("Step 2: Predicting Skinning Weights...")
        skin_cmd = ["python", "launch/inference/generate_skin.py", "--input", abs_skeleton_output_path, "--source", input_glb_path, "--output", abs_skin_output_path]
        run_unirig_command(skin_cmd, "Skinning Prediction")
        if not os.path.exists(abs_skin_output_path):
            raise gr.Error("Skinning prediction failed to produce an output file.")

        # Step 3: Merge Skeleton/Skin with Original Mesh
        print("Step 3: Merging Results...")
        merge_cmd = ["python", "launch/inference/merge.py", "--source", abs_skin_output_path, "--target", input_glb_path, "--output", abs_final_rigged_glb_path]
        run_unirig_command(merge_cmd, "Merging")
        if not os.path.exists(abs_final_rigged_glb_path):
            raise gr.Error("Merging process failed to produce the final rigged GLB file.")

        return abs_final_rigged_glb_path

    except gr.Error:
        if os.path.exists(processing_temp_dir):
            shutil.rmtree(processing_temp_dir)
            print(f"Cleaned up temporary directory: {processing_temp_dir}")
        raise
    except Exception as e:
        print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
        if os.path.exists(processing_temp_dir):
            shutil.rmtree(processing_temp_dir)
            print(f"Cleaned up temporary directory: {processing_temp_dir}")
        raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")

# --- Gradio Interface ---
theme = gr.themes.Soft(
    primary_hue=gr.themes.colors.sky,
    secondary_hue=gr.themes.colors.blue,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__":
    print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")

iface = gr.Interface(
    fn=rig_glb_mesh_multistep,
    inputs=gr.File(label="Upload .glb Mesh File", type="filepath"),
    outputs=gr.Model3D(label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0]),
    title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
    description=(
        "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh.\n"
        "Running on: {DEVICE.upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
        "UniRig Source: https://github.com/VAST-AI-Research/UniRig"
    ),
    cache_examples=False,
    theme=theme
)

if __name__ == "__main__":
    if not os.path.isdir(UNIRIG_REPO_DIR):
        print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}.")
    iface.launch()