import shlex
import subprocess
import tempfile
import time

import numpy as np
import rembg
import spaces
import torch
from PIL import Image

subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))

from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation

from flask import Flask, flash, request
from flask_session import Session


app = Flask(__name__)
app.config["SESSION_PERMANENT"] = False
app.config["SESSION_TYPE"] = "filesystem"

Session(app)


if torch.cuda.is_available():
    device = "cuda:0"
else:
    device = "cpu"

model = TSR.from_pretrained(
    "stabilityai/TripoSR",
    config_name="config.yaml",
    weight_name="model.ckpt",
)
model.renderer.set_chunk_size(131072)
model.to(device)

rembg_session = rembg.new_session()


def check_input_image(input_image):
    if input_image is None:
        raise ValueError('Please provide an input image.')


def preprocess(input_image, do_remove_background, foreground_ratio):
    def fill_background(image):
        image = np.array(image).astype(np.float32) / 255.0
        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
        image = Image.fromarray((image * 255.0).astype(np.uint8))
        return image

    if do_remove_background:
        image = input_image.convert("RGB")
        image = remove_background(image, rembg_session)
        image = resize_foreground(image, foreground_ratio)
        image = fill_background(image)
    else:
        image = input_image
        if image.mode == "RGBA":
            image = fill_background(image)
    return image


@spaces.GPU
def generate(image, mc_resolution, formats=["obj", "glb"]):
    scene_codes = model(image, device=device)
    mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
    mesh = to_gradio_3d_orientation(mesh)

    mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
    mesh.export(mesh_path_glb.name)

    mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
    mesh.apply_scale([-1, 1, 1])  # Otherwise the visualized .obj will be flipped
    mesh.export(mesh_path_obj.name)
    
    return mesh_path_obj.name, mesh_path_glb.name

def run_example(image_pil):
    preprocessed = preprocess(image_pil, False, 0.9)
    mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
    return preprocessed, mesh_name_obj, mesh_name_glb


@app.route("/", methods=['GET', 'POST'])
def hello():
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('No file part')
            return {"status": "Failed", "message": "Please Provide file name(file)."}
        file = request.files['file']
        image = Image.open(file)
        preprocess_image = run_example(image)
        print(preprocess_image)
        return {"status": "Success", "message": "You can download the 3D model.", "data": preprocess_image}
        
    else: 
        return {
            "status": "Success",
            "message":"You can upload an image file to get the 3D model."
        }



if __name__ == "__main__":
  app.run()

# with gr.Blocks() as demo:
#     gr.Markdown(HEADER)
#     with gr.Row(variant="panel"):
#         with gr.Column():
#             with gr.Row():
#                 input_image = gr.Image(
#                     label="Input Image",
#                     image_mode="RGBA",
#                     sources="upload",
#                     type="pil",
#                     elem_id="content_image",
#                 )
#                 processed_image = gr.Image(label="Processed Image", interactive=False)
#             with gr.Row():
#                 with gr.Group():
#                     do_remove_background = gr.Checkbox(
#                         label="Remove Background", value=True
#                     )
#                     foreground_ratio = gr.Slider(
#                         label="Foreground Ratio",
#                         minimum=0.5,
#                         maximum=1.0,
#                         value=0.85,
#                         step=0.05,
#                     )
#                     mc_resolution = gr.Slider(
#                         label="Marching Cubes Resolution",
#                         minimum=32,
#                         maximum=320,
#                         value=256,
#                         step=32
#                      )
#             with gr.Row():
#                 submit = gr.Button("Generate", elem_id="generate", variant="primary")
#         with gr.Column():
#             with gr.Tab("OBJ"):
#                 output_model_obj = gr.Model3D(
#                     label="Output Model (OBJ Format)",
#                     interactive=False,
#                 )
#                 gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
#             with gr.Tab("GLB"):
#                 output_model_glb = gr.Model3D(
#                     label="Output Model (GLB Format)",
#                     interactive=False,
#                 )
#                 gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
#     with gr.Row(variant="panel"):
#         gr.Examples(
#             examples=[
#                 os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
#             ],
#             inputs=[input_image],
#             outputs=[processed_image, output_model_obj, output_model_glb],
#             cache_examples=True,
#             fn=partial(run_example),
#             label="Examples",
#             examples_per_page=20
#         )
#     submit.click(fn=check_input_image, inputs=[input_image]).success(
#         fn=preprocess,
#         inputs=[input_image, do_remove_background, foreground_ratio],
#         outputs=[processed_image],
#     ).success(
#         fn=generate,
#         inputs=[processed_image, mc_resolution],
#         outputs=[output_model_obj, output_model_glb],
#     )

# demo.queue(max_size=10)
# demo.launch()