"""
Assets from polyhaven.com (CC0) in the station file: models/anchorage/source_polyhaven, fetched by
build/fetch_polyhaven.py.

    Haven().place(asset, matrix, collection)   one instance of a model; its mesh is imported once (glTF) and shared
    Haven().surface(asset, size)               a material from a texture set, `size` metres to one repeat (by UVs)
    Haven().world(asset)                       a world lit by an HDRI

A model's parts are joined into one mesh ["ph_asset"], so an instance is one object, like a kit piece.
"""
import json
import math
import os

import bmesh
import bpy
from mathutils import Vector

HERE = os.path.dirname(os.path.abspath(__file__))
SRC = os.path.join(os.path.dirname(HERE), "source_polyhaven")
UV_TILE = 4.4


class Haven(object):
    def __init__(self, log=print):
        self.log = log
        self._mesh = {me["ph_asset"]: me for me in bpy.data.meshes if me.get("ph_asset")}
        self._mat = {m["ph_surface"]: m for m in bpy.data.materials if m.get("ph_surface")}

    def has(self, asset):
        return os.path.isdir(os.path.join(SRC, "models", asset))

    def mesh(self, asset):
        if asset in self._mesh:
            return self._mesh[asset]
        folder = os.path.join(SRC, "models", asset)
        gltf = [f for f in os.listdir(folder) if f.endswith(".gltf")] if os.path.isdir(folder) else []
        if not gltf:
            self.log("Poly Haven model not fetched: %s" % asset)
            self._mesh[asset] = None
            return None
        before = set(bpy.data.objects)
        bpy.ops.import_scene.gltf(filepath=os.path.join(folder, gltf[0]))
        bpy.context.view_layer.update()
        new = [o for o in bpy.data.objects if o not in before]
        bm = bmesh.new()
        mats = []
        for ob in new:
            if ob.type != "MESH":
                continue
            remap = []
            for m in ob.data.materials:
                if m not in mats:
                    mats.append(m)
                remap.append(mats.index(m))
            tmp = ob.data.copy()
            tmp.transform(ob.matrix_world)
            n0 = len(bm.faces)
            bm.from_mesh(tmp)
            bm.faces.ensure_lookup_table()
            if remap:
                for f in bm.faces[n0:]:
                    f.material_index = remap[min(f.material_index, len(remap) - 1)]
            bpy.data.meshes.remove(tmp)
        me = bpy.data.meshes.new("PH_" + asset)
        bm.to_mesh(me)
        bm.free()
        for m in mats:
            me.materials.append(m)
            if m is not None:
                m["ph_material_of"] = asset
        try:
            me.set_sharp_from_angle(angle=math.radians(42.0))
        except Exception:  # noqa: BLE001
            pass
        me["ph_asset"] = asset
        me["ph_source"] = "source_polyhaven/models/%s/%s" % (asset, gltf[0])
        me.use_fake_user = True
        for ob in new:
            old = ob.data if ob.type == "MESH" else None
            bpy.data.objects.remove(ob)
            if old is not None and old.users == 0:
                bpy.data.meshes.remove(old)
        self._mesh[asset] = me
        return me

    def box(self, asset):
        me = self.mesh(asset)
        if me is None:
            return None
        lo = Vector((min(v.co.x for v in me.vertices), min(v.co.y for v in me.vertices), min(v.co.z for v in me.vertices)))
        hi = Vector((max(v.co.x for v in me.vertices), max(v.co.y for v in me.vertices), max(v.co.z for v in me.vertices)))
        return lo, hi

    def place(self, asset, matrix, col, name=None):
        me = self.mesh(asset)
        if me is None:
            return None
        ob = bpy.data.objects.new(name or asset, me)
        ob.matrix_world = matrix
        col.objects.link(ob)
        ob["ph_asset"] = asset
        return ob

    def surface(self, asset, size):
        """A material from a texture set: colour (darkened by its occlusion map), roughness, normal (OpenGL)."""
        key = "%s@%.2f" % (asset, size)
        if key in self._mat:
            return self._mat[key]
        folder = os.path.join(SRC, "textures", asset)
        files = {}
        for dp, dn, fn in os.walk(folder):
            for f in fn:
                low = f.lower()
                for tag, k in (("_diff_", "diff"), ("_nor_gl_", "nor"), ("_rough_", "rough"), ("_ao_", "ao"), ("_arm_", "arm")):
                    if tag in low:
                        files[k] = os.path.join(dp, f)
        mat = bpy.data.materials.new("PH_" + asset)
        mat["ph_surface"] = key
        mat["ph_asset"] = asset
        mat["ph_size"] = size
        # for Unreal: a material of the station's own, made from these files (build/export_layout.py lists it)
        root = os.path.dirname(HERE)
        mat["ue_material"] = "/Game/Anchorage/Materials/M_ANCH_PH_" + asset
        mat["ue_scan_size"] = size
        mat["ue_scan_files"] = json.dumps({u: os.path.relpath(files[k], root).replace("\\", "/")
                                           for u, k in (("BaseColor", "diff"), ("Normal", "nor"), ("Roughness", "rough"), ("AO", "ao"))
                                           if k in files})
        mat["ue_note"] = "polyhaven.com texture set %s (CC0)" % asset
        mat.use_nodes = True
        nt = mat.node_tree
        nt.nodes.clear()
        out = nt.nodes.new("ShaderNodeOutputMaterial")
        out.location = (400, 0)
        bsdf = nt.nodes.new("ShaderNodeBsdfPrincipled")
        bsdf.location = (60, 0)
        nt.links.new(bsdf.outputs["BSDF"], out.inputs["Surface"])
        tc = nt.nodes.new("ShaderNodeTexCoord")
        tc.location = (-1300, 0)
        mp = nt.nodes.new("ShaderNodeMapping")
        mp.location = (-1100, 0)
        mp.inputs["Scale"].default_value = (UV_TILE / size, UV_TILE / size, 1.0)
        nt.links.new(tc.outputs["UV"], mp.inputs["Vector"])

        def image(k, y, colour):
            if k not in files:
                return None
            node = nt.nodes.new("ShaderNodeTexImage")
            node.location = (-860, y)
            node.image = bpy.data.images.load(files[k], check_existing=True)
            node.image.colorspace_settings.name = "sRGB" if colour else "Non-Color"
            nt.links.new(mp.outputs["Vector"], node.inputs["Vector"])
            return node
        base, ao = image("diff", 320, True), image("ao", 40, False)
        if base is not None:
            src = base.outputs["Color"]
            if ao is not None:
                mix = nt.nodes.new("ShaderNodeMix")
                mix.data_type, mix.blend_type = "RGBA", "MULTIPLY"
                mix.location = (-420, 240)
                mix.inputs[0].default_value = 0.7
                nt.links.new(src, mix.inputs[6])
                nt.links.new(ao.outputs["Color"], mix.inputs[7])
                src = mix.outputs[2]
            nt.links.new(src, bsdf.inputs["Base Color"])
        arm = image("arm", -240, False)
        if arm is not None:
            sep = nt.nodes.new("ShaderNodeSeparateColor")
            sep.location = (-420, -240)
            nt.links.new(arm.outputs["Color"], sep.inputs["Color"])
            nt.links.new(sep.outputs["Green"], bsdf.inputs["Roughness"])
            nt.links.new(sep.outputs["Blue"], bsdf.inputs["Metallic"])
        else:
            rgh = image("rough", -240, False)
            if rgh is not None:
                nt.links.new(rgh.outputs["Color"], bsdf.inputs["Roughness"])
        nrm = image("nor", -540, False)
        if nrm is not None:
            nm = nt.nodes.new("ShaderNodeNormalMap")
            nm.location = (-420, -540)
            nt.links.new(nrm.outputs["Color"], nm.inputs["Color"])
            nt.links.new(nm.outputs["Normal"], bsdf.inputs["Normal"])
        self._mat[key] = mat
        return mat

    def world(self, asset, strength=1.0, turn=0.0, rebuild=False):
        """A world lit by the HDRI. What the camera itself sees is a plain overcast sky, lighter at the horizon:
        the picture has the land it was taken on round its horizon (trees, poles), which has no place at sea.
        Light and reflections come from the picture."""
        name = "PH_World_" + asset
        w = bpy.data.worlds.get(name)
        if w is not None and not rebuild:
            return w
        folder = os.path.join(SRC, "hdris", asset)
        hdr = [f for f in os.listdir(folder) if f.endswith(".hdr")] if os.path.isdir(folder) else []
        if not hdr:
            self.log("Poly Haven sky not fetched: %s" % asset)
            return None
        if w is None:
            w = bpy.data.worlds.new(name)
        w.use_nodes = True
        nt = w.node_tree
        nt.nodes.clear()
        out = nt.nodes.new("ShaderNodeOutputWorld")
        out.location = (600, 0)
        bg = nt.nodes.new("ShaderNodeBackground")
        bg.location = (100, 120)
        env = nt.nodes.new("ShaderNodeTexEnvironment")
        env.location = (-220, 120)
        mp = nt.nodes.new("ShaderNodeMapping")
        mp.location = (-420, 120)
        tc = nt.nodes.new("ShaderNodeTexCoord")
        tc.location = (-620, 0)
        env.image = bpy.data.images.load(os.path.join(folder, hdr[0]), check_existing=True)
        mp.inputs["Rotation"].default_value = (0.0, 0.0, math.radians(turn))
        bg.inputs["Strength"].default_value = strength
        nt.links.new(tc.outputs["Generated"], mp.inputs["Vector"])
        nt.links.new(mp.outputs["Vector"], env.inputs["Vector"])
        nt.links.new(env.outputs["Color"], bg.inputs["Color"])
        # the sky the camera sees: a gradient by height
        sep = nt.nodes.new("ShaderNodeSeparateXYZ")
        sep.location = (-420, -220)
        nt.links.new(tc.outputs["Generated"], sep.inputs["Vector"])
        ramp = nt.nodes.new("ShaderNodeValToRGB")
        ramp.location = (-220, -220)
        ramp.color_ramp.elements[0].position = 0.0
        ramp.color_ramp.elements[0].color = (0.80, 0.83, 0.86, 1.0)
        ramp.color_ramp.elements[1].position = 0.55
        ramp.color_ramp.elements[1].color = (0.56, 0.61, 0.68, 1.0)
        nt.links.new(sep.outputs["Z"], ramp.inputs["Fac"])
        sky = nt.nodes.new("ShaderNodeBackground")
        sky.location = (100, -220)
        sky.inputs["Strength"].default_value = strength
        nt.links.new(ramp.outputs["Color"], sky.inputs["Color"])
        lp = nt.nodes.new("ShaderNodeLightPath")
        lp.location = (100, 380)
        mix = nt.nodes.new("ShaderNodeMixShader")
        mix.location = (360, 0)
        nt.links.new(lp.outputs["Is Camera Ray"], mix.inputs["Fac"])
        nt.links.new(bg.outputs["Background"], mix.inputs[1])
        nt.links.new(sky.outputs["Background"], mix.inputs[2])
        nt.links.new(mix.outputs["Shader"], out.inputs["Surface"])
        w["ph_asset"] = asset
        w.use_fake_user = True
        return w