import gradio as gr
from pygltflib import GLTF2, Texture, Image
import base64
import io

import tempfile

def modify_texture():
    # Load the GLB file
    glb = GLTF2().load("train.glb")

    # Load the default texture image
    with open("defect.jpg", "rb") as f:
        default_texture = Image(uri=f"data:image/jpeg;base64,{base64.b64encode(f.read()).decode()}")

    # Assuming there's at least one image in the original GLB
    if glb.images:
        # Replace the first image with the default texture
        glb.images[0] = default_texture
    else:
        # Add default image if none exists
        glb.images.append(default_texture)
        # Update texture to point to the default image
        default_texture = Texture(source=len(glb.images)-1)
        glb.textures.append(default_texture)

    # Save the modified GLB to a temporary file
    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmp_file:
        tmp_filename = tmp_file.name
        glb.save(tmp_filename)

    # Read the temporary file back into memory as BytesIO
    with open(tmp_filename, "rb") as tmp_file:
        output_file = io.BytesIO(tmp_file.read())

    return output_file


gr.Interface(
    fn=modify_texture,
    inputs=[],
    outputs=gr.Model3D(label="Modified 3D Model"),
    title="GLB Texture Modifier"
).launch()