"""
World / look track, part 2 (module of the import runner): M_Terrain and its instance MI_Terrain (DESIGN_UE.md section 7.8).

Opaque lit, world-space normal. Inputs: vertex colour (rgb = sRGB-encoded rock colour, a = fresh-cut amount), UV0 (ore type
1..3, ore strength), UV1.x (dense basalt). One Custom node ports the v2 terrain shader (04_terrain.js lines 11-39) with value noise
written as macros (HLSL has no nested functions): strata brightness, fresh cuts cooler, faceted ore with a glint, a noise-
perturbed normal, roughness 0.9 / ore 0.18 / basalt 0.25, and the ambient term of section 7.3 as emissive.

Fix list after M3 (Docs/status/fixes_m3.md, requirement V1): surface detail so that rock in the headlamp reads as rock (mottling,
bedding layers, cracks, a stronger and finer normal, roughness 0.58-0.82 instead of 0.9), and the dense basalt's columnar jointing.
"""
import importlib.util
import os

import unreal

CODE = """
#define H3(q) frac(sin(dot((q), float3(12.9898, 78.233, 37.719))) * 43758.5453)
#define SM3(f) ((f) * (f) * (3.0 - 2.0 * (f)))
#define VN(p) lerp(lerp(lerp(H3(floor(p)), H3(floor(p) + float3(1,0,0)), SM3(frac(p)).x), lerp(H3(floor(p) + float3(0,1,0)), H3(floor(p) + float3(1,1,0)), SM3(frac(p)).x), SM3(frac(p)).y), lerp(lerp(H3(floor(p) + float3(0,0,1)), H3(floor(p) + float3(1,0,1)), SM3(frac(p)).x), lerp(H3(floor(p) + float3(0,1,1)), H3(floor(p) + float3(1,1,1)), SM3(frac(p)).x), SM3(frac(p)).y), SM3(frac(p)).z)
float3 Pm = WP * 0.01;
float3 vc = max(VC.rgb, 0.0);   // the engine already decodes the 8-bit sRGB vertex colour to linear
float cutAmt = Cut;
float3 s1 = float3(Pm.z * 0.42, Pm.x * 0.015, Pm.y * 0.015);
float3 s2 = Pm * 0.31;
float strata = VN(s1) * 0.55 + VN(s2) * 0.45;
float3 col = vc * (0.72 + 0.5 * strata);
col = lerp(col, col * float3(0.75, 0.82, 0.9), cutAmt * 0.5);
// Surface detail (fix list after M3, requirement V1). A lamp at the eye lights a wall head-on: Lambert is flat there and nothing casts
// a visible shadow, so rock can only read through what is ON the surface: mottling from 60 cm down to 8 cm; bedding (layers about
// 0.9 m thick, each with its own tone, a dark seam between them); a fracture network (warped cells about 0.9 m across, each block
// with its own tone and its own facet direction, dark cracks between them); and below a finer normal with a wet-rock sheen.
float3 d1 = Pm * 1.7;
float f1 = VN(d1);
float3 d2 = Pm * 5.3;
float f2 = VN(d2);
float3 d3 = Pm * 13.0;
float f3 = VN(d3);
float mott = f1 * 0.5 + f2 * 0.3 + f3 * 0.2;
float3 w1 = Pm * 0.23;
float bz = Pm.z * 1.1 + VN(w1) * 3.0;
float blayer = floor(bz);
float bfr = frac(bz);
float btone = H3(float3(blayer, 3.0, 7.0));
float bed = smoothstep(0.07, 0.0, min(bfr, 1.0 - bfr));
float3 c1 = Pm * 0.8 + float3(13.1, 7.7, 3.3);
float3 c2 = Pm * 0.9 + float3(5.2, 1.3, 9.1);
float3 vp = Pm * 1.1 + (float3(f1, VN(c1), VN(c2)) - 0.5) * 0.9;
float3 vi = floor(vp);
float3 vf = frac(vp);
float vF1 = 8.0;
float vF2 = 8.0;
float3 vid = float3(0.0, 0.0, 0.0);
for (int vz = -1; vz <= 1; vz++)
{
	for (int vy = -1; vy <= 1; vy++)
	{
		for (int vx = -1; vx <= 1; vx++)
		{
			float3 vg = float3(vx, vy, vz);
			float3 vc3 = vi + vg;
			float3 vo = float3(H3(vc3), H3(vc3 + 17.3), H3(vc3 + 41.7));
			float3 vr = vg + vo - vf;
			float vd = dot(vr, vr);
			if (vd < vF1) { vF2 = vF1; vF1 = vd; vid = vc3; }
			else if (vd < vF2) { vF2 = vd; }
		}
	}
}
float vcrk = VN(c1);
float crack = smoothstep(0.03 + 0.10 * vcrk, 0.0, sqrt(vF2) - sqrt(vF1)) * (0.25 + 0.75 * smoothstep(0.3, 0.6, VN(c2)));   // cracks of uneven width, strong in places and faint in others
float3 blockRnd = float3(H3(vid + 3.1), H3(vid + 9.2), H3(vid + 27.9));
col *= (0.3 + 1.4 * mott) * (0.62 + 0.76 * btone) * (0.74 + 0.52 * blockRnd.x) * (1.0 - 0.35 * bed) * (1.0 - 0.8 * crack);
float oreT = UV0.x;
float oreS = UV0.y;
float isOre = step(0.5, oreT);
float3 oc = oreT < 1.5 ? float3(1.0, 0.45, 0.12) : (oreT < 2.5 ? float3(0.15, 0.75, 1.0) : float3(1.0, 0.25, 0.8));
float3 cell = floor(Pm * 2.2);
float facet = H3(cell);
float k = smoothstep(0.05, 0.6, oreS) * isOre;
col = lerp(col, oc * (0.35 + 0.65 * facet), k * 0.85);
float basalt = saturate(UV1.x);
// Dense basalt: glassy black columns. The jointing lines (hexagonal in plan, so vertical stripes on a wall) give it a shape in the lamp;
// before, a mirror-smooth black surface with a noise normal showed only bright squiggles of specular.
float2 hp = Pm.xy * 0.9;
float2 hq = float2(hp.x * 1.1547, hp.y + hp.x * 0.57735);
float2 hi = floor(hq);
float2 hf = frac(hq);
float hv = floor(frac((hi.x + hi.y) / 3.0 + 0.001) * 3.0);
float hca = step(1.0, hv);
float hcb = step(2.0, hv);
float2 hma = step(hf.xy, hf.yx);
float hexd = dot(hma, 1.0 - hf.yx + hca * (hf.x + hf.y - 1.0) + hcb * (hf.yx - 2.0 * hf.xy));
float joint = smoothstep(0.07, 0.0, hexd);
col = lerp(col, float3(0.030, 0.034, 0.042) * (0.6 + 0.8 * mott) * (1.0 - 0.75 * joint), basalt);
float e = 0.12;
float3 p1 = Pm * 0.55;
float n0 = VN(p1);
float3 q1 = p1 + float3(e, 0, 0); float nx = VN(q1);
float3 q2 = p1 + float3(0, e, 0); float ny = VN(q2);
float3 q3 = p1 + float3(0, 0, e); float nz = VN(q3);
float3 g1 = float3(nx - n0, ny - n0, nz - n0) / e;
float3 p2 = Pm * 2.1;
float m0 = VN(p2);
float3 r1 = p2 + float3(e, 0, 0); float mx = VN(r1);
float3 r2 = p2 + float3(0, e, 0); float my = VN(r2);
float3 r3 = p2 + float3(0, 0, e); float mz = VN(r3);
float3 g2 = float3(mx - m0, my - m0, mz - m0) / e;
float3 t1 = d2 + float3(e, 0, 0);
float3 t2 = d2 + float3(0, e, 0);
float3 t3 = d2 + float3(0, 0, e);
float3 g3 = float3(VN(t1) - f2, VN(t2) - f2, VN(t3) - f2) / e;
float3 gg = (g1 * 0.45 + g2 * 0.22 + g3 * 0.10 + (blockRnd - 0.5) * 0.55) * lerp(1.0, 0.45, basalt);
float3 fn = normalize(float3(H3(cell) - 0.5, H3(cell + 7.1) - 0.5, H3(cell + 3.3) - 0.5));
gg += fn * 0.9 * smoothstep(0.1, 0.7, oreS) * isOre;
float3 Nw = normalize(N);
Nw = normalize(Nw - gg + dot(gg, Nw) * Nw);
Nrm = Nw;
Rough = lerp(lerp(lerp(0.58, 0.82, mott), 0.18, k), 0.5, basalt);   // wet rock: a sheen that follows the bumps (was 0.9, basalt 0.25)
float3 amb = lerp(Ground, Sky, 0.5 + 0.5 * Nw.z);
float spark = pow(H3(cell + floor(Tm * 0.7)), 18.0) * isOre;
float glow = smoothstep(0.1, 0.8, oreS) * (0.05 + 0.1 * sin(Tm * 1.3 + facet * 6.28) + 1.5 * spark) * isOre;
Emis = col * amb + oc * glow;
#undef VN
#undef SM3
#undef H3
return col;
"""


