#!/usr/bin/env python3
"""Render a model's card thumbnail without a GPU.

Usage: render_thumbnail.py <input.stl|.obj|.3mf> <output.png> [--size 512]
Prints {png, faces} (or {error}) as JSON.

3MF projects show their first plate with parts in their filament colours. Big
meshes are simplified first, then triangles are painted back to front with
Pillow, which is plenty for a card image.
"""
import argparse
import json

import fast_simplification
import numpy as np
import trimesh
from PIL import Image, ImageDraw

import threemf_stream

MAX_FACES = 150_000
GREY = (196, 202, 214)
SUPERSAMPLE = 2


def simplify(vertices, faces, ratio):
    if ratio >= 1.0 or len(faces) <= 400:
        return vertices, faces
    v, f = fast_simplification.simplify(
        np.asarray(vertices, dtype=np.float32), np.asarray(faces, dtype=np.int32), target_reduction=1 - ratio)
    return np.asarray(v, dtype=np.float64), np.asarray(f)


def hex_rgb(colour):
    return tuple(int(colour[i:i + 2], 16) for i in (1, 3, 5)) if colour else GREY


def load_3mf(path):
    meshes, placements = threemf_stream.load(path)
    project = threemf_stream.read_project(path)
    placements = threemf_stream.printable(project, placements)
    if len(project["plates"]) > 1:
        placements = threemf_stream.on_plate(project, placements, project["plates"][0]["index"]) or placements
    if not placements:
        raise ValueError("no geometry found")

    # Simplify each distinct part once, by the same ratio, so instances stay cheap.
    total = sum(len(meshes[p.key][1]) for p in placements)
    ratio = min(1.0, MAX_FACES / total)
    small = {}
    vs, fs, cs, offset = [], [], [], 0
    for p in placements:
        if p.key not in small:
            small[p.key] = simplify(*meshes[p.key], ratio)
        v, f = small[p.key]
        t = p.transform
        if np.linalg.det(t[:3, :3]) < 0:
            f = f[:, ::-1]  # a mirrored placement flips winding; flip it back so faces point outward
        vs.append(v @ t[:3, :3].T + t[:3, 3])
        fs.append(f + offset)
        cs.append(np.tile(hex_rgb(threemf_stream.filament_colour(project, p)), (len(f), 1)))
        offset += len(v)
    return np.vstack(vs), np.vstack(fs), np.vstack(cs)


def load_mesh(path):
    mesh = trimesh.load(path, force="mesh")
    if mesh is None or len(mesh.faces) == 0:
        raise ValueError("no geometry found")
    v, f = simplify(np.asarray(mesh.vertices), np.asarray(mesh.faces), min(1.0, MAX_FACES / len(mesh.faces)))
    if path.lower().endswith(".obj"):
        v = np.column_stack([v[:, 0], -v[:, 2], v[:, 1]])  # OBJ is Y-up; everything else here is Z-up
    return v, f, np.tile(GREY, (len(f), 1))


def render(vertices, faces, colours, size):
    # Three-quarter view from the front right, a little above.
    az, el = np.radians(-40), np.radians(30)
    forward = -np.array([np.cos(el) * np.sin(az), -np.cos(el) * np.cos(az), np.sin(el)])
    right = np.cross(forward, [0, 0, 1])
    right /= np.linalg.norm(right)
    up = np.cross(right, forward)

    a, b, c = vertices[faces[:, 0]], vertices[faces[:, 1]], vertices[faces[:, 2]]
    normals = np.cross(b - a, c - a)
    # A mesh exported inside-out has negative volume: turn its normals around.
    if np.einsum("ij,ij->i", a, np.cross(b, c)).sum() < 0:
        normals = -normals
    normals /= np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12

    # Only faces pointing at the camera are drawn. Painting back faces too lets a
    # large far triangle land on top of a near one when sorted by centre depth.
    facing = normals @ forward < 0
    if facing.any():
        a, b, c, normals, colours = a[facing], b[facing], c[facing], normals[facing], colours[facing]

    light = -forward + 0.6 * up - 0.4 * right
    light /= np.linalg.norm(light)
    shade = 0.32 + 0.68 * np.clip(normals @ light, 0, 1)
    fills = np.clip(colours * shade[:, None], 0, 255).astype(np.uint8)

    xs = np.stack([a @ right, b @ right, c @ right], axis=1)
    ys = np.stack([a @ up, b @ up, c @ up], axis=1)
    lo = np.array([xs.min(), ys.min()])
    span = np.array([xs.max(), ys.max()]) - lo
    width = size * SUPERSAMPLE
    scale = width * 0.86 / max(span.max(), 1e-9)
    offset = (width - span * scale) / 2
    px = (xs - lo[0]) * scale + offset[0]
    py = width - ((ys - lo[1]) * scale + offset[1])

    image = Image.new("RGBA", (width, width), (0, 0, 0, 0))
    draw = ImageDraw.Draw(image)
    depth = ((a + b + c) / 3) @ forward
    for i in np.argsort(-depth):  # far to near
        fill = tuple(int(x) for x in fills[i]) + (255,)
        draw.polygon([(px[i, 0], py[i, 0]), (px[i, 1], py[i, 1]), (px[i, 2], py[i, 2])], fill=fill, outline=fill)
    return image.resize((size, size), Image.LANCZOS)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--size", type=int, default=512)
    args = parser.parse_args()
    try:
        load = load_3mf if args.input.lower().endswith(".3mf") else load_mesh
        vertices, faces, colours = load(args.input)
        render(vertices, faces, colours, args.size).save(args.output, "PNG")
    except Exception as exc:
        print(json.dumps({"error": f"could not render thumbnail: {str(exc)[:200]}"}))
        return 1
    print(json.dumps({"png": args.output, "faces": int(len(faces))}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())