# UE-side screenshot renderer for visual review (commandlet, real RHI).
#
# Builds a small review scene in the default world (floor, sky, key light,
# MPFB character, the three scale-test cubes) and renders three fixed camera
# angles synchronously via SceneCapture2D render targets, writing PNGs
# directly into <repo>/Screenshots/ (gitignored).
#
# Invoked by Tools/render_screenshots.sh (UnrealEditor-Cmd -run=pythonscript).
# NOTE: no -nullrhi — rendering needs the real Vulkan RHI. In 5.8 the
# commandlet swallows print(); results go to /tmp/ue-screenshots-result.txt.
# Do NOT load_map here: after a map load the commandlet's spawn context dies
# and every spawn_actor call returns None.

import os

import unreal

RESULT = "/tmp/ue-screenshots-result.txt"
lines = []


def log(msg):
    lines.append(str(msg))


REPO = os.path.abspath(os.path.join(unreal.Paths.project_dir(), "..", ".."))
OUT_DIR = os.path.join(REPO, "Screenshots")
WORLD = unreal.EditorLevelLibrary.get_editor_world()


def spawn(cls, loc, label=None, scale=None, rot=None):
    a = unreal.EditorLevelLibrary.spawn_actor_from_class(cls, loc)
    if a is None:
        log(f"spawn {cls.__name__} returned None")
        return None
    if rot is not None:
        a.set_actor_rotation(rot, False)
    if scale is not None:
        a.set_actor_scale3d(scale)
    if label:
        a.set_actor_label(label)
    return a


def spawn_mesh(mesh, loc, label=None, scale=None):
    a = unreal.EditorLevelLibrary.spawn_actor_from_object(mesh, loc)
    if a is None:
        log(f"spawn mesh {label} returned None")
        return None
    if scale is not None:
        a.set_actor_scale3d(scale)
    if label:
        a.set_actor_label(label)
    return a


# --- stage content -----------------------------------------------------------

try:
    floor_mesh = unreal.EditorAssetLibrary.load_asset("/Game/Test/ScaleTest/ScaleTest_A")
    if floor_mesh:
        spawn_mesh(floor_mesh, unreal.Vector(450, 450, -50.0), "ReviewFloor", scale=(45.0, 45.0, 0.1))
        log("staged: floor (45 m plane from ScaleTest_A)")

    dl = spawn(unreal.DirectionalLight, unreal.Vector(0, 0, 900), "RenderKeyLight",
               rot=unreal.Rotator(pitch=-38.0, yaw=35.0, roll=0.0))
    if dl:
        log("staged: DirectionalLight")
    if spawn(unreal.SkyLight, unreal.Vector(0, 0, 1200), "RenderFillSky"):
        log("staged: SkyLight")
    if spawn(unreal.SkyAtmosphere, unreal.Vector(0, 0, 0), "RenderSky"):
        log("staged: SkyAtmosphere")

    char_mesh = unreal.EditorAssetLibrary.load_asset("/Game/Test/Character/character")
    if char_mesh:
        if spawn_mesh(char_mesh, unreal.Vector(0, 0, 0), "ReviewCharacter"):
            log("staged: MPFB character at origin")
    else:
        log("MISSING: /Game/Test/Character/character")

    for name, x, y in (("ScaleTest_A", 0.0, 0.0), ("ScaleTest_B", 400.0, 0.0), ("ScaleTest_C", 0.0, 400.0)):
        m = unreal.EditorAssetLibrary.load_asset(f"/Game/Test/ScaleTest/{name}")
        if m:
            spawn_mesh(m, unreal.Vector(x, y, 50.0), f"Cube_{name}")
            log(f"staged: {name} at ({x}, {y}, 50)")
except Exception as exc:  # noqa
    log(f"staging failed: {exc!r}")

# --- render ------------------------------------------------------------------


def render_shot(name, loc, target, fov=70.0):
    rt = unreal.RenderingLibrary.create_render_target2d(
        WORLD, 1920, 1080, unreal.TextureRenderTargetFormat.RTF_RGBA8_SRGB,
        unreal.LinearColor(0.0, 0.0, 0.0, 1.0))
    cap = spawn(unreal.SceneCapture2D, loc)
    if cap is None:
        log(f"{name}: no capture actor")
        return
    cap.set_actor_rotation(unreal.MathLibrary.find_look_at_rotation(loc, target), False)
    comp = cap.get_component_by_class(unreal.SceneCaptureComponent2D)
    comp.set_editor_property("texture_target", rt)
    comp.set_editor_property("capture_source", unreal.SceneCaptureSource.SCS_FINAL_COLOR_LDR)
    comp.set_editor_property("fov_angle", fov)
    comp.set_editor_property("capture_every_frame", False)
    comp.set_editor_property("capture_on_movement", False)
    try:
        pps = comp.get_editor_property("post_process_settings")
        pps.set_editor_property("auto_exposure_method", unreal.AutomaticExposureMethod.AEM_MANUAL)
        pps.set_editor_property("auto_exposure_bias", 1.0)
        comp.set_editor_property("post_process_settings", pps)
    except Exception as exc:  # noqa
        log(f"{name}: exposure control skipped ({exc!r})")
    comp.capture_scene()
    os.makedirs(OUT_DIR, exist_ok=True)
    ok = unreal.RenderingLibrary.export_render_target(WORLD, rt, OUT_DIR, name)
    out = os.path.join(OUT_DIR, f"{name}.png")
    size = os.path.getsize(out) if os.path.exists(out) else 0
    log(f"rendered {name}: export ok={ok}, {size} bytes")
    unreal.EditorLevelLibrary.destroy_actor(cap)


SHOTS = [
    ("01-overview", unreal.Vector(900, -1100, 500), unreal.Vector(200, 100, 80), 70.0),
    ("02-character", unreal.Vector(180, -260, 160), unreal.Vector(0, 0, 85), 55.0),
    ("03-scale-cubes", unreal.Vector(550, -700, 150), unreal.Vector(200, 200, 50), 60.0),
]

for shot_name, loc, target, fov in SHOTS:
    try:
        render_shot(shot_name, loc, target, fov)
    except Exception as exc:  # noqa
        log(f"shot {shot_name} failed: {exc!r}")

with open(RESULT, "w") as f:
    f.write("\n".join(lines) + "\n")