"""Seamless PBR texture bakes for Champions of Luneraria (U5A, docs/port/U5-VISUAL.md).

Procedural material bakes in Blender (Cycles, 1 sample, emission of the map), not bitmaps
scaled up. Every set is seamless on all four edges (torus mapping like phase10/bake_bark.py).
U5.3: the two torus radii set the grain. Unequal radii stretch the noise along v (the bark
grain of phase10); U5 used about 10:1 for every set, which drew vertical streaks into stone,
earth, leaves and the river. Sets without a grain now use one radius (isotropic()), equal in
area to the former pair, so the mean feature size stays; bark, wood and rock strata keep theirs.
Writes 8-bit PNGs to Assets/Luneraria/Art/Textures/:

  <name>-albedo-2k.png    base colour
  <name>-normal-2k.png    tangent-space normal from the height field (Sobel)
  <name>-orm-2k.png       R occlusion, G roughness, B metallic (Unity-style ORM)

Sets: limestone (Kluft stone), soil (packed earth), bark (oak ridges), leaf (foliage),
rock (eroded strata), water (wave normals).

Usage (from the Unity project root):
  sh tools/blender/bake_pbr.sh
  sh tools/blender/bake_pbr.sh limestone,soil
Exit code 0 only if every set passed.
"""
import bpy
import math
import struct
import sys
import zlib
from pathlib import Path

UNITY_ROOT = Path(__file__).resolve().parents[2]
OUT = UNITY_ROOT / "Assets" / "Luneraria" / "Art" / "Textures"
SIZE = 2048


def args():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    result = {"sets": "all"}
    for i in range(0, len(argv) - 1, 2):
        result[argv[i].lstrip("-")] = argv[i + 1]
    return result


def linear(hexcolor):
    rgb = [int(hexcolor[i:i + 2], 16) / 255 for i in (0, 2, 4)]
    return tuple((c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4) for c in rgb) + (1.0,)


def math_node(nodes, links, op, a, b=None):
    obj = nodes.new("ShaderNodeMath")
    obj.operation = op
    if isinstance(a, (float, int)):
        obj.inputs[0].default_value = float(a)
    else:
        links.new(a, obj.inputs[0])
    if b is not None:
        if isinstance(b, (float, int)):
            obj.inputs[1].default_value = float(b)
        else:
            links.new(b, obj.inputs[1])
    return obj.outputs[0]


def torus_uv(nodes, links, scale_u, scale_v):
    """UV -> a circle in (cos u, sin u, cos v) so noise wraps on every edge."""
    uv = nodes.new("ShaderNodeTexCoord")
    split = nodes.new("ShaderNodeSeparateXYZ")
    links.new(uv.outputs["UV"], split.inputs[0])
    u = math_node(nodes, links, "MULTIPLY", split.outputs["X"], math.tau)
    v = math_node(nodes, links, "MULTIPLY", split.outputs["Y"], math.tau)
    co = nodes.new("ShaderNodeCombineXYZ")
    links.new(math_node(nodes, links, "MULTIPLY", math_node(nodes, links, "COSINE", u), scale_u), co.inputs["X"])
    links.new(math_node(nodes, links, "MULTIPLY", math_node(nodes, links, "SINE", u), scale_u), co.inputs["Y"])
    links.new(math_node(nodes, links, "MULTIPLY", math_node(nodes, links, "COSINE", v), scale_v), co.inputs["Z"])
    w = math_node(nodes, links, "MULTIPLY", math_node(nodes, links, "SINE", v), scale_v)
    return co.outputs[0], w


def isotropic(scale_u, scale_v):
    """One torus radius with the area of (scale_u, scale_v): same mean feature size, no grain."""
    radius = math.sqrt(scale_u * scale_v)
    return radius, radius


