"""Surfaces for the machine parts kit (exec'd through the Blender MCP; long steps run on bpy.app.timers).

Like the docking arm (memory: props-realistic-with-good-textures), a mechanical asset gets in its first version:
  * two UV layers. UVMap (UV0): one unique unwrap of ALL kit meshes packed together into one 4096 atlas (even texel
    density, about 1.5 mm per texel). UV_Tile (UV1): box projection in metres, for the tiling sets.
  * maps baked from the model in Cycles into that atlas: textures/T_MPK_Masks.png (R ambient occlusion, G edges,
    B cavity) and textures/T_MPK_BevelNormal.png (every hard edge rounded by 2.5 mm, tangent space).
  * the tiling sets of the user's Season 2 pack, read in place from models/vessel-interior/textures/ModSciInteriors:
    ScratchedMetal under the paint, StainedSteel under the dark coat, Glass_DirtRough as a smudge mask.
The sheet's look: an even tan paint, a little lighter wear on exposed edges, a few small brown chips, dark ports.
"""
import math, os, time
import bpy, bmesh
import numpy as np

ROOT = r"D:\HADAL - The Descent\models\machine-parts-kit"
TEX = os.path.join(ROOT, "textures")
PACK = r"D:\HADAL - The Descent\models\vessel-interior\textures\ModSciInteriors\Textures\Base"
SCR = ("ScratchedMetal/T_ScratchedMetal_Diff.png", "ScratchedMetal/T_ScratchedMetal_Rough.png",
       "ScratchedMetal/T_ScratchedMetal_Normal.png")
STN = ("StainedSteel/T_StainSteeled_Diff.png", "StainedSteel/T_StainSteeled_Rough.png")
DIRT = "Glass/T_Glass_DirtRough.png"
SIZE = 4096
NS = bpy.app.driver_namespace


def parts():
    return [o for o in bpy.data.collections["MPK | Parts"].all_objects if o.type == 'MESH']


# ------------------------------------------------------------------ UVs
def uv_tile(ob, tile=1.0):
    """UV_Tile: each face projected on the plane its normal faces most (metres, 1 unit = 1 m)"""
    me = ob.data
    if "UV_Tile" not in me.uv_layers:
        me.uv_layers.new(name="UV_Tile")
    bm = bmesh.new(); bm.from_mesh(me)
    lay = bm.loops.layers.uv["UV_Tile"]
    for f in bm.faces:
        n = f.normal
        a = max(range(3), key=lambda i: abs(n[i]))
        i, j = {0: (1, 2), 1: (0, 2), 2: (0, 1)}[a]
        s = 1 if n[a] >= 0 else -1
        for l in f.loops:
            co = l.vert.co
            u = co[i] * (s if a != 1 else -s)
            l[lay].uv = (u / tile, co[j] / tile)
    bm.to_mesh(me); bm.free()


def uv_atlas_start():
    """unwrap every kit mesh into one shared UV space (multi-object edit mode packs them together)"""
    obs = parts()
    for o in obs:
        me = o.data
        if "UVMap" not in me.uv_layers:
            me.uv_layers.new(name="UVMap")
        me.uv_layers.active = me.uv_layers["UVMap"]
        me.uv_layers["UVMap"].active_render = True
    for o in bpy.context.view_layer.objects:
        o.select_set(False)
    for o in obs:
        o.select_set(True)
    bpy.context.view_layer.objects.active = obs[0]
    NS["mpk_surf"] = "uv: unwrapping"

    def _go():
        t = time.time()
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.uv.smart_project(angle_limit=math.radians(60), island_margin=0.0015, area_weight=0.0,
                                 correct_aspect=True, scale_to_bounds=False)
        bpy.ops.uv.pack_islands(rotate=True, margin_method='FRACTION', margin=0.0012)
        bpy.ops.object.mode_set(mode='OBJECT')
        for o in obs:
            uv_tile(o)
            # UVMap first: Unreal builds tangents from UV0
            names = [l.name for l in o.data.uv_layers]
            assert names[0] == "UVMap" or "UVMap" in names, names
        NS["mpk_surf"] = f"uv: done {time.time() - t:.0f}s"
        return None
    bpy.app.timers.register(_go, first_interval=0.1)


