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

# --- Configuration ---
# Path to the cloned UniRig repository directory within the Space
# Ensure this path is correct relative to app.py
UNIRIG_REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "UniRig")) # Get absolute path

# Absolute path to the Blender installation provided in the Space environment
BLENDER_INSTALL_DIR = "/opt/blender-4.2.0-linux-x64"
BLENDER_PYTHON_VERSION_DIR = "4.2" # Corresponds to Blender 4.2.x
BLENDER_PYTHON_VERSION = "python3.11" # Blender 4.2 uses Python 3.11

# Construct paths dynamically based on the above
BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
BLENDER_PYTHON_BIN_DIR = os.path.join(BLENDER_PYTHON_DIR, "bin") # Directory containing python executable
BLENDER_PYTHON_EXEC = os.path.join(BLENDER_PYTHON_BIN_DIR, BLENDER_PYTHON_VERSION) # Full path to executable
BLENDER_PYTHON_LIB_PATH = os.path.join(BLENDER_PYTHON_DIR, "lib", BLENDER_PYTHON_VERSION)
BLENDER_PYTHON_SITE_PACKAGES = os.path.join(BLENDER_PYTHON_LIB_PATH, "site-packages")

# Path to the setup script (executed only if Blender isn't found initially)
SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")

# --- Initial Checks ---
print("--- Environment Checks ---")

# Check if Blender Python executable exists (needed for environment setup)
# Run setup script if Blender isn't found (assuming setup script handles installation)
if not os.path.exists(BLENDER_PYTHON_EXEC):
    print(f"Blender Python executable not found at {BLENDER_PYTHON_EXEC}. Running setup script...")
    if os.path.exists(SETUP_SCRIPT):
        try:
            # Run the setup script using bash
            setup_result = subprocess.run(["bash", SETUP_SCRIPT], check=True, capture_output=True, text=True)
            print("Setup script executed successfully.")
            print(f"Setup STDOUT:\n{setup_result.stdout}")
            if setup_result.stderr:
                 print(f"Setup STDERR:\n{setup_result.stderr}")
            # Re-check if the executable exists after running the script
            if not os.path.exists(BLENDER_PYTHON_EXEC):
                 raise RuntimeError(f"Setup script ran but Blender Python still not found at {BLENDER_PYTHON_EXEC}.")
        except subprocess.CalledProcessError as e:
             print(f"ERROR running setup script: {SETUP_SCRIPT}")
             print(f"Return Code: {e.returncode}")
             print(f"Stdout: {e.stdout}")
             print(f"Stderr: {e.stderr}")
             # Raise a more informative error for Gradio if setup fails
             raise gr.Error(f"Failed to execute setup script. Check logs. Stderr: {e.stderr[-500:]}")
        except Exception as e:
             raise gr.Error(f"Unexpected error running setup script '{SETUP_SCRIPT}': {e}")
    else:
        # If setup script is missing, raise a Gradio error
        raise gr.Error(f"Blender Python not found and setup script missing: {SETUP_SCRIPT}")
else:
    print(f"Blender Python executable found: {BLENDER_PYTHON_EXEC}")

