"""Surfaces for the hard-surface kit (memory: props-realistic-with-good-textures; the method of
models/docking-arm/build/surfaces.py, cut down to two materials).

  * M_Kit_Paint: the sheet's beige cast paint over the Season 2 pack's ScratchedMetal set (read in place from
    models/vessel-interior/textures, already in the Unreal project), on UV_Tile (metres); smudges from the pack's
    Glass_DirtRough; dark scuffs on exposed edges and shade in crevices from the baked masks.
  * M_Kit_Dark: sockets, drillings, slot floors: dark matte coating over the same roughness map.
  * textures/T_Kit_Masks.png, one 4096 atlas baked in Cycles over all 39 meshes on UVMap (UV0):
    R ambient occlusion, G edges, B crevices.
"""
import os
import time
import bpy
import bmesh
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
PACK = os.path.normpath(os.path.join(HERE, "..", "..", "vessel-interior", "textures", "ModSci_Engineer", "Textures"))
SCR = ("Base/ScratchedMetal/T_ScratchedMetal_Diff.png", "Base/ScratchedMetal/T_ScratchedMetal_Rough.png",
       "Base/ScratchedMetal/T_ScratchedMetal_Normal.png")
DIRT = "Base/Glass/T_Glass_DirtRough.png"
ATLAS = int(os.environ.get("KIT_ATLAS", "4096"))
EDGE_R = 0.006            # metres: the parts are small, the edge mask looks this far round a corner

# tint: linear colour of the paint (the sheet's beige); shade: the diffuse map darkens/lightens it between these;
# rough: roughness at the map's dark and bright ends; wear: how readily exposed edges show wear_col.
SURF = {
    "M_Kit_Paint": dict(tint=(0.40, 0.345, 0.26), shade=(0.90, 1.05), rough=(0.66, 0.44), normal=0.30, tile=0.45,
                        dirt=0.05, dirt_tile=0.9, grime=0.38, wear=0.42, wear_col=(0.075, 0.066, 0.058), wear_rough=0.55),
    "M_Kit_Dark":  dict(tint=(0.022, 0.021, 0.020), shade=(1.0, 1.0), rough=(0.80, 0.60), normal=0.15, tile=0.35,
                        dirt=0.05, dirt_tile=0.6, grime=0.10, wear=0.0),
}


def use_gpu():
    if os.environ.get("KIT_CPU"):
        bpy.context.scene.render.engine = "CYCLES"
        bpy.context.scene.cycles.device = "CPU"
        return "CPU"
    sc = bpy.context.scene
    sc.render.engine = 'CYCLES'
    try:
        prefs = bpy.context.preferences.addons['cycles'].preferences
        for kind in ('OPTIX', 'CUDA', 'HIP', 'ONEAPI'):
            try:
                prefs.compute_device_type = kind
            except Exception:
                continue
            (prefs.refresh_devices if hasattr(prefs, "refresh_devices") else prefs.get_devices)()
            if any(d.type == kind for d in prefs.devices):
                for d in prefs.devices:
                    d.use = d.type == kind
                sc.cycles.device = 'GPU'
                return kind
    except Exception as ex:
        print("GPU not set:", ex)
    sc.cycles.device = 'CPU'
    return 'CPU'


def parts():
    return [o for o in bpy.data.objects if o.type == 'MESH' and o.name.startswith("SM_Kit_")]


