"""
Helpers of the world / look track for the asset scripts (not a module of the import runner: the leading underscore keeps
import_assets.py from running it). Builds the shared look assets under /Game/Materials (DESIGN_UE.md sections 7.1, 7.3, 7.9):

  MPC_Descent        material parameter collection written once per frame by ADescentEnvironment
  M_DescentSurface   opaque lit master: BaseColor, Metallic, Roughness, EmissiveColor x EmissiveStrength, AmbientMix + the water ambient term
  M_DescentGlass     translucent glass master (V10: real glass)

Every function is safe to call again: it rebuilds the cheap node graphs and leaves instances alone. A module loads this
file with `load(ctx)` because the runner does not put Tools/ue on sys.path.
"""
import importlib.util
import os

import unreal

MAT_DIR = "/Game/Materials"

# Parameters of MPC_Descent (section 7.1). Scalars have a default; vectors have an RGBA default.
MPC_SCALARS = [
    ("Time", 0.0), ("SurfaceLight", 0.36), ("FogDensity", 0.0177), ("FloorGlow", 0.0), ("Pulse", 0.0), ("HullRim", 0.0),
    ("BeamAmt", 0.0), ("Jack", 0.0), ("Hunger", 0.0), ("Damage", 0.0), ("Fade", 0.0), ("Flash", 0.0), ("Glitch", 0.0),
    ("Pressure", 0.0), ("Grain", 0.0),
    # WaterAmount scales the post-process water (0 = off, Descent.Look.Set WaterPP 0); SnowAmount is the marine snow switch.
    ("WaterAmount", 1.0), ("SnowAmount", 1.0),
    # 1 while the camera is in dry air (inside the carrier): only then does PP_Water spare the stencil-1 pixels (section 7.4).
    ("DryCamera", 0.0),
]
MPC_VECTORS = [
    ("WaterTop", (0.028, 0.275, 0.323, 1.0)), ("WaterDeep", (0.0018, 0.0159, 0.0232, 1.0)), ("SunDir", (-0.12, 0.20, 1.0, 0.0)),
    ("AmbientWaterSky", (0.0, 0.0, 0.0, 1.0)), ("AmbientWaterGround", (0.0, 0.0, 0.0, 1.0)), ("AmbientDry", (0.0, 0.0, 0.0, 1.0)),
    ("CarrierOrigin", (0.0, 0.0, 0.0, 0.0)), ("CarrierFwd", (1.0, 0.0, 0.0, 0.0)), ("CarrierUp", (0.0, 0.0, 1.0, 0.0)),
    ("BeamPos", (0.0, 0.0, 0.0, 0.0)), ("BeamDir", (1.0, 0.0, 0.0, 0.0)), ("ShearVel", (0.0, 0.0, 0.0, 0.0)), ("FlashColor", (1.0, 1.0, 1.0, 0.0)),
]


def load(ctx):
    """Returns this module (for assets_*.py files that cannot import it by name)."""
    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


def scalar_param(ctx, mat, name, default, x, y):
    n = ctx.mel.create_material_expression(mat, unreal.MaterialExpressionScalarParameter, x, y)
    n.set_editor_property("parameter_name", name)
    n.set_editor_property("default_value", float(default))
    return n


def vec_param(ctx, mat, name, default, x, y):
    n = ctx.mel.create_material_expression(mat, unreal.MaterialExpressionVectorParameter, x, y)
    n.set_editor_property("parameter_name", name)
    n.set_editor_property("default_value", unreal.LinearColor(default[0], default[1], default[2], default[3] if len(default) > 3 else 1.0))
    return n


def mpc_param(ctx, mat, mpc, name, x, y):
    """A node that reads one parameter of MPC_Descent."""
    n = ctx.mel.create_material_expression(mat, unreal.MaterialExpressionCollectionParameter, x, y)
    n.set_editor_property("collection", mpc)
    n.set_editor_property("parameter_name", name)
    return n


def custom_node(ctx, mat, code, input_names, out_type, x, y, extra_outputs=(), description="Custom"):
    """A Custom (HLSL) node. extra_outputs = [(name, unreal.CustomMaterialOutputType)]."""
    c = ctx.mel.create_material_expression(mat, unreal.MaterialExpressionCustom, x, y)
    c.set_editor_property("code", code)
    c.set_editor_property("description", description)
    c.set_editor_property("output_type", out_type)
    ins = []
    for n in input_names:
        ci = unreal.CustomInput()
        ci.set_editor_property("input_name", n)
        ins.append(ci)
    c.set_editor_property("inputs", ins)
    if extra_outputs:
        outs = []
        for name, t in extra_outputs:
            co = unreal.CustomOutput()
            co.set_editor_property("output_name", name)
            co.set_editor_property("output_type", t)
            outs.append(co)
        c.set_editor_property("additional_outputs", outs)
    return c