# Verify Blender Python site-packages path and bpy module presence
bpy_found = False
if os.path.exists(BLENDER_PYTHON_SITE_PACKAGES):
    print(f"Blender Python site-packages found at: {BLENDER_PYTHON_SITE_PACKAGES}")
    # Check 1: 'bpy' directory with __init__.py
    bpy_module_dir = os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy")
    if os.path.isdir(bpy_module_dir) and os.path.exists(os.path.join(bpy_module_dir, "__init__.py")):
        print("Found 'bpy' module directory in site-packages.")
        bpy_found = True
    # Check 2: 'bpy.so' file (less common structure)
    elif os.path.exists(os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy.so")):
         print("Found 'bpy.so' in site-packages.")
         bpy_found = True

    if not bpy_found:
        print("WARNING: Blender Python 'bpy' module indicator not found in site-packages. Imports might fail.")
else:
    print(f"WARNING: Blender Python site-packages directory not found at {BLENDER_PYTHON_SITE_PACKAGES}. Check paths.")


# Check for UniRig repository
if not os.path.isdir(UNIRIG_REPO_DIR):
    print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
    # Raise Gradio error if critical component is missing
    raise gr.Error(f"UniRig repository missing at: {UNIRIG_REPO_DIR}. Ensure it's cloned correctly.")
else:
    print(f"UniRig repository found at: {UNIRIG_REPO_DIR}")

# Check PyTorch and CUDA
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Gradio environment using device: {DEVICE}")
if DEVICE.type == 'cuda':
    try:
        print(f"Gradio CUDA Device Name: {torch.cuda.get_device_name(0)}")
        print(f"Gradio PyTorch CUDA Version: {torch.version.cuda}")
    except Exception as e:
        print(f"Could not get CUDA device details: {e}")
else:
    print("Warning: Gradio environment CUDA not available or not detected by PyTorch.")
    print("UniRig subprocess will attempt to use GPU via Blender's Python environment (invoked by .sh scripts).")

print("--- End Environment Checks ---")

# --- Helper Functions ---

def patch_asset_py():
    """Temporary patch to fix type hinting error in UniRig's asset.py"""
    # This patch might still be needed if the .sh scripts call Python code that uses asset.py
    asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
    try:
        # Check if file exists before trying to open
        if not os.path.exists(asset_py_path):
            print(f"Warning: asset.py not found at {asset_py_path}, skipping patch.")
            return # Don't raise error, maybe it's not needed

        with open(asset_py_path, "r") as f:
            content = f.read()

        problematic_line = "meta: Union[Dict[str, ...], None]=None"
        corrected_line = "meta: Union[Dict[str, Any], None]=None"
        typing_import = "from typing import Any"

        if corrected_line in content:
            print("Patch already applied to asset.py")
            return
        if problematic_line not in content:
             print("Problematic line not found in asset.py, patch might be unnecessary.")
             return

        print("Applying patch to asset.py...")
        content = content.replace(problematic_line, corrected_line)
        if typing_import not in content:
            if "from typing import" in content:
                 content = content.replace("from typing import", f"{typing_import}\nfrom typing import", 1)
            else:
                 content = f"{typing_import}\n{content}"

        with open(asset_py_path, "w") as f:
            f.write(content)
        print("Successfully patched asset.py")

    except Exception as e:
        # Log error but don't necessarily stop the app, maybe patch isn't critical
        print(f"ERROR: Failed to patch asset.py: {e}. Proceeding cautiously.")
        # raise gr.Error(f"Failed to apply necessary patch to UniRig code: {e}") # Optional: make it fatal

@spaces.GPU
def run_unirig_command(script_path: str, args: List[str], step_name: str):
    """
    Runs a specific UniRig SHELL script (.sh) using bash in a subprocess,
    ensuring the correct environment (PYTHONPATH, etc.) is set.

    Args:
        script_path: Absolute path to the .sh script to execute.
        args: A list of command-line arguments for the script.
        step_name: Name of the step for logging.
    """
    # Command is now bash + script path + arguments
    cmd = ["bash", script_path] + args

    print(f"\n--- Running UniRig Step: {step_name} ---")
    print(f"Command: {' '.join(cmd)}")

    # Prepare the environment for the subprocess (shell script)
    process_env = os.environ.copy()
    unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")

    # 1. Set PYTHONPATH: Blender's site-packages + UniRig source
    pythonpath_parts = [
        BLENDER_PYTHON_SITE_PACKAGES,
        unirig_src_dir,
        UNIRIG_REPO_DIR
    ]
    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
    print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")

    # 2. Set LD_LIBRARY_PATH: Include Blender's Python library directory
    blender_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
    # Prepend Blender lib path to existing LD_LIBRARY_PATH if it exists
    existing_ld_path = process_env.get('LD_LIBRARY_PATH', '')
    process_env["LD_LIBRARY_PATH"] = f"{blender_lib_path}{os.pathsep}{existing_ld_path}" if existing_ld_path else blender_lib_path
    print(f"Subprocess LD_LIBRARY_PATH: {process_env['LD_LIBRARY_PATH']}")

    # 3. Set PATH: *Prepend* Blender's Python bin directory to the system PATH.
    # This makes it the *first* place the shell looks for 'python' or 'python3.11'.
    existing_path = process_env.get('PATH', '')
    process_env["PATH"] = f"{BLENDER_PYTHON_BIN_DIR}{os.pathsep}{existing_path}"
    print(f"Subprocess PATH: {process_env['PATH']}")


    try:
        # Execute the shell script.
        # cwd=UNIRIG_REPO_DIR ensures the script runs from the repo's root,
        result = subprocess.run(
            cmd,
            cwd=UNIRIG_REPO_DIR,
            capture_output=True,
            text=True,
            check=True, # Raises CalledProcessError on non-zero exit codes
            env=process_env # Pass the modified environment
        )
        print(f"{step_name} STDOUT:\n{result.stdout}")
        if result.stderr:
            print(f"{step_name} STDERR (Warnings/Info):\n{result.stderr}")

    except subprocess.CalledProcessError as e:
        print(f"ERROR during {step_name}: Subprocess failed!")
        print(f"Command: {' '.join(e.cmd)}")
        print(f"Return code: {e.returncode}")
        print(f"--- {step_name} STDOUT ---:\n{e.stdout}")
        print(f"--- {step_name} STDERR ---:\n{e.stderr}")
        error_summary = e.stderr.strip().splitlines()
        last_lines = "\n".join(error_summary[-5:]) if error_summary else "No stderr output."
        # Check specifically for bpy/torch import errors within the subprocess stderr
        if "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to import Blender's 'bpy' module. Check environment setup.")
        elif "ImportError: Failed to load PyTorch C extensions" in e.stderr:
             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to load PyTorch extensions. Check environment and PyTorch installation within Blender's Python.")
        else:
             raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")


    except FileNotFoundError:
        # This error means 'bash' or the script_path wasn't found
        print(f"ERROR: Could not find executable 'bash' or script '{script_path}' for {step_name}.")
        print(f"Attempted command: {' '.join(cmd)}")
        raise gr.Error(f"Setup error for UniRig '{step_name}'. 'bash' or script '{os.path.basename(script_path)}' not found.")

    except Exception as e_general:
        print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")

    print(f"--- Finished UniRig Step: {step_name} ---")