# ------------------------------------------------------------------ images
def image(name, colour=False, fill=(0, 0, 0, 1)):
    im = bpy.data.images.get(name)
    if im is not None:
        bpy.data.images.remove(im)
    im = bpy.data.images.new(name, SIZE, SIZE, alpha=False, float_buffer=False)
    im.colorspace_settings.name = 'sRGB' if colour else 'Non-Color'
    im.generated_color = fill
    return im


def pack_tex(rel, colour):
    p = os.path.join(PACK, rel)
    im = bpy.data.images.load(p, check_existing=True)
    im.colorspace_settings.name = 'sRGB' if colour else 'Non-Color'
    return im


# ------------------------------------------------------------------ bakes
BAKE_STEPS = ("AO", "EDGE", "CAVITY", "NORMAL")


def _bake_material_swap(kind):
    """temporarily route every kit material to the shader the bake needs; returns a restore function"""
    saved = []
    for mn in ("MPK_Paint", "MPK_Dark", "MPK_Steel"):
        m = bpy.data.materials[mn]
        nt = m.node_tree
        out = next(n for n in nt.nodes if n.type == 'OUTPUT_MATERIAL' and n.is_active_output)
        old = out.inputs["Surface"].links[0].from_socket if out.inputs["Surface"].links else None
        tmp = []
        if kind in ("EDGE", "CAVITY", "NORMAL"):
            bev = nt.nodes.new("ShaderNodeBevel"); tmp.append(bev)
            bev.samples = 16
            bev.inputs["Radius"].default_value = 0.0025 if kind == "NORMAL" else 0.004
            geo = nt.nodes.new("ShaderNodeNewGeometry"); tmp.append(geo)
            if kind == "NORMAL":
                sh = nt.nodes.new("ShaderNodeBsdfDiffuse"); tmp.append(sh)
                nt.links.new(bev.outputs[0], sh.inputs["Normal"])
            else:
                dp = nt.nodes.new("ShaderNodeVectorMath"); dp.operation = 'DOT_PRODUCT'; tmp.append(dp)
                nt.links.new(bev.outputs[0], dp.inputs[0]); nt.links.new(geo.outputs["Normal"], dp.inputs[1])
                mr = nt.nodes.new("ShaderNodeMapRange"); tmp.append(mr)
                nt.links.new(dp.outputs["Value"], mr.inputs[0])
                mr.inputs[1].default_value = 0.985 if kind == "EDGE" else 0.90   # from: dot near 1 = flat
                mr.inputs[2].default_value = 0.80 if kind == "EDGE" else 0.999
                mr.inputs[3].default_value = 0.0; mr.inputs[4].default_value = 1.0
                if kind == "CAVITY":
                    # cavity = inner corners only: the bevel normal leans away from the pointiness
                    pt = nt.nodes.new("ShaderNodeNewGeometry"); tmp.append(pt)
                    inv = nt.nodes.new("ShaderNodeMapRange"); tmp.append(inv)
                    nt.links.new(pt.outputs["Pointiness"], inv.inputs[0])
                    inv.inputs[1].default_value = 0.5; inv.inputs[2].default_value = 0.42
                    inv.inputs[3].default_value = 0.0; inv.inputs[4].default_value = 1.0
                    mr2 = nt.nodes.new("ShaderNodeMath"); mr2.operation = 'MULTIPLY'; tmp.append(mr2)
                    one = nt.nodes.new("ShaderNodeMath"); one.operation = 'SUBTRACT'; tmp.append(one)
                    one.inputs[0].default_value = 1.0
                    nt.links.new(mr.outputs[0], one.inputs[1])
                    nt.links.new(one.outputs[0], mr2.inputs[0]); nt.links.new(inv.outputs[0], mr2.inputs[1])
                    val = mr2.outputs[0]
                else:
                    val = mr.outputs[0]
                sh = nt.nodes.new("ShaderNodeEmission"); tmp.append(sh)
                nt.links.new(val, sh.inputs["Color"])
            nt.links.new(sh.outputs[0], out.inputs["Surface"])
        tgt = nt.nodes.get("BAKE_TARGET")
        if tgt is None:
            tgt = nt.nodes.new("ShaderNodeTexImage"); tgt.name = tgt.label = "BAKE_TARGET"
            tgt.location = (-1400, 800)
        nt.nodes.active = tgt
        saved.append((m, out, old, tmp, tgt))

    def restore():
        for m, out, old, tmp, tgt in saved:
            nt = m.node_tree
            for n in tmp:
                nt.nodes.remove(n)
            if old is not None:
                nt.links.new(old, out.inputs["Surface"])
    return saved, restore