def noise(nodes, links, co, w, scale, detail, roughness):
    tex = nodes.new("ShaderNodeTexNoise")
    tex.noise_dimensions = "4D"
    links.new(co, tex.inputs["Vector"])
    links.new(w, tex.inputs["W"])
    tex.inputs["Scale"].default_value = scale
    tex.inputs["Detail"].default_value = detail
    tex.inputs["Roughness"].default_value = roughness
    return tex.outputs["Fac"]


def ramp(nodes, links, source, stops):
    """ValToRGB with (position, hex) stops in sRGB; output linear colours."""
    node_ramp = nodes.new("ShaderNodeValToRGB")
    elements = node_ramp.color_ramp.elements
    while len(elements) > 2:
        elements.remove(elements[-1])
    elements[0].position = stops[0][0]
    elements[0].color = linear(stops[0][1])
    elements[1].position = stops[-1][0]
    elements[1].color = linear(stops[-1][1])
    for position, hexcolor in stops[1:-1]:
        elements.new(position).color = linear(hexcolor)
    links.new(source, node_ramp.inputs[0])
    return node_ramp.outputs[0]


def new_scene(name):
    scene = bpy.data.scenes.new(name)
    bpy.context.window.scene = scene
    scene.render.engine = "CYCLES"
    scene.cycles.samples = 1
    scene.render.resolution_x = SIZE
    scene.render.resolution_y = SIZE
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = "PNG"
    scene.render.image_settings.color_mode = "RGB"
    scene.render.image_settings.color_depth = "8"
    scene.render.image_settings.compression = 90
    scene.render.dither_intensity = 0
    scene.view_settings.view_transform = "Standard"
    scene.view_settings.look = "None"
    scene.render.film_transparent = False
    bpy.ops.mesh.primitive_plane_add(size=2)
    plane = bpy.context.object
    plane.name = name + "_surface"
    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    plane.data.materials.append(mat)
    nodes = mat.node_tree.nodes
    nodes.clear()
    links = mat.node_tree.links
    emit = nodes.new("ShaderNodeEmission")
    output = nodes.new("ShaderNodeOutputMaterial")
    links.new(emit.outputs[0], output.inputs["Surface"])
    bpy.ops.object.camera_add(location=(0, 0, 3))
    camera = bpy.context.object
    camera.data.type = "ORTHO"
    camera.data.ortho_scale = 2
    scene.camera = camera
    return scene, nodes, links, emit


def render(scene, emit, source, path):
    mat = bpy.data.materials.get(scene.name)
    links = mat.node_tree.links
    for link in list(emit.inputs["Color"].links):
        links.remove(link)
    if isinstance(source, (float, int)):
        emit.inputs["Color"].default_value = (float(source), float(source), float(source), 1.0)
    else:
        links.new(source, emit.inputs["Color"])
    scene.render.filepath = str(path)
    bpy.ops.render.render(write_still=True)
    print("PBR_WRITE", path.name)


# ---------------------------------------------------------------- PNG helpers (no numpy)
def read_png_gray(path):
    """Returns (width, height, list of 0..255 luminance samples, row-major)."""
    image = bpy.data.images.load(str(path))
    width, height = image.size
    pixels = list(image.pixels)
    bpy.data.images.remove(image)
    values = [0] * (width * height)
    for i in range(width * height):
        r = pixels[i * 4]
        g = pixels[i * 4 + 1]
        b = pixels[i * 4 + 2]
        values[i] = int(round((0.2126 * r + 0.7152 * g + 0.0722 * b) * 255))
    return width, height, values


def write_png_rgb(path, width, height, pixels_rgb):
    """pixels_rgb: bytes-like length width*height*3, top-left origin."""
    raw = bytearray()
    stride = width * 3
    for y in range(height):
        raw.append(0)  # filter none
        raw.extend(pixels_rgb[y * stride:(y + 1) * stride])
    compressed = zlib.compress(bytes(raw), 9)

    def chunk(tag, data):
        return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)

    png = b"\x89PNG\r\n\x1a\n"
    png += chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
    png += chunk(b"IDAT", compressed)
    png += chunk(b"IEND", b"")
    path.write_bytes(png)
    print("PBR_WRITE", path.name)