@spaces.GPU
def rig_glb_mesh_multistep(input_glb_file_obj):
    """
    Main Gradio function to rig a GLB mesh using UniRig's multi-step process.
    Orchestrates calls to run_unirig_command for each step, executing .sh scripts.
    """
    try:
        patch_asset_py() # Attempt patch, might still be relevant
    except gr.Error as e:
         print(f"Ignoring patch error: {e}") # Or just log it
    except Exception as e:
         print(f"Ignoring unexpected patch error: {e}")

    # --- Input Validation ---
    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}")

    if not os.path.exists(input_glb_path):
         raise gr.Error(f"Input file path does not exist: {input_glb_path}")
    if not input_glb_path.lower().endswith(".glb"):
         raise gr.Error("Invalid file type. Please upload a .glb file.")

    # --- Setup Temporary Directory ---
    processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
    print(f"Using temporary processing directory: {processing_temp_dir}")

    try:
        # --- Define File Paths ---
        base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
        # Ensure paths passed to scripts are absolute
        abs_input_glb_path = os.path.abspath(input_glb_path)
        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")

        # --- Define Absolute Paths to UniRig SHELL Scripts ---
        skeleton_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/generate_skeleton.sh")
        skin_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/generate_skin.sh")
        merge_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/merge.sh")

        # --- Execute UniRig Steps ---

        # Step 1: Skeleton Prediction
        print("\nStarting Step 1: Predicting Skeleton...")
        skeleton_args = [
            "--input", abs_input_glb_path,
            "--output", abs_skeleton_output_path
        ]
        if not os.path.exists(skeleton_script_path):
            raise gr.Error(f"Skeleton script not found at: {skeleton_script_path}")
        run_unirig_command(skeleton_script_path, skeleton_args, "Skeleton Prediction")
        if not os.path.exists(abs_skeleton_output_path):
            # Check if the error wasn't already raised by run_unirig_command
            raise gr.Error("Skeleton prediction failed. Output file not created. Check logs.")
        print("Step 1: Skeleton Prediction completed.")

        # Step 2: Skinning Weight Prediction
        print("\nStarting Step 2: Predicting Skinning Weights...")
        skin_args = [
            "--input", abs_skeleton_output_path, # Input is the skeleton from step 1
            "--source", abs_input_glb_path,    # Source mesh
            "--output", abs_skin_output_path
        ]
        if not os.path.exists(skin_script_path):
            raise gr.Error(f"Skinning script not found at: {skin_script_path}")
        run_unirig_command(skin_script_path, skin_args, "Skinning Prediction")
        if not os.path.exists(abs_skin_output_path):
            raise gr.Error("Skinning prediction failed. Output file not created. Check logs.")
        print("Step 2: Skinning Prediction completed.")

        # Step 3: Merge Skeleton/Skin with Original Mesh
        print("\nStarting Step 3: Merging Results...")
        merge_args = [
            # Determine which source to use based on what exists or a user choice?
            # Assuming skin output is the desired source if it exists.
            "--source", abs_skin_output_path,
            "--target", abs_input_glb_path,
            "--output", abs_final_rigged_glb_path
        ]
        if not os.path.exists(merge_script_path):
            raise gr.Error(f"Merging script not found at: {merge_script_path}")
        run_unirig_command(merge_script_path, merge_args, "Merging")
        if not os.path.exists(abs_final_rigged_glb_path):
            raise gr.Error("Merging process failed. Final rigged GLB file not created. Check logs.")
        print("Step 3: Merging completed.")

        # --- Return Result ---
        print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
        return abs_final_rigged_glb_path

    except gr.Error as e:
        print(f"Gradio Error occurred: {e}")
        if os.path.exists(processing_temp_dir):
            shutil.rmtree(processing_temp_dir)
            print(f"Cleaned up temporary directory: {processing_temp_dir}")
        raise e

    except Exception as e:
        print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
        import traceback
        traceback.print_exc()
        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 Definition ---
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"],
)

