import argparse, io, math, os, sys
from pathlib import Path
from typing import Tuple

import numpy as np, requests
from PIL import Image
import trimesh, pyrender
from trimesh.transformations import translation_matrix, rotation_matrix  # NEW

os.environ.setdefault("PYOPENGL_PLATFORM", "egl")


def download_to_memory(url: str) -> bytes:
    resp = requests.get(url, timeout=30)
    resp.raise_for_status()
    return resp.content


def load_mesh(data: bytes) -> trimesh.Trimesh:
    return trimesh.load(io.BytesIO(data), file_type="glb", force="mesh")


def build_scene(mesh: trimesh.Trimesh) -> Tuple[pyrender.Scene, float]:
    tm_mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
    scene = pyrender.Scene(bg_color=[1, 1, 1, 0])
    scene.add(tm_mesh)

    bb = mesh.bounding_box_oriented.extents          # fixed name
    if not np.all(bb):
        bb = mesh.extents
    radius = np.linalg.norm(bb) * 0.6
    return scene, radius


def add_lighting(scene: pyrender.Scene, radius: float) -> None:
    key  = pyrender.PointLight(color=np.ones(3), intensity=40.0)
    fill = pyrender.PointLight(color=np.ones(3), intensity=20.0)
    back = pyrender.PointLight(color=np.ones(3), intensity=10.0)

    scene.add(key,  pose=translation_matrix([ radius,  radius,  radius]))
    scene.add(fill, pose=translation_matrix([-radius,  radius,  radius]))
    scene.add(back, pose=translation_matrix([     0, -radius, -radius]))


def setup_camera(scene: pyrender.Scene, radius: float) -> None:
    cam = pyrender.PerspectiveCamera(yfov=np.radians(45.0))
    cam_pose = translation_matrix([0, 0, radius * 2.5])
    scene.add(cam, pose=cam_pose)


def render_thumbnail(scene: pyrender.Scene, size: int) -> Image.Image:
    r = pyrender.OffscreenRenderer(viewport_width=size, viewport_height=size)
    try:
        color, _ = r.render(scene)
    finally:
        r.delete()
    return Image.fromarray(color)


def generate_thumbnail(url: str, out_path: Path, size: int = 512) -> None:
    raw = download_to_memory(url)
    mesh = load_mesh(raw)
    
    # Get mesh dimensions before any transformations
    original_extents = mesh.extents
    longest_dimension = np.max(original_extents)
    
    # Print dimension info
    print(f"Mesh dimensions (X, Y, Z): {original_extents}")
    print(f"Longest dimension: {longest_dimension:.4f}")
    
    # Scaling constant - you can use this to normalize models to a target size
    target_size = 2.5  # Target longest dimension
    scale_factor = target_size / longest_dimension
    print(f"Scale factor to normalize to {target_size}: {scale_factor:.4f}")
    
    # Calculate radius BEFORE scaling for consistent camera/lighting positioning
    bb = mesh.bounding_box_oriented.extents
    if not np.all(bb):
        bb = mesh.extents
    fixed_radius = np.linalg.norm(bb) * 0.6
    
    mesh.apply_translation(-mesh.bounding_box.centroid)
    
    # Apply the scaling transformation
    mesh.apply_scale(scale_factor)
    
    # Rotate 45 degrees to the left (around Y-axis)
    rotation = rotation_matrix(np.radians(30), [0.3, -0.5, 0])
    mesh.apply_transform(rotation)

    # Build scene but use fixed radius for camera/lighting
    tm_mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
    scene = pyrender.Scene(bg_color=[0.15, 0.15, 0.15, 1])  # Gray background
    scene.add(tm_mesh)
    
    add_lighting(scene, fixed_radius)
    setup_camera(scene, fixed_radius)
    img = render_thumbnail(scene, size)

    img.save(out_path, "PNG")
    print(f"Saved thumbnail → {out_path.resolve()}")


def parse_args(argv):
    p = argparse.ArgumentParser(description="Render a GLB file to a PNG thumbnail")
    p.add_argument("url"), p.add_argument("output")
    p.add_argument("size", nargs="?", type=int, default=512)
    return p.parse_args(argv)


if __name__ == "__main__":
    args = parse_args(sys.argv[1:])
    generate_thumbnail(args.url, Path(args.output), args.size)