"""Aethel-9 surfaces, run on the built file:  blender -b Tender_Aethel.blend --python build/ae_surfaces.py
* every mesh gets a second UV layer 'UV_Tile' (box projection, metres) for the tiling sets; UV0 ('UVMap') stays: on the hull
  it is the surface parameter that the painted maps use, on the parts the lofts' arc lengths.
* tiling sets from the user's own Season 2 pack, read in place from models/vessel-interior/textures (nothing downloaded):
  ScratchedMetal under the paints, StainedSteel on bare metal, the CC0 leather scan (vessel-interior/source_textures) on seats.
* the hull: the painted set (ae_textures.py) plus an occlusion map baked in Cycles from the whole model onto its UV0 and
  written into the ORM's red channel (T_Aethel_Hull_ORM.png).
"""
import os
import sys
import bpy
import numpy as np

ROOT = os.path.dirname(bpy.data.filepath)
HERE = os.path.join(ROOT, 'build')
TEX = os.path.join(ROOT, 'textures')
VI = os.path.normpath(os.path.join(ROOT, '..', 'vessel-interior'))
PACK = os.path.join(VI, 'textures', 'ModSci_Engineer', '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')
LEATHER = os.path.join(VI, 'source_textures', 'leather_white')
argv = sys.argv[sys.argv.index('--') + 1:] if '--' in sys.argv else []
BAKE = 'nobake' not in argv

# name: (tint, metallic, roughness, tile metres, detail set, detail strength, extra)
# concept pass (2026-10-22): dark gunmetal paint worn to bare metal in scratches (extra['scratch']: how much of the set's
# dark scratch lines show as bright metal), black frames, dark metals, silver joints
LOOK = {
    'Paint_Hull': ((0.034, 0.045, 0.066), 0.35, 0.40, 0.9, 'SCR', 0.20, {'scratch': 0.45}),
    'Paint_Graphite': ((0.022, 0.026, 0.034), 0.45, 0.42, 0.9, 'SCR', 0.25, {'scratch': 0.6}),
    'HullDark': ((0.02, 0.022, 0.026), 0.4, 0.5, 1.0, 'SCR', 0.3, {}),
    'Metal_Dark': ((0.030, 0.033, 0.040), 0.85, 0.36, 0.7, 'STN', 0.45, {'scratch': 0.4}),
    'Metal_Light': ((0.55, 0.57, 0.60), 1.0, 0.24, 0.8, 'STN', 0.35, {}),
    'Frame': ((0.012, 0.013, 0.016), 0.3, 0.40, 0.7, 'SCR', 0.30, {'scratch': 0.35}),
    'Lining': ((0.10, 0.13, 0.17), 0.2, 0.55, 1.0, 'SCR', 0.35, {}),
    'Floor': ((0.05, 0.055, 0.06), 0.7, 0.5, 0.8, 'STN', 0.6, {}),
    'Console': ((0.06, 0.07, 0.08), 0.3, 0.45, 0.8, 'SCR', 0.4, {}),
    'Seat': ((0.03, 0.03, 0.035), 0.0, 0.6, 0.5, 'LEATHER', 0.6, {}),
    'Logo': ((0.80, 0.82, 0.86), 0.0, 0.45, 0.8, 'SCR', 0.25, {}),
}
WORN_METAL = (0.17, 0.18, 0.20)
# per tiling set: linear-luminance median (the detail multiplier is normalised by it) and the scratch band (the set's darkest
# few per cent: ScratchedMetal 5 % ~ 0.624, median 0.658; StainedSteel 5 % ~ 0.180, median 0.245)
SET_STATS = {'SCR': dict(med=0.658, s0=0.632, s1=0.612), 'STN': dict(med=0.245, s0=0.186, s1=0.165),
             'LEATHER': dict(med=0.40, s0=0.0, s1=-1.0)}


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


def box_uv(ob, scale=1.0):
    me = ob.data
    if 'UV_Tile' in me.uv_layers:
        me.uv_layers.remove(me.uv_layers['UV_Tile'])
    uv = me.uv_layers.new(name='UV_Tile')
    V = np.array([v.co for v in me.vertices]) @ np.array(ob.matrix_world)[:3, :3].T + np.array(ob.matrix_world)[:3, 3]
    out = np.zeros((len(me.loops), 2))
    for p in me.polygons:
        n = np.abs(np.array(p.normal))
        ax = int(np.argmax(n))
        for li in p.loop_indices:
            c = V[me.loops[li].vertex_index]
            out[li] = (c[[1, 2]] if ax == 0 else (c[[0, 2]] if ax == 1 else c[[0, 1]])) * scale
    uv.data.foreach_set('uv', out.ravel())
    me.uv_layers.active = me.uv_layers['UVMap']
    me.uv_layers['UVMap'].active_render = True