def height_to_normal(path_height, path_normal):
    width, height, values = read_png_gray(path_height)

    def sample(x, y):
        return values[(y % height) * width + (x % width)]

    strength = 2.0
    out = bytearray(width * height * 3)
    for y in range(height):
        for x in range(width):
            left = sample(x - 1, y)
            right = sample(x + 1, y)
            up = sample(x, y - 1)
            down = sample(x, y + 1)
            dx = (left - right) * strength / 255.0
            dy = (up - down) * strength / 255.0
            # Tangent-space normal, OpenGL convention (Y up).
            nx, ny, nz = dx, dy, 1.0
            length = math.sqrt(nx * nx + ny * ny + nz * nz) or 1.0
            nx, ny, nz = nx / length, ny / length, nz / length
            i = (y * width + x) * 3
            out[i] = int(round((nx * 0.5 + 0.5) * 255))
            out[i + 1] = int(round((ny * 0.5 + 0.5) * 255))
            out[i + 2] = int(round((nz * 0.5 + 0.5) * 255))
    write_png_rgb(path_normal, width, height, out)


def pack_orm(path_ao, path_rough, path_metal, path_orm):
    _, _, ao = read_png_gray(path_ao)
    width, height, rough = read_png_gray(path_rough)
    _, _, metal = read_png_gray(path_metal)
    out = bytearray(width * height * 3)
    for i in range(width * height):
        out[i * 3] = ao[i]
        out[i * 3 + 1] = rough[i]
        out[i * 3 + 2] = metal[i]
    write_png_rgb(path_orm, width, height, out)


def save_set(name, builder):
    scene, nodes, links, emit = new_scene("PBR_" + name)
    height, albedo, roughness, metallic, ao = builder(nodes, links)
    albedo_path = OUT / (name + "-albedo-2k.png")
    height_path = OUT / (name + "-height-2k.png")
    rough_path = OUT / (name + "-roughness-2k.png")
    metal_path = OUT / (name + "-metallic-2k.png")
    ao_path = OUT / (name + "-ao-2k.png")
    render(scene, emit, albedo, albedo_path)
    render(scene, emit, height, height_path)
    render(scene, emit, roughness, rough_path)
    render(scene, emit, metallic if not isinstance(metallic, (float, int)) else float(metallic), metal_path)
    render(scene, emit, ao if not isinstance(ao, (float, int)) else float(ao), ao_path)
    height_to_normal(height_path, OUT / (name + "-normal-2k.png"))
    pack_orm(ao_path, rough_path, metal_path, OUT / (name + "-orm-2k.png"))
    height_path.unlink()
    rough_path.unlink()
    metal_path.unlink()
    ao_path.unlink()
    print("PBR_BAKE_PASS", name)


# ---------------------------------------------------------------- sets
def build_limestone(nodes, links):
    co, w = torus_uv(nodes, links, *isotropic(9.0, 0.8))
    broad = noise(nodes, links, co, w, 2.2, 6, 0.7)
    grain = noise(nodes, links, co, w, 11.0, 5, 0.62)
    cracks = noise(nodes, links, co, w, 5.5, 8, 0.8)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", broad, 0.55),
                       math_node(nodes, links, "MULTIPLY", grain, 0.35))
    height = math_node(nodes, links, "SUBTRACT", height, math_node(nodes, links, "MULTIPLY", cracks, 0.2))
    albedo = ramp(nodes, links, height, [
        (0.22, "6E675C"), (0.4, "8C8476"), (0.55, "A39A88"), (0.72, "BDB3A0"), (0.88, "D2C8B4")])
    roughness = ramp(nodes, links, height, [(0.2, "F0F0F0"), (0.5, "D0D0D0"), (0.8, "B8B8B8")])
    ao = ramp(nodes, links, height, [(0.15, "707070"), (0.45, "D0D0D0"), (0.8, "FFFFFF")])
    return height, albedo, roughness, 0.0, ao