def ensure_mpc(ctx):
    path = MAT_DIR + "/MPC_Descent"
    ctx.eal.make_directory(MAT_DIR)
    mpc = ctx.get_or_create(MAT_DIR, "MPC_Descent", unreal.MaterialParameterCollection, unreal.MaterialParameterCollectionFactoryNew())
    # Rebuilt at most once per import run: rewriting the parameter list gives the parameters new ids, and every material that
    # already holds a CollectionParameter node for the old ids then fails to compile ("CollectionParameter has invalid parameter None").
    # With -Force the modules assets_10, 20, 40 and 41 all call this, and the later calls used to break the earlier materials.
    if ctx.shared.get("mpc_done"):
        return mpc
    want_s = [n for n, _ in MPC_SCALARS]
    want_v = [n for n, _ in MPC_VECTORS]
    have_s = [str(n) for n in mpc.get_scalar_parameter_names()]
    have_v = [str(n) for n in mpc.get_vector_parameter_names()]
    if have_s == want_s and have_v == want_v and not ctx.force:
        ctx.shared["mpc_done"] = True
        return mpc
    sp = []
    for n, d in MPC_SCALARS:
        p = unreal.CollectionScalarParameter()
        p.set_editor_property("parameter_name", n)
        p.set_editor_property("default_value", float(d))
        sp.append(p)
    vp = []
    for n, d in MPC_VECTORS:
        p = unreal.CollectionVectorParameter()
        p.set_editor_property("parameter_name", n)
        p.set_editor_property("default_value", unreal.LinearColor(d[0], d[1], d[2], d[3]))
        vp.append(p)
    mpc.set_editor_property("scalar_parameters", sp)
    mpc.set_editor_property("vector_parameters", vp)
    ctx.eal.save_loaded_asset(mpc)
    ctx.shared["mpc_done"] = True
    ctx.log("MPC_Descent: %d scalars, %d vectors (names read back: %d / %d)" % (len(sp), len(vp), len(mpc.get_scalar_parameter_names()), len(mpc.get_vector_parameter_names())))
    return mpc


SURFACE_CODE = """
float3 amb = lerp(lerp(Ground, Sky, 0.5 + 0.5 * N.z), Dry, saturate(Mix));
return Base * amb + Emis;
"""


def ensure_surface_master(ctx, mpc):
    """M_DescentSurface. Parameter names are the Kazah master's (BaseColor, Metallic, Roughness, EmissiveColor, EmissiveStrength)
    so the 14 Kazah instances can be re-parented without changes; AmbientMix, EnvReflect, PanelScale and Hazard are the design's."""
    MEL = ctx.mel
    mat = ctx.get_or_create(MAT_DIR, "M_DescentSurface", 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)
    base = vec_param(ctx, mat, "BaseColor", (0.5, 0.5, 0.5, 1.0), -900, -300)
    metal = scalar_param(ctx, mat, "Metallic", 0.0, -900, -100)
    rough = scalar_param(ctx, mat, "Roughness", 0.5, -900, 100)
    emc = vec_param(ctx, mat, "EmissiveColor", (0.0, 0.0, 0.0, 1.0), -900, 300)
    ems = scalar_param(ctx, mat, "EmissiveStrength", 0.0, -900, 450)
    mix = scalar_param(ctx, mat, "AmbientMix", 0.0, -900, 600)
    scalar_param(ctx, mat, "EnvReflect", 0.0, -900, 750)   # reserved: the fake environment reflection (section 7.3) is not built yet
    scalar_param(ctx, mat, "PanelScale", 0.0, -900, 900)   # reserved for the carrier track
    scalar_param(ctx, mat, "Hazard", 0.0, -900, 1050)      # reserved for the carrier track
    emul = MEL.create_material_expression(mat, unreal.MaterialExpressionMultiply, -650, 380)
    MEL.connect_material_expressions(emc, "", emul, "A")
    MEL.connect_material_expressions(ems, "", emul, "B")
    nrm = MEL.create_material_expression(mat, unreal.MaterialExpressionPixelNormalWS, -650, 0)
    sky = mpc_param(ctx, mat, mpc, "AmbientWaterSky", -650, -450)
    gnd = mpc_param(ctx, mat, mpc, "AmbientWaterGround", -650, -350)
    dry = mpc_param(ctx, mat, mpc, "AmbientDry", -650, -250)
    c = custom_node(ctx, mat, SURFACE_CODE, ["Base", "N", "Sky", "Ground", "Dry", "Mix", "Emis"], unreal.CustomMaterialOutputType.CMOT_FLOAT3, -350, 0, description="water ambient")
    MEL.connect_material_expressions(base, "", c, "Base")
    MEL.connect_material_expressions(nrm, "", c, "N")
    MEL.connect_material_expressions(sky, "", c, "Sky")
    MEL.connect_material_expressions(gnd, "", c, "Ground")
    MEL.connect_material_expressions(dry, "", c, "Dry")
    MEL.connect_material_expressions(mix, "", c, "Mix")
    MEL.connect_material_expressions(emul, "", c, "Emis")
    MEL.connect_material_property(base, "", unreal.MaterialProperty.MP_BASE_COLOR)
    MEL.connect_material_property(metal, "", unreal.MaterialProperty.MP_METALLIC)
    MEL.connect_material_property(rough, "", unreal.MaterialProperty.MP_ROUGHNESS)
    MEL.connect_material_property(c, "", unreal.MaterialProperty.MP_EMISSIVE_COLOR)
    MEL.recompile_material(mat)
    ctx.eal.save_loaded_asset(mat)
    return mat


