import gradio as gr
import torch
import os
import sys
import tempfile
import shutil
import subprocess
import spaces

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

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

@spaces.GPU # Decorator for ZeroGPU
def run_unirig_command(command_list, step_name):
    """
    Helper function to run UniRig commands (now expecting bash scripts) using subprocess.
    command_list: The full command and its arguments, e.g., ["bash", "script.sh", "--arg", "value"]
    """
    # The command_list is now expected to be the full command, e.g., starting with "bash"
    cmd = command_list
    
    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 UNIRIG_REPO_DIR/src to PYTHONPATH for the subprocess.
    # The bash scripts will internally call Python, which needs to find the 'unirig' package.
    # Also, keep UNIRIG_REPO_DIR itself in case some scripts or modules there are run directly
    # or expect the project root to be in PYTHONPATH.
    existing_pythonpath = process_env.get('PYTHONPATH', '')
    new_pythonpath_parts = [unirig_src_dir, UNIRIG_REPO_DIR] # UniRig/src first, then UniRig/
    if existing_pythonpath:
        # Prepend our paths to existing PYTHONPATH
        new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
    
    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts)) # filter(None,...) handles empty existing_pythonpath
    print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")

    try:
        # Execute the command from the UniRig directory (UNIRIG_REPO_DIR)
        # This is crucial for the bash scripts to find their relative paths (e.g., to Python scripts)
        # and for any underlying Python/Hydra calls to find configurations (e.g., in UniRig/configs/)
        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}")
        # Provide a more user-friendly error, potentially masking long tracebacks
        error_summary = e.stderr.splitlines()[-5:] # Last 5 lines of stderr
        raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
    except FileNotFoundError:
        # This error means the executable (e.g., "bash" or the script itself) was not found.
        print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}. Is UniRig cloned correctly and 'bash' available?")
        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs, UniRig directory structure, 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 # Decorator for ZeroGPU
def rig_glb_mesh_multistep(input_glb_file_obj):
    """
    Takes an input GLB file object (from gr.File with type="filepath"),
    rigs it using the new UniRig multi-step process by calling its bash scripts,
    and returns the path to the final rigged GLB file.
    """
    if not os.path.isdir(UNIRIG_REPO_DIR):
         raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Cannot proceed. Please check Space setup.")

    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 # This is the absolute path from gr.File(type="filepath")
    print(f"Input GLB path received: {input_glb_path}")

    # Create a dedicated temporary directory for all intermediate and final files
    # The output paths for UniRig scripts will point into this directory.
    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]
        
        # Define absolute paths for intermediate files within the processing_temp_dir
        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 using generate_skeleton.sh
        print("Step 1: Predicting Skeleton...")
        skeleton_cmd = [
            "bash", "launch/inference/generate_skeleton.sh",
            "--input", input_glb_path, # Input is the original GLB
            "--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 using generate_skin.sh
        print("Step 2: Predicting Skinning Weights...")
        # generate_skin.sh requires the skeleton from step 1 as --input,
        # and the original mesh as --source.
        skin_cmd = [
            "bash", "launch/inference/generate_skin.sh",
            "--input", abs_skeleton_output_path,    # Input is the skeleton FBX from previous step
            "--source", input_glb_path,             # Source is the original GLB mesh
            "--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. Check logs for UniRig errors.")

        # Step 3: Merge Skeleton/Skin with Original Mesh using merge.sh
        print("Step 3: Merging Results...")
        # merge.sh requires the skinned FBX as --source (which contains skeleton and weights)
        # and the original GLB as --target.
        merge_cmd = [
            "bash", "launch/inference/merge.sh",
            "--source", abs_skin_output_path,      # Source is the skinned FBX from previous step
            "--target", input_glb_path,            # Target is the original GLB mesh
            "--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. Check logs for UniRig errors.")

        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}. The application will not work.")

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 by calling its provided bash scripts.\n"
        "The process involves: 1. Skeleton Prediction, 2. Skinning Weight Prediction, 3. Merging.\n"
        "This may take several minutes. Ensure your GLB has clean geometry.\n"
        f"Running on: {str(DEVICE).upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
        f"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}. Ensure it's cloned in the Space's root.")
    
    iface.launch()