def run(ctx):
    h = _load(ctx)
    MEL = ctx.mel
    mpc = h.ensure_mpc(ctx)
    mat = ctx.get_or_create(h.MAT_DIR, "M_Terrain", unreal.Material, unreal.MaterialFactoryNew())
    MEL.delete_all_material_expressions(mat)
    mat.set_editor_property("blend_mode", unreal.BlendMode.BLEND_OPAQUE)
    mat.set_editor_property("shading_model", unreal.MaterialShadingModel.MSM_DEFAULT_LIT)
    mat.set_editor_property("two_sided", False)
    ctx.try_set(mat, "tangent_space_normal", False, "world-space normal")

    vcol = MEL.create_material_expression(mat, unreal.MaterialExpressionVertexColor, -900, -300)
    uv0 = MEL.create_material_expression(mat, unreal.MaterialExpressionTextureCoordinate, -900, -150)
    uv0.set_editor_property("coordinate_index", 0)
    uv1 = MEL.create_material_expression(mat, unreal.MaterialExpressionTextureCoordinate, -900, 0)
    uv1.set_editor_property("coordinate_index", 1)
    wpos = MEL.create_material_expression(mat, unreal.MaterialExpressionWorldPosition, -900, 150)
    vnrm = MEL.create_material_expression(mat, unreal.MaterialExpressionVertexNormalWS, -900, 300)
    tm = MEL.create_material_expression(mat, unreal.MaterialExpressionTime, -900, 450)
    sky = h.mpc_param(ctx, mat, mpc, "AmbientWaterSky", -900, 600)
    gnd = h.mpc_param(ctx, mat, mpc, "AmbientWaterGround", -900, 700)

    c = h.custom_node(ctx, mat, CODE, ["VC", "Cut", "UV0", "UV1", "WP", "N", "Tm", "Sky", "Ground"], unreal.CustomMaterialOutputType.CMOT_FLOAT3, -500, 0,
                      extra_outputs=[("Emis", unreal.CustomMaterialOutputType.CMOT_FLOAT3), ("Rough", unreal.CustomMaterialOutputType.CMOT_FLOAT1), ("Nrm", unreal.CustomMaterialOutputType.CMOT_FLOAT3)],
                      description="M_Terrain shading")
    for src, out, name in ((vcol, "", "VC"), (vcol, "A", "Cut"), (uv0, "", "UV0"), (uv1, "", "UV1"), (wpos, "", "WP"), (vnrm, "", "N"), (tm, "", "Tm"), (sky, "", "Sky"), (gnd, "", "Ground")):
        if not MEL.connect_material_expressions(src, out, c, name):
            ctx.fail("M_Terrain: could not connect the input " + name)
    MEL.connect_material_property(c, "", unreal.MaterialProperty.MP_BASE_COLOR)
    MEL.connect_material_property(c, "Emis", unreal.MaterialProperty.MP_EMISSIVE_COLOR)
    MEL.connect_material_property(c, "Rough", unreal.MaterialProperty.MP_ROUGHNESS)
    MEL.connect_material_property(c, "Nrm", unreal.MaterialProperty.MP_NORMAL)
    MEL.recompile_material(mat)
    ctx.eal.save_loaded_asset(mat)

    mi = ctx.get_or_create(h.MAT_DIR, "MI_Terrain", unreal.MaterialInstanceConstant, unreal.MaterialInstanceConstantFactoryNew())
    mi.set_editor_property("parent", mat)
    MEL.update_material_instance(mi)
    ctx.eal.save_loaded_asset(mi)
    ctx.log("M_Terrain / MI_Terrain ready")


def _load(ctx):
    spec = importlib.util.spec_from_file_location("descent_helpers_look", os.path.join(ctx.project_dir, "Tools", "ue", "_helpers_look.py"))
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module