# vr_look.py - surfaces and studio for the veiled robot (Cycles).
# Armour: Poly Haven blue_metal_plate (CC0, already in models/pressure-door/source_polyhaven),
# grey-scaled and tinted to the references' dark olive paint; worn to bare steel on edges
# (Bevel-node edge mask) and darkened in crevices (AO node). Both masks are Cycles-only:
# this look is NOT baked yet.
import bpy, math
from mathutils import Vector as V

TEX = r"D:\HADAL - The Descent\models\pressure-door\source_polyhaven\textures\blue_metal_plate"


def img(name, cs):
    path = f"{TEX}\\blue_metal_plate_{name}_2k.jpg"
    im = bpy.data.images.load(path, check_existing=True)
    im.colorspace_settings.name = cs
    return im


def nodes_clear(m):
    m.use_nodes = True
    nt = m.node_tree
    for n in list(nt.nodes):
        nt.nodes.remove(n)
    return nt


def armor(m, tint=(0.052, 0.060, 0.046), scale=2.2):
    nt = nodes_clear(m); N = nt.nodes; L = nt.links.new
    out = N.new("ShaderNodeOutputMaterial"); bs = N.new("ShaderNodeBsdfPrincipled")
    tc = N.new("ShaderNodeTexCoord"); mp = N.new("ShaderNodeMapping")
    mp.inputs["Scale"].default_value = (scale, scale, scale)
    L(tc.outputs["Object"], mp.inputs["Vector"])

    def tex(name, cs):
        t = N.new("ShaderNodeTexImage"); t.image = img(name, cs)
        t.projection = 'BOX'; t.projection_blend = 0.25
        L(mp.outputs["Vector"], t.inputs["Vector"]); return t
    dif = tex("diff", "sRGB"); rou = tex("rough", "Non-Color"); nor = tex("nor_gl", "Non-Color")
    # paint colour: scan brightness only, around the tint
    bw = N.new("ShaderNodeRGBToBW"); L(dif.outputs["Color"], bw.inputs["Color"])
    paint = N.new("ShaderNodeMix"); paint.data_type = 'RGBA'
    paint.inputs["A"].default_value = (*[c * 0.55 for c in tint], 1)
    paint.inputs["B"].default_value = (*[c * 1.9 for c in tint], 1)
    L(bw.outputs["Val"], paint.inputs["Factor"])
    # edge mask from the Bevel node
    bev = N.new("ShaderNodeBevel"); bev.inputs["Radius"].default_value = 0.004; bev.samples = 6
    geo = N.new("ShaderNodeNewGeometry")
    dot = N.new("ShaderNodeVectorMath"); dot.operation = 'DOT_PRODUCT'
    L(bev.outputs["Normal"], dot.inputs[0]); L(geo.outputs["Normal"], dot.inputs[1])
    edge = N.new("ShaderNodeMapRange")
    edge.inputs["From Min"].default_value = 0.985; edge.inputs["From Max"].default_value = 0.80
    L(dot.outputs["Value"], edge.inputs["Value"])
    noi = N.new("ShaderNodeTexNoise"); noi.inputs["Scale"].default_value = 38; noi.inputs["Detail"].default_value = 8
    L(tc.outputs["Object"], noi.inputs["Vector"])
    brk = N.new("ShaderNodeMath"); brk.operation = 'MULTIPLY'
    L(edge.outputs["Result"], brk.inputs[0]); L(noi.outputs["Fac"], brk.inputs[1])
    wear = N.new("ShaderNodeMapRange")
    wear.inputs["From Min"].default_value = 0.22; wear.inputs["From Max"].default_value = 0.42
    L(brk.outputs["Value"], wear.inputs["Value"])
    # crevice dirt
    ao = N.new("ShaderNodeAmbientOcclusion"); ao.inputs["Distance"].default_value = 0.03; ao.samples = 8
    dirt = N.new("ShaderNodeMix"); dirt.data_type = 'RGBA'
    dirt.inputs["B"].default_value = (0.012, 0.012, 0.010, 1)
    inv = N.new("ShaderNodeMath"); inv.operation = 'SUBTRACT'; inv.inputs[0].default_value = 1.0
    L(ao.outputs["AO"], inv.inputs[1])
    dm = N.new("ShaderNodeMath"); dm.operation = 'MULTIPLY'; dm.inputs[1].default_value = 0.7
    L(inv.outputs["Value"], dm.inputs[0]); L(dm.outputs["Value"], dirt.inputs["Factor"])
    L(paint.outputs["Result"], dirt.inputs["A"])
    col = N.new("ShaderNodeMix"); col.data_type = 'RGBA'
    col.inputs["B"].default_value = (0.32, 0.32, 0.31, 1)            # bare steel
    L(dirt.outputs["Result"], col.inputs["A"]); L(wear.outputs["Result"], col.inputs["Factor"])
    L(col.outputs["Result"], bs.inputs["Base Color"])
    met = N.new("ShaderNodeMapRange"); met.inputs["To Min"].default_value = 0.25
    L(wear.outputs["Result"], met.inputs["Value"]); L(met.outputs["Result"], bs.inputs["Metallic"])
    rr = N.new("ShaderNodeMapRange"); rr.inputs["To Min"].default_value = 0.36; rr.inputs["To Max"].default_value = 0.62
    L(rou.outputs["Color"], rr.inputs["Value"]); L(rr.outputs["Result"], bs.inputs["Roughness"])
    nm = N.new("ShaderNodeNormalMap"); nm.inputs["Strength"].default_value = 0.35
    L(nor.outputs["Color"], nm.inputs["Color"]); L(nm.outputs["Normal"], bs.inputs["Normal"])
    bs.inputs["Coat Weight"].default_value = 0.15; bs.inputs["Coat Roughness"].default_value = 0.25
    L(bs.outputs["BSDF"], out.inputs["Surface"])
    m.diffuse_color = (*tint, 1)