# Check UniRig repo existence again before building the interface
if not os.path.isdir(UNIRIG_REPO_DIR):
     startup_error_message = (
         f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. "
         "The application cannot start. Please ensure the repository is cloned correctly."
     )
     print(startup_error_message)
     # Display an error message if Gradio tries to build the UI without UniRig
     with gr.Blocks(theme=theme) as iface:
          gr.Markdown(f"# Application Error\n\n{startup_error_message}")
else:
     # Build the normal interface if UniRig is found
     iface = gr.Interface(
         fn=rig_glb_mesh_multistep,
         inputs=gr.File(
             label="Upload .glb Mesh File",
             type="filepath", # Provides the path to the uploaded file
             file_types=[".glb"] # Restrict file types
         ),
         outputs=gr.Model3D(
             label="Rigged 3D Model (.glb)",
             clear_color=[0.8, 0.8, 0.8, 1.0] # Background color for the viewer
         ),
         title=f"UniRig Auto-Rigger (Blender {BLENDER_PYTHON_VERSION_DIR} / Python {BLENDER_PYTHON_VERSION})",
         description=(
             "Upload a 3D mesh in `.glb` format. This application uses UniRig via Blender (invoked through shell scripts) to automatically generate a skeleton and skinning weights.\n"
             f"* Running main app on Python {sys.version.split()[0]}, UniRig steps use Blender's Python {BLENDER_PYTHON_VERSION} via .sh scripts.\n"
             f"* Utilizing device: **{DEVICE.type.upper()}** (via ZeroGPU if available).\n"
             f"* UniRig Source: https://github.com/VAST-AI-Research/UniRig"
         ),
         cache_examples=False, # Disable caching if results depend heavily on GPU state or are large
         theme=theme,
         allow_flagging='never' # Disable flagging unless needed
     )

# --- Launch the Application ---
if __name__ == "__main__":
    # Ensure the interface object exists before launching
    if 'iface' in locals():
        print("Launching Gradio interface...")
        # Consider adding share=True for public link if needed, or server_name="0.0.0.0"
        iface.launch()
    else:
        # This case should only happen if the UniRig repo check failed above
        print("ERROR: Gradio interface not created due to startup errors.")