def ensure_glass_master(ctx):
    """M_DescentGlass: the M0 glass (lit translucency, per-pixel lighting, two-sided, fresnel opacity) with the design's opacity range."""
    MEL = ctx.mel
    mat = ctx.get_or_create(MAT_DIR, "M_DescentGlass", unreal.Material, unreal.MaterialFactoryNew())
    MEL.delete_all_material_expressions(mat)
    mat.set_editor_property("blend_mode", unreal.BlendMode.BLEND_TRANSLUCENT)
    mat.set_editor_property("shading_model", unreal.MaterialShadingModel.MSM_DEFAULT_LIT)
    mat.set_editor_property("translucency_lighting_mode", unreal.TranslucencyLightingMode.TLM_SURFACE_PER_PIXEL_LIGHTING)
    mat.set_editor_property("two_sided", True)
    tint = vec_param(ctx, mat, "BaseColor", (0.52, 0.87, 0.96, 1.0), -900, -300)
    rough = scalar_param(ctx, mat, "Roughness", 0.06, -900, -100)
    op_min = scalar_param(ctx, mat, "OpacityFaceOn", 0.06, -900, 100)
    op_max = scalar_param(ctx, mat, "OpacityGrazing", 0.45, -900, 250)
    fres_pow = scalar_param(ctx, mat, "FresnelExponent", 3.0, -900, 400)
    ior = scalar_param(ctx, mat, "IOR", 1.03, -900, 550)
    fres = MEL.create_material_expression(mat, unreal.MaterialExpressionFresnel, -600, 300)
    MEL.connect_material_expressions(fres_pow, "", fres, "ExponentIn")
    lerp = MEL.create_material_expression(mat, unreal.MaterialExpressionLinearInterpolate, -320, 200)
    MEL.connect_material_expressions(op_min, "", lerp, "A")
    MEL.connect_material_expressions(op_max, "", lerp, "B")
    MEL.connect_material_expressions(fres, "", lerp, "Alpha")
    spec = MEL.create_material_expression(mat, unreal.MaterialExpressionConstant, -600, -500)
    spec.set_editor_property("r", 1.0)
    metal = MEL.create_material_expression(mat, unreal.MaterialExpressionConstant, -600, -420)
    metal.set_editor_property("r", 0.0)
    MEL.connect_material_property(tint, "", unreal.MaterialProperty.MP_BASE_COLOR)
    MEL.connect_material_property(metal, "", unreal.MaterialProperty.MP_METALLIC)
    MEL.connect_material_property(spec, "", unreal.MaterialProperty.MP_SPECULAR)
    MEL.connect_material_property(rough, "", unreal.MaterialProperty.MP_ROUGHNESS)
    MEL.connect_material_property(lerp, "", unreal.MaterialProperty.MP_OPACITY)
    MEL.connect_material_property(ior, "", unreal.MaterialProperty.MP_REFRACTION)
    MEL.recompile_material(mat)
    ctx.eal.save_loaded_asset(mat)
    return mat