def bake_masks(folder):
    t0 = time.time()
    kind = use_gpu()
    sc = bpy.context.scene
    N = ATLAS
    im = bpy.data.images.get("T_Kit_Masks") or bpy.data.images.new("T_Kit_Masks", N, N, alpha=False)
    im.colorspace_settings.name = 'Non-Color'            # before the pixels (memory: blender-mcp-long-jobs-and-gotchas)
    im.pixels.foreach_set(np.tile(np.array((1.0, 0.0, 1.0, 1.0), np.float32), N * N))
    print("IMAGE", round(time.time() - t0, 1), flush=True)
    mm = bpy.data.materials.new("_bake_masks")
    mm.use_nodes = True
    nt = mm.node_tree
    nt.nodes.clear()
    ao = nt.nodes.new('ShaderNodeAmbientOcclusion'); ao.samples = 12; ao.only_local = True
    ao.inputs["Distance"].default_value = 0.08
    cv = nt.nodes.new('ShaderNodeAmbientOcclusion'); cv.samples = 8; cv.only_local = True
    cv.inputs["Distance"].default_value = 0.012
    bv = nt.nodes.new('ShaderNodeBevel'); bv.samples = 6; bv.inputs["Radius"].default_value = EDGE_R
    geo = nt.nodes.new('ShaderNodeNewGeometry')
    dot = nt.nodes.new('ShaderNodeVectorMath'); dot.operation = 'DOT_PRODUCT'
    nt.links.new(bv.outputs["Normal"], dot.inputs[0]); nt.links.new(geo.outputs["Normal"], dot.inputs[1])
    inv = nt.nodes.new('ShaderNodeMath'); inv.operation = 'SUBTRACT'; inv.inputs[0].default_value = 1.0
    nt.links.new(dot.outputs["Value"], inv.inputs[1])
    edg = nt.nodes.new('ShaderNodeMath'); edg.operation = 'MULTIPLY'; edg.inputs[1].default_value = 3.5; edg.use_clamp = True
    nt.links.new(inv.outputs[0], edg.inputs[0])
    cc = nt.nodes.new('ShaderNodeCombineColor')
    nt.links.new(ao.outputs["AO"], cc.inputs[0]); nt.links.new(edg.outputs[0], cc.inputs[1]); nt.links.new(cv.outputs["AO"], cc.inputs[2])
    em = nt.nodes.new('ShaderNodeEmission')
    nt.links.new(cc.outputs[0], em.inputs["Color"])
    om = nt.nodes.new('ShaderNodeOutputMaterial')
    nt.links.new(em.outputs[0], om.inputs["Surface"])
    ti = nt.nodes.new('ShaderNodeTexImage'); ti.image = im
    nt.nodes.active = ti
    sc.cycles.samples = int(os.environ.get("KIT_SAMPLES", "24"))
    sc.cycles.use_denoising = False
    sc.render.bake.margin = 4
    sc.render.bake.margin_type = 'EXTEND'      # 'ADJACENT_FACES' (default) took minutes over thousands of islands
    sc.render.bake.use_clear = False

    def run():
        # Each part bakes into a small image the size of its own square of the atlas, then lands in the atlas:
        # Cycles walks the whole target image per object, so baking straight into the 4K atlas took over an hour.
        arr = np.empty(N * N * 4, np.float32)
        im.pixels.foreach_get(arr)
        arr = arr.reshape(N, N, 4)
        for p in parts():
            x, y, s = p["atlas_cell_px"]
            cell = bpy.data.images.new("_cell", s, s, alpha=False)
            cell.colorspace_settings.name = 'Non-Color'
            ti.image = cell
            me = p.data.copy()
            me.materials.clear()
            me.materials.append(mm)
            for poly in me.polygons:
                poly.material_index = 0
            uv = me.uv_layers["UVMap"]
            for d in uv.data:
                d.uv = ((d.uv.x * N - x) / s, (d.uv.y * N - y) / s)
            me.uv_layers.active = uv
            uv.active_render = True
            tmp = bpy.data.objects.new("_bake_" + p.name, me)
            sc.collection.objects.link(tmp)
            tmp.matrix_world = p.matrix_world
            for o2 in bpy.context.view_layer.objects:
                o2.select_set(False)
            tmp.select_set(True)
            bpy.context.view_layer.objects.active = tmp
            try:
                bpy.ops.object.bake(type='EMIT', margin=4, margin_type='EXTEND', use_clear=True)
                px = np.empty(s * s * 4, np.float32)
                cell.pixels.foreach_get(px)
                arr[y:y + s, x:x + s] = px.reshape(s, s, 4)
                print("BAKED", p.name, s, round(time.time() - t0, 1), flush=True)
            finally:
                bpy.data.objects.remove(tmp)
                bpy.data.meshes.remove(me)
                bpy.data.images.remove(cell)
        im.pixels.foreach_set(arr.ravel())
        ti.image = im
    try:
        run()
    except Exception as ex:
        if kind == 'CPU':
            raise
        print("bake on %s failed (%s), again on the CPU" % (kind, ex))
        kind = 'CPU'
        sc.cycles.device = 'CPU'
        run()
    os.makedirs(folder, exist_ok=True)
    im.filepath_raw = os.path.join(folder, "T_Kit_Masks.png")
    im.file_format = 'PNG'
    im.save()
    bpy.data.materials.remove(mm)
    return im, dict(device=kind, seconds=round(time.time() - t0, 1), size=N)


