import sys
import spaces
sys.path.append("flash3d")  # Add the flash3d directory to the system path for importing local modules

from omegaconf import OmegaConf
import gradio as gr
import torch
import torchvision.transforms as TT
import torchvision.transforms.functional as TTF
from huggingface_hub import hf_hub_download
import numpy as np

from networks.gaussian_predictor import GaussianPredictor
from util.vis3d import save_ply

def main():
    print("[INFO] Starting main function...")
    # Determine if CUDA (GPU) is available and set the device accordingly
    if torch.cuda.is_available():
        device = "cuda:0"
        print("[INFO] CUDA is available. Using GPU device.")
    else:
        device = "cpu"
        print("[INFO] CUDA is not available. Using CPU device.")

    # Download model configuration and weights from Hugging Face Hub
    print("[INFO] Downloading model configuration...")
    model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                     filename="config_re10k_v1.yaml")
    print("[INFO] Downloading model weights...")
    model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                 filename="model_re10k_v1.pth")

    # Load model configuration using OmegaConf
    print("[INFO] Loading model configuration...")
    cfg = OmegaConf.load(model_cfg_path)
    
    # Initialize the GaussianPredictor model with the loaded configuration
    print("[INFO] Initializing GaussianPredictor model...")
    model = GaussianPredictor(cfg)
    device = torch.device(device)
    model.to(device)  # Move the model to the specified device (CPU or GPU)
    
    # Load the pre-trained model weights
    print("[INFO] Loading model weights...")
    model.load_model(model_path)

    # Define transformation functions for image preprocessing
    pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))  # Padding to augment the image borders
    to_tensor = TT.ToTensor()  # Convert image to tensor

    # Function to check if an image is uploaded by the user
    def check_input_image(input_image):
        print("[DEBUG] Checking input image...")
        if input_image is None:
            print("[ERROR] No image uploaded!")
            raise gr.Error("No image uploaded!")
        print("[INFO] Input image is valid.")

    # Function to preprocess the input image before passing it to the model
    def preprocess(image):
        print("[DEBUG] Preprocessing image...")
        # Resize the image to the desired height and width specified in the configuration
        image = TTF.resize(
            image, (cfg.dataset.height, cfg.dataset.width), 
            interpolation=TT.InterpolationMode.BICUBIC
        )
        # Apply padding to the image
        image = pad_border_fn(image)
        print("[INFO] Image preprocessing complete.")
        return image

    # Function to reconstruct the 3D model from the input image and export it as a PLY file
    @spaces.GPU(duration=120)  # Decorator to allocate a GPU for this function during execution
    def reconstruct_and_export(image):
        """
        Passes image through model, outputs reconstruction in form of a dict of tensors.
        """
        print("[DEBUG] Starting reconstruction and export...")
        # Convert the preprocessed image to a tensor and move it to the specified device
        image = to_tensor(image).to(device).unsqueeze(0)
        inputs = {
            ("color_aug", 0, 0): image,
        }

        # Pass the image through the model to get the output
        print("[INFO] Passing image through the model...")
        outputs = model(inputs)

        # Export the reconstruction to a PLY file
        print(f"[INFO] Saving output to {ply_out_path}...")
        save_ply(outputs, ply_out_path, num_gauss=2)
        print("[INFO] Reconstruction and export complete.")

        return ply_out_path
    
    # Path to save the output PLY file
    ply_out_path = f'./mesh.ply'

    # CSS styling for the Gradio interface
    css = """
        h1 {
            text-align: center;
            display:block;
        }
        """

    # Create the Gradio user interface
    with gr.Blocks(css=css) as demo:
        gr.Markdown(
            """
            # Flash3D
            """
        )
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                with gr.Row():
                    # Input image component for the user to upload an image
                    input_image = gr.Image(
                        label="Input Image",
                        image_mode="RGBA",
                        sources="upload",
                        type="pil",
                        elem_id="content_image",
                    )
                with gr.Row():
                    # Button to trigger the generation process
                    submit = gr.Button("Generate", elem_id="generate", variant="primary")

                with gr.Row(variant="panel"): 
                    # Examples panel to provide sample images for users
                    gr.Examples(
                        examples=[
                            './demo_examples/bedroom_01.png',
                            './demo_examples/kitti_02.png',
                            './demo_examples/kitti_03.png',
                            './demo_examples/re10k_04.jpg',
                            './demo_examples/re10k_05.jpg',
                            './demo_examples/re10k_06.jpg',
                        ],
                        inputs=[input_image],
                        cache_examples=False,
                        label="Examples",
                        examples_per_page=20,
                    )

                with gr.Row():
                    # Display the preprocessed image (after resizing and padding)
                    processed_image = gr.Image(label="Processed Image", interactive=False)

            with gr.Column(scale=2):
                with gr.Row():
                    with gr.Tab("Reconstruction"):
                        # 3D model viewer to display the reconstructed model
                        output_model = gr.Model3D(
                            height=512,
                            label="Output Model",
                            interactive=False
                        )

        # Define the workflow for the Generate button
        submit.click(fn=check_input_image, inputs=[input_image]).success(
            fn=preprocess,
            inputs=[input_image],
            outputs=[processed_image],
        ).success(
            fn=reconstruct_and_export,
            inputs=[processed_image],
            outputs=[output_model],
        )

    # Queue the requests to handle them sequentially (to avoid GPU resource conflicts)
    demo.queue(max_size=1)
    print("[INFO] Launching Gradio demo...")
    demo.launch(share=True)  # Launch the Gradio interface and allow public sharing

if __name__ == "__main__":
    print("[INFO] Running application...")
    main()