def build_soil(nodes, links):
    co, w = torus_uv(nodes, links, *isotropic(7.0, 0.7))
    broad = noise(nodes, links, co, w, 1.6, 5, 0.7)
    grain = noise(nodes, links, co, w, 14.0, 4, 0.55)
    pebbles = noise(nodes, links, co, w, 6.5, 6, 0.75)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", broad, 0.5),
                       math_node(nodes, links, "MULTIPLY", grain, 0.3))
    height = math_node(nodes, links, "ADD", height, math_node(nodes, links, "MULTIPLY", pebbles, 0.2))
    albedo = ramp(nodes, links, height, [
        (0.2, "2A241C"), (0.4, "3C3428"), (0.58, "4A4030"), (0.75, "5C503C"), (0.9, "6E6048")])
    roughness = ramp(nodes, links, height, [(0.2, "FAFAFA"), (0.6, "E8E8E8"), (0.9, "D8D8D8")])
    ao = ramp(nodes, links, height, [(0.15, "808080"), (0.5, "D8D8D8"), (0.85, "FFFFFF")])
    return height, albedo, roughness, 0.0, ao


def build_bark(nodes, links):
    co, w = torus_uv(nodes, links, 10.0, 0.7)
    grain = noise(nodes, links, co, w, 2.6, 7, 0.73)
    ridges = noise(nodes, links, co, w, 7.5, 5, 0.8)
    height = ramp(nodes, links, grain, [
        (0.27, "121212"), (0.43, "242424"), (0.51, "8C8C8C"), (0.67, "CCCCCC")])
    bump = math_node(nodes, links, "ADD", height, math_node(nodes, links, "MULTIPLY", ridges, 0.25))
    albedo = ramp(nodes, links, height, [
        (0.0, "302B23"), (0.18, "4B3C2A"), (0.51, "766041"), (0.74, "94855D"), (1.0, "B09A6D")])
    roughness = ramp(nodes, links, ridges, [(0.2, "F5F5F5"), (0.7, "E0E0E0")])
    ao = ramp(nodes, links, height, [(0.1, "606060"), (0.5, "C8C8C8"), (0.85, "FFFFFF")])
    return bump, albedo, roughness, 0.0, ao


def build_leaf(nodes, links):
    co, w = torus_uv(nodes, links, *isotropic(8.0, 0.8))
    broad = noise(nodes, links, co, w, 1.8, 4, 0.65)
    vein = noise(nodes, links, co, w, 9.0, 5, 0.7)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", broad, 0.6),
                       math_node(nodes, links, "MULTIPLY", vein, 0.4))
    albedo = ramp(nodes, links, height, [
        (0.2, "1E3224"), (0.4, "2F4A32"), (0.58, "456640"), (0.75, "5F8554"), (0.9, "7FA66A")])
    roughness = ramp(nodes, links, height, [(0.2, "D8D8D8"), (0.7, "C0C0C0")])
    ao = ramp(nodes, links, height, [(0.15, "707070"), (0.55, "D0D0D0"), (0.85, "FFFFFF")])
    return height, albedo, roughness, 0.0, ao


def build_rock(nodes, links):
    co, w = torus_uv(nodes, links, 6.0, 0.6)
    strata = noise(nodes, links, co, w, 3.5, 7, 0.78)
    grain = noise(nodes, links, co, w, 16.0, 3, 0.5)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", strata, 0.7),
                       math_node(nodes, links, "MULTIPLY", grain, 0.3))
    albedo = ramp(nodes, links, height, [
        (0.2, "3A3C40"), (0.4, "4A4C50"), (0.58, "5A5C60"), (0.75, "6A6C70"), (0.9, "7A7C80")])
    roughness = ramp(nodes, links, height, [(0.2, "FAFAFA"), (0.7, "ECECEC")])
    ao = ramp(nodes, links, height, [(0.12, "606060"), (0.5, "D0D0D0"), (0.85, "FFFFFF")])
    return height, albedo, roughness, 0.0, ao