def nodes_detail(nt, set_name, tile, uv='UV_Tile'):
    """Image nodes of a tiling set on UV_Tile; returns (diff colour socket, rough socket, normal-map socket or None)."""
    n_uv = nt.nodes.new('ShaderNodeUVMap'); n_uv.uv_map = uv
    mp = nt.nodes.new('ShaderNodeMapping'); mp.inputs['Scale'].default_value = (1 / tile, 1 / tile, 1 / tile)
    nt.links.new(n_uv.outputs['UV'], mp.inputs['Vector'])
    if set_name == 'LEATHER':
        files = [os.path.join(LEATHER, f) for f in ('leather_white_diff_2k.jpg', 'leather_white_rough_2k.jpg',
                                                     'leather_white_nor_gl_2k.jpg')]
    else:
        st = SCR if set_name == 'SCR' else STN
        files = [os.path.join(PACK, f) for f in st] + ([None] if len(st) == 2 else [])
    out = []
    for f, col in zip(files, (True, False, False)):
        if f is None or not os.path.exists(f):
            out.append(None); continue
        t = nt.nodes.new('ShaderNodeTexImage'); t.image = load(f, col)
        nt.links.new(mp.outputs['Vector'], t.inputs['Vector'])
        out.append(t.outputs['Color'])
    return out


def build_material(mat, look, base_socket=None, rough_socket=None, normal_socket=None):
    tint, met, rough, tile, set_name, k, extra = look
    nt = mat.node_tree
    p = nt.nodes['Principled BSDF']
    diff, rtex, ntex = nodes_detail(nt, set_name, tile)
    # base: tint (or the painted map) x the set's luminance around 1 (strength k)
    st = SET_STATS.get(set_name, dict(med=0.45, s0=0.0, s1=-1.0))
    lum = nt.nodes.new('ShaderNodeRGBToBW'); nt.links.new(diff, lum.inputs['Color'])
    m1 = nt.nodes.new('ShaderNodeMath'); m1.operation = 'MULTIPLY_ADD'
    m1.inputs[1].default_value = k / st['med']; m1.inputs[2].default_value = 1 - k
    nt.links.new(lum.outputs['Val'], m1.inputs[0])
    mix = nt.nodes.new('ShaderNodeMix'); mix.data_type = 'RGBA'; mix.blend_type = 'MULTIPLY'
    mix.inputs['Factor'].default_value = 1.0
    if base_socket is not None:
        nt.links.new(base_socket, mix.inputs['A'])
    else:
        mix.inputs['A'].default_value = (*tint, 1)
    nt.links.new(m1.outputs['Value'], mix.inputs['B'])
    base_out = mix.outputs['Result']
    scr_k = extra.get('scratch', 0.0)
    if scr_k > 0:
        # the set's dark scratch lines -> a mask (1 in a scratch), mixed towards worn bare metal
        sm = nt.nodes.new('ShaderNodeMapRange'); sm.clamp = True
        sm.inputs['From Min'].default_value = st['s0']; sm.inputs['From Max'].default_value = st['s1']
        sm.inputs['To Min'].default_value = 0.0; sm.inputs['To Max'].default_value = scr_k
        nt.links.new(lum.outputs['Val'], sm.inputs['Value'])
        wm = nt.nodes.new('ShaderNodeMix'); wm.data_type = 'RGBA'; wm.blend_type = 'MIX'
        nt.links.new(sm.outputs['Result'], wm.inputs['Factor'])
        nt.links.new(base_out, wm.inputs['A']); wm.inputs['B'].default_value = (*WORN_METAL, 1)
        base_out = wm.outputs['Result']
        mm = nt.nodes.new('ShaderNodeMapRange'); mm.clamp = True
        mm.inputs['To Min'].default_value = met if base_socket is None else 0.0
        mm.inputs['To Max'].default_value = 0.95
        nt.links.new(sm.outputs['Result'], mm.inputs['Value'])
        nt.links.new(mm.outputs['Result'], p.inputs['Metallic'])
    else:
        p.inputs['Metallic'].default_value = met
    nt.links.new(base_out, p.inputs['Base Color'])
    # roughness: base (or painted) + the set's variation
    r1 = nt.nodes.new('ShaderNodeMath'); r1.operation = 'MULTIPLY_ADD'
    r1.inputs[1].default_value = 0.30; r1.inputs[2].default_value = -0.15
    nt.links.new(rtex if rtex is not None else lum.outputs['Val'], r1.inputs[0])
    r2 = nt.nodes.new('ShaderNodeMath'); r2.operation = 'ADD'; r2.use_clamp = True
    nt.links.new(r1.outputs['Value'], r2.inputs[0])
    if rough_socket is not None:
        nt.links.new(rough_socket, r2.inputs[1])
    else:
        r2.inputs[1].default_value = rough
    nt.links.new(r2.outputs['Value'], p.inputs['Roughness'])
    # normal: the set's normal (or a bump from its roughness) over the given normal
    if ntex is not None:
        nm = nt.nodes.new('ShaderNodeNormalMap'); nm.uv_map = 'UV_Tile'; nm.inputs['Strength'].default_value = 0.35
        nt.links.new(ntex, nm.inputs['Color'])
        if normal_socket is not None:
            bump = nt.nodes.new('ShaderNodeBump'); bump.inputs['Strength'].default_value = 0.0
            nt.links.new(normal_socket, bump.inputs['Normal'])
            # the painted normal is the finer detail on the hull: keep it, add the set as a light bump
            b2 = nt.nodes.new('ShaderNodeBump'); b2.inputs['Strength'].default_value = 0.04; b2.inputs['Distance'].default_value = 0.002
            nt.links.new(rtex if rtex is not None else lum.outputs['Val'], b2.inputs['Height'])
            nt.links.new(normal_socket, b2.inputs['Normal'])
            nt.links.new(b2.outputs['Normal'], p.inputs['Normal'])
        else:
            nt.links.new(nm.outputs['Normal'], p.inputs['Normal'])
    else:
        b2 = nt.nodes.new('ShaderNodeBump'); b2.inputs['Strength'].default_value = 0.05; b2.inputs['Distance'].default_value = 0.002
        nt.links.new(rtex if rtex is not None else lum.outputs['Val'], b2.inputs['Height'])
        if normal_socket is not None:
            nt.links.new(normal_socket, b2.inputs['Normal'])
        nt.links.new(b2.outputs['Normal'], p.inputs['Normal'])