def simple(m, rgb, metal, rough, rough_var=0.12, coat=0.0, emit=None):
    nt = nodes_clear(m); N = nt.nodes; L = nt.links.new
    out = N.new("ShaderNodeOutputMaterial"); bs = N.new("ShaderNodeBsdfPrincipled")
    bs.inputs["Base Color"].default_value = (*rgb, 1)
    bs.inputs["Metallic"].default_value = metal
    tc = N.new("ShaderNodeTexCoord")
    noi = N.new("ShaderNodeTexNoise"); noi.inputs["Scale"].default_value = 60; noi.inputs["Detail"].default_value = 6
    L(tc.outputs["Object"], noi.inputs["Vector"])
    mr = N.new("ShaderNodeMapRange")
    mr.inputs["To Min"].default_value = rough - rough_var / 2; mr.inputs["To Max"].default_value = rough + rough_var / 2
    L(noi.outputs["Fac"], mr.inputs["Value"]); L(mr.outputs["Result"], bs.inputs["Roughness"])
    bs.inputs["Coat Weight"].default_value = coat
    if emit:
        bs.inputs["Emission Color"].default_value = (*emit, 1)
        bs.inputs["Emission Strength"].default_value = 12.0
    L(bs.outputs["BSDF"], out.inputs["Surface"])
    m.diffuse_color = (*rgb, 1)


def studio(dark=False):
    sc = bpy.context.scene
    st = bpy.data.collections.get("VR | Studio")
    sc.render.engine = 'CYCLES'
    sc.cycles.samples = 160; sc.cycles.use_denoising = True
    try:
        sc.cycles.device = 'GPU'
    except Exception:
        pass
    sc.view_settings.view_transform = 'AgX'; sc.view_settings.look = 'AgX - Medium High Contrast'
    w = sc.world or bpy.data.worlds.new("World"); sc.world = w; w.use_nodes = True
    bg = w.node_tree.nodes.get("Background")
    g = 0.02 if dark else 0.05
    bg.inputs["Color"].default_value = (g, g, g, 1); bg.inputs["Strength"].default_value = 1.0
    floor = bpy.data.objects.get("VR_Studio_Floor")
    if floor is None:
        me = bpy.data.meshes.new("VR_Studio_Floor")
        import bmesh
        bm = bmesh.new()
        # a cyclorama: floor curving up into a back wall
        prof = [(y, 0.0) for y in (-8, -2, 0, 5.0)] + [(5.0 + 1.5 * math.sin(a), 1.5 - 1.5 * math.cos(a))
                                                       for a in [i * math.pi / 2 / 8 for i in range(1, 9)]] + [(6.5, 8)]
        rows = [[bm.verts.new((x, y, z)) for x in (-12, 12)] for y, z in prof]
        for a, b in zip(rows, rows[1:]):
            bm.faces.new((a[0], a[1], b[1], b[0]))
        bm.to_mesh(me); bm.free()
        for p in me.polygons:
            p.use_smooth = True
        floor = bpy.data.objects.new("VR_Studio_Floor", me); st.objects.link(floor)
        fm = bpy.data.materials.new("M_VR_StudioFloor"); floor.data.materials.append(fm)
    fm = floor.data.materials[0]; fm.use_nodes = True
    p = fm.node_tree.nodes.get("Principled BSDF")
    v = 0.012 if dark else 0.16
    p.inputs["Base Color"].default_value = (v, v, v, 1); p.inputs["Roughness"].default_value = 0.35 if dark else 0.8
    for o in [o for o in st.objects if o.type == 'LIGHT']:
        bpy.data.objects.remove(o)
    def area(name, loc, energy, size, rgb=(1, 1, 1)):
        ld = bpy.data.lights.new(name, 'AREA'); ld.energy = energy; ld.size = size; ld.color = rgb
        o = bpy.data.objects.new(name, ld); st.objects.link(o); o.location = loc
        d = V((0, 0, 1.1)) - V(loc); o.rotation_euler = d.to_track_quat('-Z', 'Y').to_euler()
    area("VR_Key", (-2.2, -3.0, 3.2), 320, 2.0)
    area("VR_Fill", (3.0, -2.4, 1.6), 90, 3.0)
    area("VR_RimL", (-2.4, 2.6, 2.6), 380, 1.2)
    area("VR_RimR", (2.4, 2.6, 2.4), 380, 1.2)
    area("VR_Top", (0, 0.4, 4.0), 110, 1.5)


def apply_all(dark=False):
    armor(bpy.data.materials["M_VR_Armor"])
    simple(bpy.data.materials["M_VR_Mech"], (0.010, 0.010, 0.011), 0.6, 0.32, coat=0.2)
    simple(bpy.data.materials["M_VR_Red"], (0.40, 0.010, 0.008), 0.0, 0.26, coat=0.6)
    simple(bpy.data.materials["M_VR_Eye"], (0.2, 0.02, 0.0), 0.0, 0.3, emit=(1.0, 0.16, 0.02))
    studio(dark)
    return "ok"