import gradio as gr
import os
import subprocess
import shutil
import torch # SHARP relies on PyTorch

# The SHARP repository is expected to be cloned into a folder named 'SHARP'
SHARP_REPO_PATH = "SHARP"
# The model checkpoint is expected to be in the root directory of the Space
MODEL_CHECKPOINT_PATH = "sharp_2572gikvuh.pt"

# Function to run the sharp predict command
def run_sharp_predict(input_files):
    if not input_files:
        return "Please upload at least one image file."

    # Create a temporary directory for input images
    input_dir = "temp_input_images"
    os.makedirs(input_dir, exist_ok=True)

    # Save uploaded images to the temporary input directory
    for i, img_file in enumerate(input_files):
        original_filename = getattr(img_file, 'orig_name', f"image_{i}.png")
        dest_path = os.path.join(input_dir, original_filename)
        shutil.copyfile(img_file.name, dest_path)

    # Create an output directory for gaussians
    output_dir = "temp_output_gaussians"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    sharp_script_path = os.path.join(SHARP_REPO_PATH, "sharp", "predict.py")

    if not os.path.exists(sharp_script_path):
        return f"Error: SHARP project not found at '{SHARP_REPO_PATH}' or 'predict.py' is missing. Please ensure the repository is cloned and set up correctly."

    command = [
        "python", sharp_script_path,
        "-i", input_dir,
        "-o", output_dir,
        "-c", MODEL_CHECKPOINT_PATH
    ]

    result_message = ""
    try:
        process = subprocess.run(command, capture_output=True, text=True, check=True)
        result_message += f"Prediction successful! Output in: {output_dir}\n\nSTDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}"

        generated_files = [f for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))]
        if generated_files:
            result_message += "\n\nGenerated files:\n" + "\n".join(generated_files)
        else:
            result_message += "\n\nNo output files found in the output directory."

    except subprocess.CalledProcessError as e:
        result_message = f"Prediction failed!\n\nSTDOUT:\n{e.stdout}\nSTDERR:\n{e.stderr}\nError: {e}"
    except FileNotFoundError:
        result_message = f"Error: 'python' or '{sharp_script_path}' not found. Ensure Python is in PATH and SHARP is correctly set up."
    except Exception as e:
        result_message = f"An unexpected error occurred: {e}"
    finally:
        if os.path.exists(input_dir):
            shutil.rmtree(input_dir)

    return result_message


# Create a Gradio interface for image upload and sharp predict
gradio_interface = gr.Interface(
    fn=run_sharp_predict,
    inputs=gr.File(file_count="multiple", type="filepath", label="Upload Input Images (e.g., from an MVS dataset)"),
    outputs=gr.Textbox(label="Prediction Output and Results"),
    title="SHARP: Image-to-3D Gaussian Splatting Demo",
    description="Upload a set of images to generate a 3D Gaussian Splatting representation using the apple/Sharp model. This will run the `sharp predict` command."
)

if __name__ == "__main__":
    gradio_interface.launch(share=True)