def bake_all_start(steps=BAKE_STEPS, samples=64):
    sc = bpy.context.scene
    sc.render.engine = 'CYCLES'
    sc.cycles.device = 'GPU'
    sc.render.bake.margin = 12
    sc.render.bake.use_selected_to_active = False
    sc.render.bake.use_clear = True
    hidden = []
    for o in bpy.data.objects:
        if o not in parts() and not o.hide_render:
            o.hide_render = True; hidden.append(o.name)
    obs = parts()
    for o in bpy.context.view_layer.objects:
        o.select_set(o in obs)
    bpy.context.view_layer.objects.active = obs[0]
    q = list(steps)
    NS["mpk_surf"] = "bake: queued " + ",".join(q)

    def _step():
        if not q:
            for n in hidden:
                bpy.data.objects[n].hide_render = False
            if "FAILED" not in NS.get("mpk_surf", ""):
                NS["mpk_surf"] = "bake: done"
            return None
        kind = q.pop(0)
        t = time.time()
        im = image("T_MPK_" + {"AO": "AO", "EDGE": "Edge", "CAVITY": "Cavity", "NORMAL": "BevelNormal"}[kind],
                   fill=(0.5, 0.5, 1, 1) if kind == "NORMAL" else (1, 1, 1, 1))
        saved, restore = _bake_material_swap(kind)
        for s in saved:
            s[4].image = im
        try:
            if kind == "AO":
                sc.cycles.samples = samples
                sc.world.light_settings.distance = 0.06
                bpy.ops.object.bake(type='AO')
            elif kind == "NORMAL":
                sc.cycles.samples = 16
                bpy.ops.object.bake(type='NORMAL', normal_space='TANGENT')
            else:
                sc.cycles.samples = 16
                bpy.ops.object.bake(type='EMIT')
            im.filepath_raw = os.path.join(TEX, im.name + ".png"); im.file_format = 'PNG'
            im.save()
            NS["mpk_surf"] = f"bake: {kind} {time.time() - t:.0f}s; left {q}"
        except Exception as e:
            import traceback
            NS["mpk_surf"] = f"bake: {kind} FAILED {e} | {traceback.format_exc()[-600:]}"
            q.clear()
        finally:
            restore()
        return 0.2
    bpy.app.timers.register(_step, first_interval=0.2)


def combine_masks():
    """R = AO, G = edges, B = cavity -> T_MPK_Masks.png (the one mask file the game material reads)"""
    ch = []
    for n in ("T_MPK_AO", "T_MPK_Edge", "T_MPK_Cavity"):
        im = bpy.data.images.get(n) or bpy.data.images.load(os.path.join(TEX, n + ".png"))
        a = np.empty(SIZE * SIZE * 4, np.float32); im.pixels.foreach_get(a)
        ch.append(a.reshape(-1, 4)[:, 0])
    out = bpy.data.images.get("T_MPK_Masks")
    if out is None:
        out = bpy.data.images.new("T_MPK_Masks", SIZE, SIZE, alpha=False)
    out.colorspace_settings.name = 'Non-Color'          # before the pixels (memory: blender-mcp gotchas)
    px = np.ones((SIZE * SIZE, 4), np.float32)
    px[:, 0], px[:, 1], px[:, 2] = ch
    out.pixels.foreach_set(px.ravel())
    out.filepath_raw = os.path.join(TEX, "T_MPK_Masks.png"); out.file_format = 'PNG'
    out.save()
    return out


# ------------------------------------------------------------------ materials
def _n(nt, t, x, y, **kw):
    n = nt.nodes.new(t); n.location = (x, y)
    for k, v in kw.items():
        setattr(n, k, v)
    return n


