import os
import gradio as gr
import torch
import trimesh
from PIL import Image
from rembg import remove
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground

# Check for hardware acceleration on the server host
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading Generative 3D Transformer Model on: {device.upper()}...")
# Initialize the model weights globally once when the web server starts up
model = TSR.from_pretrained(
    "stabilityai/TripoSR",
    config_name="config.yaml",
    weight_name="model.ckpt"
)
model.to(device)
model.eval()

def process_image_to_stl(input_image):
    if input_image is None:
        return None, "Error: No image uploaded."
        
    try:
        # Step 1: Strip image backgrounds natively on the server
        print("Executing background removal layers...")
        no_bg_image = remove(input_image)
        
        # Step 2: Clear artifacts and scale to the neural network's bounding box
        processed_img = remove_background(no_bg_image, "white")
        processed_img = resize_foreground(processed_img, 0.85)

        # Step 3: Run the Large Reconstruction Model to infer 3D spatial values
        print("Processing 3D tensor field reconstruction...")
        with torch.no_grad():
            scene_codes = model([processed_img], device=device)
            # Use Marching Cubes algorithm at a standard 256^3 resolution grid
            meshes = model.extract_mesh(scene_codes, resolution=256)
            ai_mesh = meshes

        # Step 4: Extract the vertex mathematical arrays
        vertices = ai_mesh.vertices.cpu().numpy()
        faces = ai_mesh.faces.cpu().numpy()
        
        # Step 5: Convert vertex points from local coordinates to a true 3D printable bed layout
        vertices[:, [1, 2]] = vertices[:, [2, 1]]  # Swap Y and Z axes
        vertices[:, 1] *= -1                      # Correct face-up inversion

        # Create a solid geometry object
        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        mesh.process(validate=True)               # Remove overlapping nodes

        # Snap the absolute bottom boundary of the 3D mesh flat to Z=0 coordinate
        z_min = mesh.bounds[0][2]
        mesh.apply_translation([0, 0, -z_min])

        # Step 6: Write out a local binary file on the server partition
        output_filename = "generated_model.stl"
        mesh.export(output_filename, file_type='stl')
        
        status_msg = f"Success! Polygon Count: {len(mesh.faces)} | Solid Manifold: {mesh.is_watertight}"
        return output_filename, status_msg

    except Exception as e:
        return None, f"An algorithmic pipeline error occurred: {str(e)}"

# Define the HTML/CSS user portal via Gradio framework
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Local-Engine AI Image-to-STL Converter")
    gr.Markdown("Upload any image (objects, shapes, drawings) to synthesize a watertight, 3D-printable solid model without external API subscriptions.")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_img_slot = gr.Image(type="pil", label="Step 1: Upload Source Image")
            submit_btn = gr.Button("Generate 3D STL Mesh", variant="primary")
            
        with gr.Column(scale=1):
            output_file_slot = gr.File(label="Step 2: Download Ready-to-Print STL File")
            execution_log = gr.Textbox(label="System Pipeline Output Logs", interactive=False)

    # Bind elements to backend trigger functions
    submit_btn.click(
        fn=process_image_to_stl,
        inputs=[input_img_slot],
        outputs=[output_file_slot, execution_log]
    )

# Fire up the local webserver link on port 7860
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)