import gradio as gr
from image_dataset import ImageDataset
from image_wgan import ImageWgan
import os
from os.path import exists
from PIL import Image, ImageEnhance
import numpy as np
import random
from pygltflib import GLTF2
from pygltflib.utils import ImageFormat, Texture, Material, Image as GLTFImage
import time
import threading


clicked_image_index = 0
print(gr.__version__)
def init():
    generated_samples_folder = "."
    discriminator_saved_model = "discriminator64.model"
    generator_saved_model = "generator64.model"
    latent_space = 100
    image_wgan = ImageWgan(
        image_shape = (4,64,64),
        latent_space_dimension=latent_space,
        generator_saved_model=generator_saved_model if exists(generator_saved_model) else None,
        discriminator_saved_model=discriminator_saved_model if exists(discriminator_saved_model) else None
    )
    image_wgan.generate(
        sample_folder=generated_samples_folder
    )
    crop()


def crop():

    import generator
    res = 64
    if res != 0:
        results = "generated.png"
        img = Image.open(results)

        width,height = img.size


        top = 2
        bottom = 2
        for i in range(1):
            left = (res+2)*i +2
            right = width-(res+2)*i
            imgcrop = img.crop((left,top,left+res,res+2))


            imgcrop.save(str(i)+".png")    
init()



def gen(seed=None):
    if seed is None or seed == "Surprise Me" or seed == "":
        seed = random.randint(0, 1000000000)
    else:
        seed = int(seed)
    np.random.seed(seed)
    init()  # Initialize and generate samples
    crop()  # Crop the generated images
    imgArr = []
    for i in range(1):
        img = Image.open(f"{i}.png").convert('RGBA')
        img = img.resize((64, 64), Image.NEAREST)
        imgArr.append(img)
    return imgArr, seed  # Return both images and the seed used

def update_model(index):
    filename = "models/player_model.glb"
    gltf = GLTF2().load(filename)

    # Step 1: Find the index of the existing texture you want to replace
    # Let's assume the texture you want to replace is at index 1 (you need to replace 1 with the actual index)
    existing_texture_index = 0

    # Check if the existing_texture_index is valid
    if existing_texture_index < len(gltf.textures):
        # Step 2: Remove the old texture and its associated image from the GLB
        # Remove the texture
        gltf.textures.pop(existing_texture_index)

        # Remove the image associated with the texture
        existing_image_index = gltf.materials[0].pbrMetallicRoughness.baseColorTexture.index
        gltf.images.pop(existing_image_index)


        # Step 3: Add the new image and texture to the GLB
        # Create and add a new image to the glTF (same as before)
        new_image = GLTFImage()
        new_image.uri = os.path.join(os.path.dirname(__file__), f"{index}.png")
        gltf.images.append(new_image)



        # Create a new texture and associate it with the added image
        new_texture = Texture()
        new_texture.source = len(gltf.images) - 1  # Index of the newly added image
        new_texture.sampler = 0
        # set to nearest neighbor

        gltf.textures.append(new_texture)

        # Step 4: Assign the new texture to the appropriate material(s) or mesh(es)
        # Assuming you have a mesh/primitive that was using the old texture and you want to apply the new texture to it, you need to set the material index for that mesh/primitive.
        # Replace 0 with the actual index of the mesh/primitive you want to update.
        gltf.meshes[0].primitives[0].material = len(gltf.materials) - 1


        


        # Now you can convert the images to data URIs and save the updated GLB
        gltf.convert_images(ImageFormat.DATAURI)

        filename2 = "models/temp_model.glb"
        gltf.save(filename2)
    else:
        print("Invalid existing_texture_index")
    

running = False
time_since_last_run = 0


# Function to handle the input from the interface
def handle_input(seed_text, surprise_me):
    if surprise_me:
        images, used_seed = gen("Surprise Me")
    else:
        images, used_seed = gen(seed_text)


    loadMinecraftSkinViewer(0)








    return images, f"Seed Used: {used_seed}", os.path.join(os.path.dirname(__file__), "models/temp_model.glb")

# Function to load the Minecraft skin viewer
def loadMinecraftSkinViewer(index):
    update_model(index)






head = f"""
<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
<script>
function loadMinecraftSkinViewer(index) {{
    gr.interface.trigger("generate", index);
}}
</script>
"""



iface = gr.Interface(
    fn=handle_input,
    inputs=[
        gr.Text(label="Seed", placeholder="Enter a seed"),  # Text input for seed
        gr.Checkbox(label="Surprise Me", value=False),  # Checkbox for 'Surprise Me' functionality
        
    ],
    outputs=[
        gr.Gallery(label="Generated Skins"),
        gr.Textbox(label="Used Seed"),  # Display the seed used to generate the images
        gr.Model3D(
            clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model"),

        
        
    ],
    title = "Minecraft Skin Generator <style>img{image-rendering: pixelated;}</style>", #<-- EWW GROSS IK IK IM SORRY, BUT THAT'S THE ONLY WAY I FOUND IT TO WORK
    head=head, # Add the Babylon.js library to the interface
    allow_flagging=False,
    

    
)

iface.launch()