# ------------------------------------------------------------------ node helpers (as in docking-arm/build/surfaces.py)
def tex_image(path, colorspace="sRGB"):
    path = os.path.normpath(path)
    for im in bpy.data.images:
        if im.filepath and os.path.normpath(bpy.path.abspath(im.filepath)) == path:
            return im
    im = bpy.data.images.load(path)
    im.colorspace_settings.name = colorspace
    return im


_RANGE = {}


def tex_range(im, lo=3, hi=97):
    """the brightness the map really spans (3rd to 97th percentile), so a range set below means what it says"""
    if im.name not in _RANGE:
        n, ch = im.size[0] * im.size[1], im.channels
        px = np.empty(n * ch, np.float32)
        im.pixels.foreach_get(px)
        px = px.reshape(n, ch)
        if ch >= 3:
            c = px[:, :3]
            if im.colorspace_settings.name == 'sRGB':
                c = np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
            v = c @ np.array([0.2126, 0.7152, 0.0722], np.float32)
        else:
            v = px[:, 0]
        _RANGE[im.name] = (round(float(np.percentile(v, lo)), 4), round(float(np.percentile(v, hi)), 4))
    return _RANGE[im.name]


class Tree:
    def __init__(self, mat):
        mat.use_nodes = True
        self.nt = mat.node_tree
        self.nt.nodes.clear()

    def node(self, kind, label=None, **props):
        n = self.nt.nodes.new(kind)
        if label:
            n.label = label
        for k, v in props.items():
            setattr(n, k, v)
        i = len(self.nt.nodes) - 1
        n.location = (-2000 + 240 * (i // 4), 500 - 320 * (i % 4))
        return n

    def link(self, a, b):
        self.nt.links.new(a, b)

    def put(self, sock, v):
        if isinstance(v, bpy.types.NodeSocket):
            self.nt.links.new(v, sock)
        else:
            sock.default_value = v

    def math(self, op, a, b=None, clamp=False, label=None):
        n = self.node('ShaderNodeMath', label, operation=op, use_clamp=clamp)
        self.put(n.inputs[0], a)
        if b is not None:
            self.put(n.inputs[1], b)
        return n.outputs[0]

    def ranged(self, v, lo, hi, label=None, f0=0.0, f1=1.0, smooth=False):
        n = self.node('ShaderNodeMapRange', label)
        if smooth:
            n.interpolation_type = 'SMOOTHSTEP'
        self.put(n.inputs["Value"], v)
        n.inputs["From Min"].default_value = f0
        n.inputs["From Max"].default_value = f1
        n.inputs["To Min"].default_value = lo
        n.inputs["To Max"].default_value = hi
        return n.outputs["Result"]

    def mixc(self, fac, a, b, label=None):
        n = self.node('ShaderNodeMix', label, data_type='RGBA')
        self.put(n.inputs[0], fac)
        self.put(n.inputs[6], a if isinstance(a, bpy.types.NodeSocket) else tuple(a) + (1.0,))
        self.put(n.inputs[7], b if isinstance(b, bpy.types.NodeSocket) else tuple(b) + (1.0,))
        return n.outputs[2]

    def mixf(self, fac, a, b, label=None):
        n = self.node('ShaderNodeMix', label, data_type='FLOAT')
        self.put(n.inputs[0], fac)
        self.put(n.inputs[2], a)
        self.put(n.inputs[3], b)
        return n.outputs[0]

    def uv(self, name):
        return self.node('ShaderNodeUVMap', uv_map=name).outputs["UV"]

    def mapped(self, src, scale, label=None):
        n = self.node('ShaderNodeMapping', label)
        self.nt.links.new(src, n.inputs["Vector"])
        n.inputs["Scale"].default_value = (scale, scale, scale)
        return n.outputs["Vector"]

    def tex(self, im, vec, label=None):
        n = self.node('ShaderNodeTexImage', label, image=im)
        self.nt.links.new(vec, n.inputs["Vector"])
        return n


def dress(mat, r, masks):
    t = Tree(mat)
    b = t.node('ShaderNodeBsdfPrincipled')
    out = t.node('ShaderNodeOutputMaterial')
    t.link(b.outputs[0], out.inputs["Surface"])
    u0 = t.uv("UV_Tile")
    mp = t.mapped(u0, 1.0 / r["tile"], "one repeat = %.2f m" % r["tile"])
    dirt = t.tex(tex_image(os.path.join(PACK, DIRT), "Non-Color"), t.mapped(u0, 1.0 / r["dirt_tile"]), "smudges").outputs["Color"]
    sep = t.node('ShaderNodeSeparateColor')
    t.link(t.tex(masks, t.uv("UVMap"), "baked: R occlusion, G edges, B crevices").outputs["Color"], sep.inputs[0])
    ao, edge, crev = sep.outputs[0], sep.outputs[1], sep.outputs[2]
    im = tex_image(os.path.join(PACK, SCR[0]))
    lo, hi = tex_range(im)
    bw = t.node('ShaderNodeRGBToBW')
    t.link(t.tex(im, mp, "surface, diffuse").outputs["Color"], bw.inputs[0])
    vm = t.node('ShaderNodeVectorMath', "tint x shade", operation='SCALE')
    vm.inputs[0].default_value = tuple(r["tint"])
    t.link(t.ranged(bw.outputs[0], r["shade"][0], r["shade"][1], "shade the tint", lo, hi), vm.inputs["Scale"])
    col = vm.outputs["Vector"]
    grime = t.math('MULTIPLY', t.math('POWER', t.math('SUBTRACT', 1.0, crev), 1.4), r["grime"], clamp=True,
                   label="shade in crevices")
    col = t.mixc(grime, col, (0.03, 0.026, 0.02), "darken crevices")
    rim = tex_image(os.path.join(PACK, SCR[1]), "Non-Color")
    rlo, rhi = tex_range(rim)
    ro = t.tex(rim, mp, "surface, roughness").outputs["Color"]
    rr = t.ranged(ro, r["rough"][0], r["rough"][1], "roughness range", rlo, rhi)
    rr = t.math('ADD', rr, t.math('MULTIPLY', dirt, r["dirt"]), clamp=True, label="smudges dull the sheen")
    rr = t.math('ADD', rr, t.math('MULTIPLY', grime, 0.3), clamp=True, label="crevices are matte")
    if r["wear"] > 0:                                   # dark scuffs on exposed edges, broken up by the scratch map
        breakup = t.ranged(ro, 1.25, 0.55, "scratches break the wear up", rlo, rhi)
        exposed = t.ranged(ao, 0.0, 1.0, "only where the part is exposed", 0.70, 0.93, smooth=True)
        w = t.math('MULTIPLY', t.math('MULTIPLY', t.math('MULTIPLY', edge, breakup), exposed), r["wear"])
        w = t.ranged(w, 0.0, 1.0, "scuffed edge", 0.36, 0.70, smooth=True)
        col = t.mixc(w, col, tuple(r["wear_col"]), "scuffs")
        rr = t.mixf(w, rr, r["wear_rough"])
    t.put(b.inputs["Base Color"], col)
    t.put(b.inputs["Roughness"], rr)
    t.put(b.inputs["Metallic"], 0.0)
    nd = t.node('ShaderNodeNormalMap', "fine relief", uv_map="UV_Tile")
    nd.inputs["Strength"].default_value = r["normal"]
    t.link(t.tex(tex_image(os.path.join(PACK, SCR[2]), "Non-Color"), mp, "surface, normal").outputs["Color"], nd.inputs["Color"])
    t.link(nd.outputs["Normal"], b.inputs["Normal"])
    mat.diffuse_color = tuple(r["tint"]) + (1.0,)


def run(out):
    masks, info = bake_masks(os.path.join(out, "textures"))
    for name, r in SURF.items():
        dress(bpy.data.materials[name], r, masks)
    try:
        bpy.ops.file.make_paths_relative()
    except Exception as ex:
        print("paths:", ex)
    print("SURF", info)
    return info