# timeforge/utils.py

import tempfile
import trimesh
from trimesh.exchange.gltf import export_glb
import numpy as np
from PIL import Image

def apply_gradient_color(mesh_text):
    """
    Apply a gradient color to the mesh vertices and return as GLB.
    Args:
        mesh_text (str): The input mesh in OBJ format as a string.
    Returns:
        str: Path to the GLB file with gradient colors applied.
    """
    # Load the mesh
    temp_file =  tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
    with open(temp_file+".obj", "w") as f:
        f.write(mesh_text)
    mesh = trimesh.load_mesh(temp_file+".obj", file_type='obj')

    # Get vertex coordinates
    vertices = mesh.vertices
    y_values = vertices[:, 1]  # Y-axis values

    # Normalize Y values to range [0, 1] for color mapping
    y_normalized = (y_values - y_values.min()) / (y_values.max() - y_values.min())

    # Generate colors: Map normalized Y values to RGB gradient (e.g., blue to red)
    colors = np.zeros((len(vertices), 4))  # RGBA
    colors[:, 0] = y_normalized  # Red channel
    colors[:, 2] = 1 - y_normalized  # Blue channel
    colors[:, 3] = 1.0  # Alpha channel (fully opaque)

    # Attach colors to mesh vertices
    mesh.visual.vertex_colors = colors

    # Export to GLB format
    glb_path = temp_file+".glb"
    with open(glb_path, "wb") as f:
        f.write(export_glb(mesh))
    
    return glb_path

def create_image_grid(images):
    images = [Image.fromarray((img * 255).astype("uint8")) for img in images]

    width, height = images[0].size
    grid_img = Image.new("RGB", (2 * width, 2 * height))

    grid_img.paste(images[0], (0, 0))
    grid_img.paste(images[1], (width, 0))
    grid_img.paste(images[2], (0, height))
    grid_img.paste(images[3], (width, height))

    return grid_img