"""
Blender export script for CNN v3 G-buffer training data.
Enables render passes and uses native OPEN_EXR_MULTILAYER output.

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

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

Output: one multilayer EXR per frame at {output}/0001.exr, 0002.exr, ...
Passes: Combined (beauty), DiffCol (albedo), Normal, Depth, IndexOB, Shadow.
"""

import contextlib
import os
import sys
import argparse

import bpy


@contextlib.contextmanager
def _suppress_blender_stdout():
    """Redirect C-level stdout to /dev/null to suppress Fra:/Mem: render spam."""
    sys.stdout.flush()
    saved = os.dup(1)
    null = os.open('/dev/null', os.O_WRONLY)
    os.dup2(null, 1)
    os.close(null)
    try:
        yield
    finally:
        sys.stdout.flush()
        os.dup2(saved, 1)
        os.close(saved)


def parse_args():
    # Blender passes its own argv; our args follow '--'.
    argv = sys.argv
    argv = argv[argv.index("--") + 1:] if "--" in argv else []
    parser = argparse.ArgumentParser(
        description="Configure Blender render passes and export multi-layer EXR."
    )
    parser.add_argument(
        "--output", default="//renders/",
        help="Output directory; frames are written as 0001.exr, 0002.exr, … "
             "Use // for blend file directory. Default: //renders/",
    )
    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). Pass '?' to list available layers.",
    )
    return parser.parse_args(argv)


def _die(msg, candidates=None):
    print(f"ERROR: {msg}")
    if candidates is not None:
        print("Candidate attributes:", candidates)
    sys.exit(1)


def _blender_version():
    """Return (major, minor) tuple, e.g. (4, 5) for Blender 4.5."""
    return bpy.app.version[:2]


def _is_blender5():
    """True on Blender 5.0+."""
    return _blender_version()[0] >= 5


def _resolve_view_layer(scene, name):
    """Return the requested view layer, defaulting to index 0. '?' lists and exits."""
    available = list(scene.view_layers.keys())
    if name == "?":
        print("Available view layers:", available)
        sys.exit(0)
    if name is None:
        return scene.view_layers[0]
    try:
        return scene.view_layers[name]
    except KeyError:
        _die(f"view layer '{name}' not found. Available: {available}")



def configure_scene(args):
    """Configure render settings and enable G-buffer passes."""
    scene = bpy.context.scene

    major, minor = _blender_version()
    print(f"[blender_export] Blender {major}.{minor}")
    if _is_blender5():
        print("  NOTE: Blender 5.x — using native OPEN_EXR_MULTILAYER output.")

    scene.render.engine = "CYCLES"
    scene.render.resolution_x = args.width
    scene.render.resolution_y = args.height
    scene.render.resolution_percentage = 100
    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

    vl = _resolve_view_layer(scene, args.view_layer)
    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
    # Alpha is available via the combined pass alpha channel.

    # Use native multilayer EXR output — all enabled passes are included
    # automatically.  This is more reliable than the compositor FileOutput
    # approach (which requires per-pass socket wiring that varies by Blender
    # version and may silently omit passes).
    scene.render.image_settings.file_format = "OPEN_EXR_MULTILAYER"
    scene.render.image_settings.exr_codec = "ZIP"
    out_abs = os.path.abspath(args.output)
    os.makedirs(out_abs, exist_ok=True)
    scene.render.filepath = out_abs + os.sep  # trailing sep → {out}/0001.exr

    # Disable compositor post-processing so we capture raw G-buffer data.
    if hasattr(scene.render, "use_compositing"):
        scene.render.use_compositing = False

    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}")
    print(f"  Output: {out_abs}/")


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

    # Render with render-progress output suppressed.  Per-frame status is
    # printed via a render_post handler to stderr (always visible).
    out_abs = os.path.abspath(args.output)

    def _on_render_post(scene, depsgraph=None):
        print(f"[blender_export] Frame {scene.frame_current} done.", file=sys.stderr, flush=True)

    bpy.app.handlers.render_post.append(_on_render_post)
    try:
        with _suppress_blender_stdout():
            bpy.ops.render.render(animation=True)
    finally:
        bpy.app.handlers.render_post.remove(_on_render_post)

    print("[blender_export] Render complete.")

    # Print the next-step pack command.
    pack_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pack_blender_sample.py")
    print("\n[blender_export] Next step — pack EXRs into sample directories:")
    print(f"  for exr in {out_abs}/*.exr; do")
    print(f'    name=$(basename "${{exr%.exr}}")')
    print(f'    python3 {pack_script} --exr "$exr" --output {out_abs}/$name/')
    print( "  done")


if __name__ == "__main__":
    main()