def reset(mat):
    nt = mat.node_tree
    for n in list(nt.nodes):
        if n.type not in ('BSDF_PRINCIPLED', 'OUTPUT_MATERIAL'):
            nt.nodes.remove(n)
    return nt


def bake_hull_ao(hull, size=2048):
    sc = bpy.context.scene
    sc.render.engine = 'CYCLES'
    sc.cycles.samples = 128
    try:
        prefs = bpy.context.preferences.addons['cycles'].preferences
        for t in ('OPTIX', 'CUDA', 'HIP'):
            try:
                prefs.compute_device_type = t; prefs.get_devices()
                if any(d.type == t for d in prefs.devices):
                    for d in prefs.devices:
                        d.use = d.type == t
                    sc.cycles.device = 'GPU'
                    break
            except Exception:
                continue
    except Exception:
        pass
    img = bpy.data.images.new('T_Aethel_Hull_AO_bake', size, size, alpha=False, float_buffer=False)
    img.colorspace_settings.name = 'Non-Color'
    mat = hull.material_slots[0].material
    targets = []
    for slot in hull.material_slots:
        nt = slot.material.node_tree
        n = nt.nodes.new('ShaderNodeTexImage'); n.image = img; n.name = '_bake'
        nt.nodes.active = n; targets.append((nt, n))
    for o in bpy.context.view_layer.objects:
        o.select_set(False)
    hull.select_set(True); bpy.context.view_layer.objects.active = hull
    hull.data.uv_layers.active = hull.data.uv_layers['UVMap']
    sc.world.light_settings.distance = 0.6 if sc.world else None
    bpy.ops.object.bake(type='AO', margin=8, use_clear=True)
    a = np.array(img.pixels[:]).reshape(size, size, 4)[..., 0]
    for nt, n in targets:
        nt.nodes.remove(n)
    bpy.data.images.remove(img)
    return a


