#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sys

try:
    import bpy
    from mathutils import Vector
except ImportError:
    print('This script must run inside Blender.')
    sys.exit(1)


def parse_args():
    argv = sys.argv
    argv = argv[argv.index('--') + 1:] if '--' in argv else []
    ap = argparse.ArgumentParser(description='Render NPZ pointcloud assets with Blender Workbench')
    ap.add_argument('--asset_npz', type=str, required=True)
    ap.add_argument('--out_png', type=str, required=True)
    ap.add_argument('--out_blend', type=str, default=None)
    ap.add_argument(
        '--background',
        choices=('light', 'dark', 'white'),
        default=None,
        help="World background: white=pure #fff (default when meta omits); light=legacy off-white; dark=legacy dark gray",
    )
    return ap.parse_args(argv)


def camera_from_coord(coord):
    import numpy as np
    mn = coord.min(axis=0)
    mx = coord.max(axis=0)
    center = (mn + mx) * 0.5
    diag = float(np.linalg.norm(mx - mn)) + 1e-6
    off = Vector((diag * 0.9, -diag * 0.85, diag * 0.55))
    loc = Vector((float(center[0]), float(center[1]), float(center[2]))) + off
    direction = Vector((float(center[0]), float(center[1]), float(center[2]))) - loc
    rot = direction.to_track_quat('-Z', 'Y')
    mat = rot.to_matrix().to_4x4()
    mat.translation = loc
    right = mat.col[0].xyz.copy()
    up = mat.col[1].xyz.copy()
    return loc, rot, right, up, diag


def build_triangle_mesh(coord, rgb_u8, right, up, diag):
    import numpy as np
    n = int(coord.shape[0])
    scale = float(diag) * 0.0025
    r = np.array([right.x, right.y, right.z], dtype=np.float64)
    u = np.array([up.x, up.y, up.z], dtype=np.float64)
    v0 = coord + (0.866 * r + 0.0 * u) * scale
    v1 = coord + (-0.433 * r + 0.75 * u) * scale
    v2 = coord + (-0.433 * r - 0.75 * u) * scale
    verts = np.stack([v0, v1, v2], axis=1).reshape(-1, 3)
    faces = [(i * 3, i * 3 + 1, i * 3 + 2) for i in range(n)]
    mesh = bpy.data.meshes.new('StoryPoints')
    mesh.from_pydata([tuple(v) for v in verts], [], faces)
    mesh.update()
    attr = mesh.color_attributes.new(name='Col', type='FLOAT_COLOR', domain='POINT')
    col = (rgb_u8.astype(np.float32) / 255.0).repeat(3, axis=0)
    for i in range(len(attr.data)):
        attr.data[i].color = (float(col[i, 0]), float(col[i, 1]), float(col[i, 2]), 1.0)
    return mesh


def setup_camera(scene, loc, rot):
    bpy.ops.object.camera_add(location=loc)
    cam = bpy.context.active_object
    cam.rotation_mode = 'QUATERNION'
    cam.rotation_quaternion = rot
    cam.data.lens = 35
    scene.camera = cam


def set_background(scene, mode):
    world = scene.world or bpy.data.worlds.new('World')
    scene.world = world
    world.use_nodes = True
    bg = world.node_tree.nodes.get('Background')
    if bg is None:
        bg = world.node_tree.nodes.new(type='ShaderNodeBackground')
    if mode == 'dark':
        bg.inputs[0].default_value = (0.02, 0.02, 0.025, 1.0)
        bg.inputs[1].default_value = 1.0
    elif mode == 'light':
        bg.inputs[0].default_value = (0.95, 0.95, 0.96, 1.0)
        bg.inputs[1].default_value = 1.0
    else:
        bg.inputs[0].default_value = (1.0, 1.0, 1.0, 1.0)
        bg.inputs[1].default_value = 1.0


def render_png(scene, path, dark):
    scene.render.engine = 'BLENDER_WORKBENCH'
    scene.render.resolution_x = 1920
    scene.render.resolution_y = 1080
    scene.render.image_settings.file_format = 'PNG'
    scene.render.image_settings.color_mode = 'RGB'
    scene.render.film_transparent = False
    scene.display.shading.light = 'STUDIO'
    scene.display.shading.color_type = 'VERTEX'
    scene.display.shading.show_cavity = True
    scene.display.shading.cavity_type = 'BOTH'
    scene.display.shading.curvature_ridge_factor = 0.5
    scene.display.shading.curvature_valley_factor = 0.5
    if dark:
        scene.display.shading.single_color = (0.05, 0.05, 0.05)
    scene.render.filepath = path
    bpy.ops.render.render(write_still=True)


def main():
    import numpy as np
    args = parse_args()
    asset = np.load(args.asset_npz)
    coord = asset['coord'].astype(np.float32)
    rgb = asset['rgb'].astype(np.uint8)
    meta_path = os.path.splitext(args.asset_npz)[0] + '.json'
    meta = {}
    if os.path.isfile(meta_path):
        meta = json.loads(open(meta_path, 'r', encoding='utf-8').read())
    bg = args.background or meta.get('background', 'light')
    if args.background is None:
        if bg in ('light', 'dark'):
            bg = 'white'

    bpy.ops.wm.read_factory_settings(use_empty=True)
    loc, rot, right, up, diag = camera_from_coord(coord)
    mesh = build_triangle_mesh(coord, rgb, right, up, diag)
    obj = bpy.data.objects.new('StoryObject', mesh)
    bpy.context.collection.objects.link(obj)
    setup_camera(bpy.context.scene, loc, rot)
    set_background(bpy.context.scene, bg)

    out_png = os.path.abspath(args.out_png)
    os.makedirs(os.path.dirname(out_png) or '.', exist_ok=True)
    render_png(bpy.context.scene, out_png, bg == 'dark')
    print('BLENDER_PNG', out_png)

    if args.out_blend:
        out_blend = os.path.abspath(args.out_blend)
        os.makedirs(os.path.dirname(out_blend) or '.', exist_ok=True)
        bpy.ops.wm.save_as_mainfile(filepath=out_blend)
        print('BLENDER_BLEND', out_blend)


if __name__ == '__main__':
    main()