def build_water(nodes, links):
    co, w = torus_uv(nodes, links, *isotropic(5.0, 0.5))
    waves = noise(nodes, links, co, w, 3.2, 6, 0.62)
    chop = noise(nodes, links, co, w, 9.5, 4, 0.55)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", waves, 0.65),
                       math_node(nodes, links, "MULTIPLY", chop, 0.35))
    albedo = ramp(nodes, links, height, [(0.25, "0A2830"), (0.55, "123A42"), (0.8, "1A4A52")])
    roughness = ramp(nodes, links, height, [(0.2, "3A3A3A"), (0.7, "555555")])
    metallic = ramp(nodes, links, height, [(0.2, "202020"), (0.7, "303030")])
    return height, albedo, roughness, metallic, 1.0


# ---------------------------------------------------------------- unit and armour families (U5.3)
def build_metal(nodes, links):
    """Hammered plate and mail: fine dents, a warm-to-cool steel grain, worn edges."""
    co, w = torus_uv(nodes, links, *isotropic(8.0, 0.7))
    dents = noise(nodes, links, co, w, 5.5, 6, 0.7)
    grain = noise(nodes, links, co, w, 18.0, 3, 0.5)
    scratches = noise(nodes, links, co, w, 11.0, 5, 0.8)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", dents, 0.6),
                       math_node(nodes, links, "MULTIPLY", grain, 0.25))
    height = math_node(nodes, links, "SUBTRACT", height, math_node(nodes, links, "MULTIPLY", scratches, 0.15))
    albedo = ramp(nodes, links, height, [
        (0.2, "4A4E55"), (0.4, "6A6F78"), (0.58, "8A909A"), (0.75, "AAB0BA"), (0.9, "C8CDD4")])
    roughness = ramp(nodes, links, height, [(0.2, "6A6A6A"), (0.5, "4A4A4A"), (0.8, "303030")])
    ao = ramp(nodes, links, height, [(0.12, "707070"), (0.5, "D0D0D0"), (0.85, "FFFFFF")])
    metallic = 0.95
    return height, albedo, roughness, metallic, ao


def build_cloth(nodes, links):
    """Woven wool with a diagonal twill and soft fibre noise."""
    co, w = torus_uv(nodes, links, *isotropic(12.0, 0.8))
    twill = noise(nodes, links, co, w, 9.0, 4, 0.55)
    fibre = noise(nodes, links, co, w, 22.0, 3, 0.5)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", twill, 0.65),
                       math_node(nodes, links, "MULTIPLY", fibre, 0.35))
    albedo = ramp(nodes, links, height, [
        (0.2, "3A342C"), (0.4, "524A3E"), (0.58, "6A6050"), (0.75, "827662"), (0.9, "9A8C76")])
    roughness = ramp(nodes, links, height, [(0.2, "F5F5F5"), (0.7, "E2E2E2")])
    ao = ramp(nodes, links, height, [(0.15, "787878"), (0.55, "D4D4D4"), (0.85, "FFFFFF")])
    metallic = 0.0
    return height, albedo, roughness, metallic, ao


