#! /usr/bin/python3

import argparse
import contextlib
import math
import os
import os.path
import subprocess
import sys

OBJ_SIZE = 10.0

# Faster than Eevee, but not as pretty
RENDER_ENGINE = "BLENDER_WORKBENCH"

SUBPROC_ENV_FLAG = "__STNE_THUMBNAIL_SUBPROC__"


def deselect():
    bpy.ops.object.select_all(action="DESELECT")


def delete_obj(obj_name):
    obj = bpy.data.objects.get(obj_name)
    if obj:
        deselect()
        obj.select_set(True)
        bpy.ops.object.delete()


def import_stl(path):
    deselect()
    bpy.ops.import_mesh.stl(filepath=path)
    bpy.ops.object.origin_set(type="GEOMETRY_ORIGIN", center="BOUNDS")

    for mesh in bpy.context.selected_objects:
        maxd = max(mesh.dimensions)
        factor = OBJ_SIZE / maxd
        mesh.scale = (factor, factor, factor)

        mat = bpy.data.materials.new("color")
        mat.diffuse_color = (0.0, 0.25, 1, 1.0)
        mesh.data.materials.append(mat)

        mesh.rotation_euler = (
            math.radians(0),
            math.radians(0),
            math.radians(0),
        )

    bpy.ops.object.transform_apply()


def position_camera():
    camera_obj = bpy.data.objects["Camera"]
    camera_obj.location = (1, 1, 1)
    camera_obj.rotation_euler = (
        math.radians(55),
        math.radians(0),
        math.radians(-35),
    )

    bpy.ops.view3d.camera_to_view_selected()

    camera = bpy.data.cameras["Camera"]
    camera.lens -= 5


def position_light():
    camera = bpy.data.objects["Camera"]
    light = bpy.data.objects["Light"]

    light.location = camera.location
    light.location[0] -= OBJ_SIZE / 3
    light.location[1] -= OBJ_SIZE / 3
    light.location[2] += OBJ_SIZE / 3

    light.rotation_euler = camera.rotation_euler

    light = bpy.data.lights["Light"]
    light.energy = 5000
    light.shadow_soft_size = 15


def set_engine():
    for scene in bpy.data.scenes:
        scene.render.engine = RENDER_ENGINE

    bpy.context.scene.render.engine = RENDER_ENGINE


def render(input, output, size):
    render = bpy.data.scenes["Scene"].render
    render.filepath = output
    render.resolution_percentage = 100
    render.resolution_x = size
    render.resolution_y = size
    render.image_settings.color_mode = "RGBA"

    delete_obj("Cube")
    import_stl(input)
    position_camera()
    position_light()
    set_engine()

    bpy.ops.render.render(write_still=True)


def main_blender():
    parser = argparse.ArgumentParser()
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("size", type=int)
    args = parser.parse_args(sys.argv[sys.argv.index("--") + 1 :])

    render(args.input, args.output, args.size)


@contextlib.contextmanager
def pipes():
    r, w = os.pipe()
    with (
        os.fdopen(r, "rt") as rf,
        os.fdopen(w, "wt") as wf,
    ):
        yield rf, wf


@contextlib.contextmanager
def x11_cleanup():
    try:
        yield
    finally:
        # Best-effort cleanup
        with contextlib.suppress(OSError):
            os.rmdir("/tmp/.X11-unix")


@contextlib.contextmanager
def display_server(proc_out):
    try:
        # Is it possible to connect to a display server with the current
        # environment?
        subprocess.run(
            ["stne-display-probe"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=True,
            timeout=5.0,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
        # Could not connect to a display server, proceed with Xvfb
        pass
    else:
        yield {}
        return

    with (
        pipes() as (r, w),
        x11_cleanup(),
        subprocess.Popen(
            [
                "Xvfb",
                "-displayfd",
                "4",
                "-screen",
                "0",
                "1280x1024x24",
                "-nolisten",
                "tcp",
            ],
            stdout=proc_out,
            stderr=proc_out,
            pass_fds=(w.fileno(),),
        ) as proc,
    ):
        w.close()
        if display := r.read().strip():
            try:
                yield {"DISPLAY": f":{display}"}
            finally:
                proc.terminate()
                return

    raise Exception("Failed to start Xvfb")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--size", default=256, type=int)
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--debug", default=False, action="store_true")
    args = parser.parse_args()

    proc_out = None if args.debug else subprocess.DEVNULL

    with display_server(proc_out) as display_env:
        subprocess.run(
            [
                "blender",
                "--background",
                "--python",
                __file__,
                "--",
                args.input,
                args.output,
                str(args.size),
            ],
            check=True,
            stdout=proc_out,
            stderr=proc_out,
            env=(
                os.environ
                | display_env
                | {
                    SUBPROC_ENV_FLAG: "1",
                }
            ),
        )


if SUBPROC_ENV_FLAG in os.environ:
    import bpy

    main_blender()
else:
    main()