def build_material(name, tint, wear_col, rough_lo, rough_hi, base_set, tile, wear_amt, chip_amt, metal=0.0):
    m = bpy.data.materials[name]
    m.use_nodes = True
    nt = m.node_tree
    keep_bake = nt.nodes.get("BAKE_TARGET")
    for n in list(nt.nodes):
        if n is not keep_bake:
            nt.nodes.remove(n)
    L = nt.links.new
    out = _n(nt, "ShaderNodeOutputMaterial", 900, 0)
    bs = _n(nt, "ShaderNodeBsdfPrincipled", 600, 0)
    L(bs.outputs[0], out.inputs[0])
    uvt = _n(nt, "ShaderNodeUVMap", -1600, 200, uv_map="UV_Tile")
    mp = _n(nt, "ShaderNodeMapping", -1400, 200)
    mp.inputs["Scale"].default_value = (1 / tile, 1 / tile, 1)
    L(uvt.outputs[0], mp.inputs[0])
    uv0 = _n(nt, "ShaderNodeUVMap", -1600, -300, uv_map="UVMap")
    dif = _n(nt, "ShaderNodeTexImage", -1150, 400, image=pack_tex(base_set[0], True))
    rgh = _n(nt, "ShaderNodeTexImage", -1150, 100, image=pack_tex(base_set[1], False))
    drt = _n(nt, "ShaderNodeTexImage", -1150, -100, image=pack_tex(DIRT, False))
    for t_ in (dif, rgh, drt):
        L(mp.outputs[0], t_.inputs[0])
    msk = _n(nt, "ShaderNodeTexImage", -1150, -350, image=bpy.data.images.get("T_MPK_Masks") or bpy.data.images.load(os.path.join(TEX, "T_MPK_Masks.png")))
    msk.image.colorspace_settings.name = 'Non-Color'
    bnm = _n(nt, "ShaderNodeTexImage", -1150, -650, image=bpy.data.images.get("T_MPK_BevelNormal") or bpy.data.images.load(os.path.join(TEX, "T_MPK_BevelNormal.png")))
    bnm.image.colorspace_settings.name = 'Non-Color'
    L(uv0.outputs[0], msk.inputs[0]); L(uv0.outputs[0], bnm.inputs[0])
    sep = _n(nt, "ShaderNodeSeparateColor", -850, -350)
    L(msk.outputs[0], sep.inputs[0])
    # albedo: tint modulated by the pack's diffuse inside its measured band
    bw = _n(nt, "ShaderNodeRGBToBW", -850, 400); L(dif.outputs[0], bw.inputs[0])
    var = _n(nt, "ShaderNodeMapRange", -650, 400)
    var.inputs[1].default_value, var.inputs[2].default_value = 0.60, 0.70     # the map's measured band (0.62-0.67)
    var.inputs[3].default_value, var.inputs[4].default_value = 0.86, 1.08
    L(bw.outputs[0], var.inputs[0])
    tin = _n(nt, "ShaderNodeMix", -450, 400, data_type='RGBA', blend_type='MULTIPLY')
    tin.inputs["Factor"].default_value = 1.0
    tin.inputs["A"].default_value = (*tint, 1)
    L(var.outputs[0], tin.inputs["B"])
    # occlusion and cavity darken
    ao = _n(nt, "ShaderNodeMapRange", -650, -250)
    ao.inputs[3].default_value, ao.inputs[4].default_value = 0.55, 1.0
    L(sep.outputs[0], ao.inputs[0])
    cav = _n(nt, "ShaderNodeMapRange", -650, -450)
    cav.inputs[3].default_value, cav.inputs[4].default_value = 1.0, 0.75
    L(sep.outputs[2], cav.inputs[0])
    occ = _n(nt, "ShaderNodeMath", -450, -300, operation='MULTIPLY'); L(ao.outputs[0], occ.inputs[0]); L(cav.outputs[0], occ.inputs[1])
    # wear on exposed edges, broken up by a seeded noise (object space, so each part differs)
    tco = _n(nt, "ShaderNodeTexCoord", -1600, -900)
    nz = _n(nt, "ShaderNodeTexNoise", -1150, -950)
    nz.inputs["Scale"].default_value = 60.0; nz.inputs["Detail"].default_value = 6.0; nz.inputs["Roughness"].default_value = 0.65
    L(tco.outputs["Object"], nz.inputs["Vector"])
    # exposed edges only: where two shells meet (a bolt on a plate) the bevel bake also reads "edge", but the AO there
    # is low, so the AO gates the wear
    gate = _n(nt, "ShaderNodeMapRange", -850, -650)
    gate.inputs[1].default_value, gate.inputs[2].default_value = 0.72, 0.92
    L(sep.outputs[0], gate.inputs[0])
    eg = _n(nt, "ShaderNodeMath", -850, -800, operation='MULTIPLY'); L(sep.outputs[1], eg.inputs[0]); L(gate.outputs[0], eg.inputs[1])
    edge_n = _n(nt, "ShaderNodeMath", -650, -800, operation='MULTIPLY'); L(eg.outputs[0], edge_n.inputs[0]); L(nz.outputs[0], edge_n.inputs[1])
    wear = _n(nt, "ShaderNodeMapRange", -450, -800)
    wear.inputs[1].default_value, wear.inputs[2].default_value = 0.55 - wear_amt, 0.62 - wear_amt
    L(edge_n.outputs[0], wear.inputs[0])
    c1 = _n(nt, "ShaderNodeMix", -250, 300, data_type='RGBA'); L(wear.outputs[0], c1.inputs["Factor"])
    L(tin.outputs["Result"], c1.inputs["A"]); c1.inputs["B"].default_value = (*wear_col, 1)
    # a few small brown chips (the sheet shows some): a second, coarser noise thresholded hard
    nz2 = _n(nt, "ShaderNodeTexNoise", -1150, -1200)
    nz2.inputs["Scale"].default_value = 140.0; nz2.inputs["Detail"].default_value = 3.0
    L(tco.outputs["Object"], nz2.inputs["Vector"])
    chip_m = _n(nt, "ShaderNodeMath", -650, -1100, operation='MULTIPLY'); L(eg.outputs[0], chip_m.inputs[0]); L(nz2.outputs[0], chip_m.inputs[1])
    chip = _n(nt, "ShaderNodeMapRange", -450, -1100)
    chip.inputs[1].default_value, chip.inputs[2].default_value = 0.72 - chip_amt, 0.74 - chip_amt
    L(chip_m.outputs[0], chip.inputs[0])
    c2 = _n(nt, "ShaderNodeMix", -100, 300, data_type='RGBA'); L(chip.outputs[0], c2.inputs["Factor"])
    L(c1.outputs["Result"], c2.inputs["A"]); c2.inputs["B"].default_value = (0.16, 0.085, 0.045, 1)
    shade = _n(nt, "ShaderNodeMix", 150, 300, data_type='RGBA', blend_type='MULTIPLY')
    shade.inputs["Factor"].default_value = 1.0
    L(c2.outputs["Result"], shade.inputs["A"]); L(occ.outputs[0], shade.inputs["B"])
    L(shade.outputs["Result"], bs.inputs["Base Color"])
    # roughness: the pack's map in a narrow band + a few hundredths of smudge, worn edges a little glossier
    rr = _n(nt, "ShaderNodeMapRange", -650, 100)
    rr.inputs[3].default_value, rr.inputs[4].default_value = rough_lo, rough_hi
    L(rgh.outputs[0], rr.inputs[0])
    sm = _n(nt, "ShaderNodeMath", -450, 100, operation='MULTIPLY_ADD')
    L(drt.outputs[0], sm.inputs[0]); sm.inputs[1].default_value = 0.04; L(rr.outputs[0], sm.inputs[2])
    rw = _n(nt, "ShaderNodeMix", 150, 50, data_type='FLOAT'); L(wear.outputs[0], rw.inputs["Factor"])
    L(sm.outputs[0], rw.inputs["A"]); rw.inputs["B"].default_value = max(rough_lo - 0.12, 0.25)
    L(rw.outputs["Result"], bs.inputs["Roughness"])
    bs.inputs["Metallic"].default_value = metal
    # normal: baked bevel normal, the pack's fine scratches as bump on top
    nm = _n(nt, "ShaderNodeNormalMap", -650, -650, uv_map="UVMap"); L(bnm.outputs[0], nm.inputs["Color"])
    bp = _n(nt, "ShaderNodeBump", 300, -400)
    bp.inputs["Strength"].default_value = 0.08; bp.inputs["Distance"].default_value = 0.0005
    L(bw.outputs[0], bp.inputs["Height"]); L(nm.outputs[0], bp.inputs["Normal"])
    L(bp.outputs[0], bs.inputs["Normal"])
    if keep_bake:
        keep_bake.location = (-1400, 800)
    return m


def build_materials():
    # tints are linear; the sheet's paint is about sRGB (180, 166, 146)
    build_material("MPK_Paint", (0.43, 0.39, 0.27), (0.60, 0.57, 0.48), 0.48, 0.66, SCR, 0.6, 0.08, 0.04)
    build_material("MPK_Steel", (0.39, 0.36, 0.25), (0.55, 0.53, 0.50), 0.42, 0.60, SCR, 0.4, 0.12, 0.04)
    build_material("MPK_Dark", (0.045, 0.042, 0.038), (0.12, 0.115, 0.11), 0.50, 0.70, STN, 0.5, -0.05, -0.2)