def build_hide(nodes, links):
    """Thick beast hide: pebbled grain, scars, slightly waxy sheen."""
    co, w = torus_uv(nodes, links, *isotropic(6.5, 0.6))
    pebble = noise(nodes, links, co, w, 7.5, 6, 0.72)
    scars = noise(nodes, links, co, w, 3.0, 5, 0.8)
    grain = noise(nodes, links, co, w, 15.0, 3, 0.5)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", pebble, 0.55),
                       math_node(nodes, links, "MULTIPLY", grain, 0.25))
    height = math_node(nodes, links, "SUBTRACT", height, math_node(nodes, links, "MULTIPLY", scars, 0.2))
    albedo = ramp(nodes, links, height, [
        (0.2, "1E2026"), (0.4, "2C2F36"), (0.58, "3A3E48"), (0.75, "4A4E58"), (0.9, "5A5E68")])
    roughness = ramp(nodes, links, height, [(0.2, "E8E8E8"), (0.7, "D0D0D0")])
    ao = ramp(nodes, links, height, [(0.12, "686868"), (0.5, "D0D0D0"), (0.85, "FFFFFF")])
    metallic = 0.0
    return height, albedo, roughness, metallic, ao


def build_bone(nodes, links):
    """Carved bone and ivory: porous grain, warm polish on the ridges."""
    co, w = torus_uv(nodes, links, *isotropic(7.0, 0.65))
    pores = noise(nodes, links, co, w, 10.0, 6, 0.7)
    grain = noise(nodes, links, co, w, 3.5, 5, 0.6)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", pores, 0.45),
                       math_node(nodes, links, "MULTIPLY", grain, 0.55))
    albedo = ramp(nodes, links, height, [
        (0.2, "8A8474"), (0.4, "A39C88"), (0.58, "BDB49C"), (0.75, "D2C8AE"), (0.9, "E6DCC2")])
    roughness = ramp(nodes, links, height, [(0.2, "B0B0B0"), (0.7, "D8D8D8")])
    ao = ramp(nodes, links, height, [(0.15, "787878"), (0.55, "D8D8D8"), (0.85, "FFFFFF")])
    metallic = 0.0
    return height, albedo, roughness, metallic, ao


def build_wood(nodes, links):
    """Tool-haft and shield plank: tight rings, a worn varnish on the high spots."""
    co, w = torus_uv(nodes, links, 9.5, 0.55)
    rings = noise(nodes, links, co, w, 4.5, 7, 0.78)
    grain = noise(nodes, links, co, w, 16.0, 3, 0.5)
    height = math_node(nodes, links, "ADD",
                       math_node(nodes, links, "MULTIPLY", rings, 0.7),
                       math_node(nodes, links, "MULTIPLY", grain, 0.3))
    albedo = ramp(nodes, links, height, [
        (0.2, "3A2A1C"), (0.4, "52402C"), (0.58, "6A563C"), (0.75, "826C4C"), (0.9, "9A825C")])
    roughness = ramp(nodes, links, height, [(0.2, "E0E0E0"), (0.7, "C8C8C8")])
    ao = ramp(nodes, links, height, [(0.12, "606060"), (0.5, "D0D0D0"), (0.85, "FFFFFF")])
    metallic = 0.0
    return height, albedo, roughness, metallic, ao


SETS = {
    "limestone": build_limestone,
    "soil": build_soil,
    "bark": build_bark,
    "leaf": build_leaf,
    "rock": build_rock,
    "water": build_water,
    "metal": build_metal,
    "cloth": build_cloth,
    "hide": build_hide,
    "bone": build_bone,
    "wood": build_wood,
}


def main():
    options = args()
    wanted = list(SETS) if options["sets"] == "all" else [s.strip() for s in options["sets"].split(",") if s.strip()]
    OUT.mkdir(parents=True, exist_ok=True)
    failed = []
    for name in wanted:
        if name not in SETS:
            failed.append(name + ": unknown set")
            continue
        try:
            save_set(name, SETS[name])
        except Exception as error:
            failed.append(name + ": " + repr(error))
    if failed:
        for line in failed:
            print("PBR_BAKE_FAIL", line)
        raise SystemExit(1)
    print("PBR_BAKE_PASS all", ",".join(wanted))


if __name__ == "__main__":
    main()