def main():
    sc = bpy.context.scene
    if sc.world is None:
        sc.world = bpy.data.worlds.new('World')
    for ob in bpy.data.objects:
        if ob.type == 'MESH':
            box_uv(ob)
    hull = bpy.data.objects['SM_Aethel_Hull']
    # occlusion bake onto the hull's UV0, written into the ORM's red channel
    orm_path = os.path.join(TEX, 'T_Aethel_Hull_ORM.png')
    if BAKE:
        ao = bake_hull_ao(hull)
        im = bpy.data.images.load(orm_path); im.colorspace_settings.name = 'Non-Color'
        w, h = im.size
        px = np.array(im.pixels[:]).reshape(h, w, 4)
        aor = ao
        if aor.shape != (h, w):
            ys = (np.arange(h) * aor.shape[0] / h).astype(int); xs = (np.arange(w) * aor.shape[1] / w).astype(int)
            aor = aor[ys][:, xs]
        px[..., 0] = np.clip(0.25 + 0.75 * aor, 0, 1)
        im.pixels[:] = px.ravel()
        im.filepath_raw = orm_path; im.file_format = 'PNG'; im.save()
        bpy.data.images.remove(im)
    # materials
    for name, look in LOOK.items():
        mat = bpy.data.materials.get(name)
        if mat is None:
            continue
        reset(mat)
        build_material(mat, look)
    # the hull: painted maps (UV0) under the paint detail; AO multiplies the base in Blender (Unreal reads ORM.R)
    mat = bpy.data.materials['Hull']
    nt = reset(mat)
    def tex(name, colour):
        t = nt.nodes.new('ShaderNodeTexImage'); t.image = load(os.path.join(TEX, name), colour)
        u = nt.nodes.new('ShaderNodeUVMap'); u.uv_map = 'UVMap'; nt.links.new(u.outputs['UV'], t.inputs['Vector'])
        t.image.reload()
        return t
    bc = tex('T_Aethel_Hull_BaseColor.png', True)
    orm = tex('T_Aethel_Hull_ORM.png', False)
    nrm = tex('T_Aethel_Hull_Normal.png', False)
    sep = nt.nodes.new('ShaderNodeSeparateColor'); nt.links.new(orm.outputs['Color'], sep.inputs['Color'])
    aom = nt.nodes.new('ShaderNodeMix'); aom.data_type = 'RGBA'; aom.blend_type = 'MULTIPLY'; aom.inputs['Factor'].default_value = 1.0
    nt.links.new(bc.outputs['Color'], aom.inputs['A'])
    nt.links.new(sep.outputs['Red'], aom.inputs['B'])
    nmap = nt.nodes.new('ShaderNodeNormalMap'); nmap.uv_map = 'UVMap'
    nt.links.new(nrm.outputs['Color'], nmap.inputs['Color'])
    build_material(mat, LOOK['Paint_Hull'], base_socket=aom.outputs['Result'], rough_socket=sep.outputs['Green'],
                   normal_socket=nmap.outputs['Normal'])
    # the painted chips are bare metal: metallic from the ORM's blue channel over the scratches' metallic
    pp = nt.nodes['Principled BSDF']
    mx = nt.nodes.new('ShaderNodeMath'); mx.operation = 'MAXIMUM'
    src = pp.inputs['Metallic'].links[0].from_socket if pp.inputs['Metallic'].is_linked else None
    if src is not None:
        nt.links.new(src, mx.inputs[0])
    else:
        mx.inputs[0].default_value = LOOK['Paint_Hull'][1]
    nt.links.new(sep.outputs['Blue'], mx.inputs[1])
    nt.links.new(mx.outputs['Value'], pp.inputs['Metallic'])
    # glass: thin (low IOR) and clear, with a faint blue glow (the concept's lit cabin seen through it)
    gp = bpy.data.materials['Glass'].node_tree.nodes['Principled BSDF']
    gp.inputs['Emission Color'].default_value = (0.10, 0.30, 0.75, 1); gp.inputs['Emission Strength'].default_value = 0.10
    g = bpy.data.materials['Glass']; p = g.node_tree.nodes['Principled BSDF']
    p.inputs['Transmission Weight'].default_value = 1.0; p.inputs['IOR'].default_value = 1.08
    p.inputs['Roughness'].default_value = 0.04; p.inputs['Alpha'].default_value = 1.0
    p.inputs['Base Color'].default_value = (0.55, 0.70, 0.85, 1)
    bpy.ops.wm.save_mainfile()
    print('SURFACES DONE')


main()