"""
Blender export script for CNN v3 G-buffer training data.
Configures render passes and a compositor File Output node,
then renders the current scene to a multi-layer EXR.

Usage (headless):
    blender -b scene.blend -P blender_export.py -- --output renders/frame_###

    # List available view layers in the blend file:
    blender -b scene.blend -P blender_export.py -- --view-layer ?

    # Use a specific view layer:
    blender -b scene.blend -P blender_export.py -- --output renders/frame_### --view-layer "MyLayer"

Each '#' in the output path is replaced by Blender with the frame number (zero-padded).
The script writes one multi-layer EXR per frame containing all required passes.

G-buffer pass mapping:
    Combined  → training target RGBA (beauty)
    DiffCol   → albedo.rgb  (pre-lighting material color)
    Normal    → normal.xy   (world-space, oct-encode in pack_blender_sample.py)
    Z         → depth       (view-space distance, normalize in pack step)
    IndexOB   → mat_id      (object index, u8 / 255)
    Shadow    → shadow      (invert: shadow=1 means fully lit)
    Alpha     → transp.     (0=opaque, 1=clear/transparent)
"""

import sys
import argparse

import bpy


def parse_args():
    # Blender passes its own argv; our args follow '--'.
    argv = sys.argv
    if "--" in argv:
        argv = argv[argv.index("--") + 1:]
    else:
        argv = []
    parser = argparse.ArgumentParser(
        description="Configure Blender render passes and export multi-layer EXR."
    )
    parser.add_argument(
        "--output",
        default="//renders/frame_###",
        help="Output path prefix (use ### for frame number padding). "
             "Default: //renders/frame_###",
    )
    parser.add_argument(
        "--width",  type=int, default=640,
        help="Render width in pixels (default: 640)"
    )
    parser.add_argument(
        "--height", type=int, default=360,
        help="Render height in pixels (default: 360)"
    )
    parser.add_argument(
        "--start-frame", type=int, default=None,
        help="First frame to render (default: scene start frame)"
    )
    parser.add_argument(
        "--end-frame", type=int, default=None,
        help="Last frame to render (default: scene end frame)"
    )
    parser.add_argument(
        "--view-layer", default=None, metavar="NAME",
        help="View layer name to use (default: first available). "
             "Use --view-layer without a value (or pass '?') to list available layers."
    )
    return parser.parse_args(argv)


def configure_scene(args):
    scene = bpy.context.scene

    # Render dimensions
    scene.render.resolution_x = args.width
    scene.render.resolution_y = args.height
    scene.render.resolution_percentage = 100

    # Frame range (optional override)
    if args.start_frame is not None:
        scene.frame_start = args.start_frame
    if args.end_frame is not None:
        scene.frame_end = args.end_frame

    # Use Cycles for best multi-pass support
    scene.render.engine = "CYCLES"

    # Enable required render passes on the selected view layer
    available = list(scene.view_layers.keys())
    if getattr(args, "view_layer", None) in (None, "?"):
        if args.view_layer == "?":
            print("Available view layers:", available)
            sys.exit(0)
        vl = scene.view_layers[0]
    else:
        try:
            vl = scene.view_layers[args.view_layer]
        except KeyError:
            print(f"ERROR: view layer '{args.view_layer}' not found.")
            print(f"Available: {available}")
            sys.exit(1)
    vl.use_pass_combined           = True   # beauty target
    vl.use_pass_diffuse_color      = True   # albedo
    vl.use_pass_normal             = True   # world normals
    vl.use_pass_z                  = True   # depth (Z)
    vl.use_pass_object_index       = True   # mat_id
    vl.use_pass_shadow             = True   # shadow catcher
    # Alpha is available via the combined pass alpha channel;
    # the compositor node below also taps it separately.

    print(f"[blender_export] Render passes configured on ViewLayer '{vl.name}'.")
    print(f"  Resolution: {args.width}x{args.height}")
    print(f"  Frames: {scene.frame_start} – {scene.frame_end}")


def configure_compositor(args):
    scene = bpy.context.scene
    scene.use_nodes = True
    tree = scene.node_tree

    # Clear all existing compositor nodes
    tree.nodes.clear()

    # Render Layers node (source of all passes)
    rl_node = tree.nodes.new("CompositorNodeRLayers")
    rl_node.location = (0, 0)

    # File Output node — multi-layer EXR (all passes in one file)
    out_node = tree.nodes.new("CompositorNodeOutputFile")
    out_node.location = (600, 0)
    out_node.format.file_format = "OPEN_EXR_MULTILAYER"
    out_node.format.exr_codec = "ZIP"
    out_node.base_path = args.output

    # Map each render pass socket to a named layer in the EXR.
    # Slot order matters: the first slot is created by default; we rename it
    # and add the rest.
    pass_sockets = [
        ("Image",           "Combined"),    # beauty / target
        ("Diffuse Color",   "DiffCol"),     # albedo
        ("Normal",          "Normal"),      # world normals
        ("Depth",           "Z"),           # view-space depth
        ("Object Index",    "IndexOB"),     # object index
        ("Shadow",          "Shadow"),      # shadow
        ("Alpha",           "Alpha"),       # transparency / alpha
    ]

    # The node starts with one default slot; configure it first.
    for i, (socket_name, layer_name) in enumerate(pass_sockets):
        if i == 0:
            # Rename the default slot
            out_node.file_slots[0].path = layer_name
        else:
            out_node.file_slots.new(layer_name)

        # Link render layer socket to file output slot
        src_socket = rl_node.outputs.get(socket_name)
        dst_socket = out_node.inputs[i]
        if src_socket:
            tree.links.new(src_socket, dst_socket)
        else:
            print(f"[blender_export] WARNING: pass socket '{socket_name}' "
                  f"not found on Render Layers node. Skipping.")

    print(f"[blender_export] Compositor configured. Output → {args.output}")
    print("  Layers: " + ", ".join(ln for _, ln in pass_sockets))


def main():
    args = parse_args()
    configure_scene(args)
    configure_compositor(args)

    # Trigger the render (only when running headless with -b)
    bpy.ops.render.render(animation=True)
    print("[blender_export] Render complete.")


if __name__ == "__main__":
    main()