FModel/FModel/Views/Snooper/Renderer.cs

using System;
using System.Linq;
using System.Numerics;
using System.Threading;
using System.Windows;
using CUE4Parse_Conversion.Animations;
using CUE4Parse_Conversion.Meshes;
using CUE4Parse.UE4.Assets.Exports;
using CUE4Parse.UE4.Assets.Exports.Animation;
using CUE4Parse.UE4.Assets.Exports.Material;
using CUE4Parse.UE4.Assets.Exports.SkeletalMesh;
using CUE4Parse.UE4.Assets.Exports.StaticMesh;
using CUE4Parse.UE4.Assets.Exports.Texture;
using CUE4Parse.UE4.Objects.Core.Math;
using CUE4Parse.UE4.Objects.Engine;
using CUE4Parse.UE4.Objects.UObject;
using FModel.Creator;
using FModel.Extensions;
using FModel.Settings;
using FModel.Views.Snooper.Buffers;
using FModel.Views.Snooper.Lights;
using FModel.Views.Snooper.Models;
using FModel.Views.Snooper.Models.Animations;
using FModel.Views.Snooper.Shading;

namespace FModel.Views.Snooper;

public class Renderer : IDisposable
{
    private readonly Skybox _skybox;
    private readonly Grid _grid;
    private Shader _shader;
    private Shader _outline;
    private Shader _light;

    public bool ShowSkybox;
    public bool ShowGrid;
    public bool ShowLights;
    public int VertexColor;

    public Camera CameraOp { get; }
    public PickingTexture Picking { get; }
    public Options Options { get; }

    public Renderer(int width, int height)
    {
        _skybox = new Skybox();
        _grid = new Grid();

        CameraOp = new Camera();
        Picking = new PickingTexture(width, height);
        Options = new Options();

        ShowSkybox = UserSettings.Default.ShowSkybox;
        ShowGrid = UserSettings.Default.ShowGrid;
        VertexColor = 0; // default
    }

    public void Load(CancellationToken cancellationToken, UObject export)
    {
        ShowLights = false;
        switch (export)
        {
            case UStaticMesh st:
                LoadStaticMesh(st);
                break;
            case USkeletalMesh sk:
                LoadSkeletalMesh(sk);
                break;
            case UMaterialInstance mi:
                LoadMaterialInstance(mi);
                break;
            case UWorld wd:
                LoadWorld(cancellationToken, wd, Transform.Identity);
                CameraOp.Mode = Camera.WorldMode.FlyCam;
                break;
        }
    }

    public void Swap(UMaterialInstance unrealMaterial)
    {
        if (!Options.TryGetModel(out var model) || !Options.TryGetSection(model, out var section)) return;

        model.Materials[section.MaterialIndex].SwapMaterial(unrealMaterial);
        Application.Current.Dispatcher.Invoke(() => model.Materials[section.MaterialIndex].Setup(Options, model.UvCount));
        Options.SwapMaterial(false);
    }

    public void Animate(UAnimSequence animSequence)
    {
        if (!Options.TryGetModel(out var model) || !model.Skeleton.IsLoaded ||
            model.Skeleton?.UnrealSkeleton.ConvertAnims(animSequence) is not { } anim || anim.Sequences.Count == 0)
            return;

        model.Skeleton.SetAnimation(anim);
        Options.AnimateMesh(false);
    }

    public void Setup()
    {
        _skybox.Setup();
        _grid.Setup();

        _shader = new Shader();
        _outline = new Shader("outline");
        _light = new Shader("light");

        Picking.Setup();
        Options.SetupModelsAndLights();
    }

    public void Render()
    {
        var viewMatrix = CameraOp.GetViewMatrix();
        var projMatrix = CameraOp.GetProjectionMatrix();

        if (ShowSkybox) _skybox.Render(viewMatrix, projMatrix);
        if (ShowGrid) _grid.Render(viewMatrix, projMatrix, CameraOp.Near, CameraOp.Far);

        _shader.Render(viewMatrix, CameraOp.Position, projMatrix);
        for (int i = 0; i < 5; i++)
            _shader.SetUniform($"bVertexColors[{i}]", i == VertexColor);

        // render model pass
        foreach (var model in Options.Models.Values)
        {
            model.UpdateMatrices(Options);
            if (!model.Show) continue;
            model.Render(_shader);
        }

        {   // light pass
            var uNumLights = Math.Min(Options.Lights.Count, 100);
            _shader.SetUniform("uNumLights", ShowLights ? uNumLights : 0);

            if (ShowLights)
                for (int i = 0; i < uNumLights; i++)
                    Options.Lights[i].Render(i, _shader);

            _light.Render(viewMatrix, projMatrix);
            for (int i = 0; i < uNumLights; i++)
                Options.Lights[i].Render(_light);
        }

        // outline pass
        if (Options.TryGetModel(out var selected) && selected.Show)
        {
            _outline.Render(viewMatrix, CameraOp.Position, projMatrix);
            selected.Render(_outline, true);
        }

        // picking pass (dedicated FBO, binding to 0 afterward)
        Picking.Render(viewMatrix, projMatrix, Options.Models);
    }

    private void LoadStaticMesh(UStaticMesh original)
    {
        var guid = original.LightingGuid;
        if (Options.TryGetModel(guid, out var model))
        {
            model.AddInstance(Transform.Identity);
            Application.Current.Dispatcher.Invoke(() => model.SetupInstances());
            return;
        }

        if (!original.TryConvert(out var mesh))
            return;

        Options.Models[guid] = new Model(original, mesh);
        Options.SelectModel(guid);
        SetupCamera(Options.Models[guid].Box);
    }

    private void LoadSkeletalMesh(USkeletalMesh original)
    {
        var guid = Guid.NewGuid();
        if (Options.Models.ContainsKey(guid) || !original.TryConvert(out var mesh)) return;

        Options.Models[guid] = new Model(original, mesh);
        Options.SelectModel(guid);
        SetupCamera(Options.Models[guid].Box);
    }

    private void LoadMaterialInstance(UMaterialInstance original)
    {
        if (!Utils.TryLoadObject("Engine/Content/EditorMeshes/EditorCube.EditorCube", out UStaticMesh editorCube))
            return;

        var guid = editorCube.LightingGuid;
        if (Options.TryGetModel(guid, out var model))
        {
            model.Materials[0].SwapMaterial(original);
            Application.Current.Dispatcher.Invoke(() => model.Materials[0].Setup(Options, model.UvCount));
            return;
        }

        if (!editorCube.TryConvert(out var mesh))
            return;

        Options.Models[guid] = new Cube(mesh, original);
        Options.SelectModel(guid);
        SetupCamera(Options.Models[guid].Box);
    }

    private void SetupCamera(FBox box) => CameraOp.Setup(box);

    private void LoadWorld(CancellationToken cancellationToken, UWorld original, Transform transform)
    {
        CameraOp.Setup(new FBox(FVector.ZeroVector, new FVector(0, 10, 10)));
        if (original.PersistentLevel.Load<ULevel>() is not { } persistentLevel)
            return;

        if (persistentLevel.TryGetValue(out FSoftObjectPath runtimeCell, "WorldPartitionRuntimeCell") &&
            Utils.TryLoadObject(runtimeCell.AssetPathName.Text.SubstringBeforeWithLast(".") + runtimeCell.SubPathString.SubstringAfterLast("."), out UObject worldPartition))
        {
            var position = worldPartition.GetOrDefault("Position", FVector.ZeroVector) * Constants.SCALE_DOWN_RATIO;
            var box = worldPartition.GetOrDefault("ContentBounds", new FBox(FVector.ZeroVector, FVector.OneVector));
            box *= MathF.Pow(Constants.SCALE_DOWN_RATIO, 2);
            CameraOp.Teleport(position, box, true);
        }

        var length = persistentLevel.Actors.Length;
        for (var i = 0; i < length; i++)
        {
            cancellationToken.ThrowIfCancellationRequested();

            if (persistentLevel.Actors[i].Load() is not { } actor ||
                actor.ExportType is "LODActor" or "SplineMeshActor")
                continue;

            Services.ApplicationService.ApplicationView.Status.UpdateStatusLabel($"{original.Name} ... {i}/{length}");
            WorldCamera(actor);
            // WorldLight(actor);
            WorldMesh(actor, transform);
            AdditionalWorlds(actor, transform.Matrix, cancellationToken);
        }
        Services.ApplicationService.ApplicationView.Status.UpdateStatusLabel($"{original.Name} ... {length}/{length}");
    }

    private void WorldCamera(UObject actor)
    {
        if (actor.ExportType != "LevelBounds" || !actor.TryGetValue(out FPackageIndex boxComponent, "BoxComponent") ||
            boxComponent.Load() is not { } boxObject) return;

        var direction = boxObject.GetOrDefault("RelativeLocation", FVector.ZeroVector) * Constants.SCALE_DOWN_RATIO;
        var position = boxObject.GetOrDefault("RelativeScale3D", FVector.OneVector) / 2f * Constants.SCALE_DOWN_RATIO;
        CameraOp.Setup(new FBox(direction, position));
    }

    private void WorldLight(UObject actor)
    {
        // if (!actor.TryGetValue(out FPackageIndex lightComponent, "LightComponent") ||
        //     lightComponent.Load() is not { } lightObject) return;
        //
        // Cache.Lights.Add(new PointLight(Cache.Icons["pointlight"], lightObject, FVector.ZeroVector));
    }

    private void WorldMesh(UObject actor, Transform transform)
    {
        if (!actor.TryGetValue(out FPackageIndex staticMeshComponent, "StaticMeshComponent", "StaticMesh", "Mesh", "LightMesh") ||
            staticMeshComponent.Load() is not { } staticMeshComp) return;

        if (!staticMeshComp.TryGetValue(out FPackageIndex staticMesh, "StaticMesh") && actor.Class is UBlueprintGeneratedClass)
        {
            foreach (var actorExp in actor.Class.Owner.GetExports())
                if (actorExp.TryGetValue(out staticMesh, "StaticMesh"))
                    break;
            var super = actor.Class.SuperStruct.Load<UBlueprintGeneratedClass>();
            if (staticMesh == null && super != null) { // look in parent struct if not found
                foreach (var actorExp in super.Owner.GetExports()) {
                    if (actorExp.TryGetValue(out staticMesh, "StaticMesh"))
                        break;
                }
            }
        }

        if (staticMesh?.Load() is not UStaticMesh m || m.Materials.Length < 1)
            return;

        var guid = m.LightingGuid;
        var t = new Transform
        {
            Relation = transform.Matrix,
            Position = staticMeshComp.GetOrDefault("RelativeLocation", FVector.ZeroVector) * Constants.SCALE_DOWN_RATIO,
            Rotation = staticMeshComp.GetOrDefault("RelativeRotation", FRotator.ZeroRotator).Quaternion(),
            Scale = staticMeshComp.GetOrDefault("RelativeScale3D", FVector.OneVector)
        };

        if (Options.TryGetModel(guid, out var model))
        {
            model.AddInstance(t);
        }
        else if (m.TryConvert(out var mesh))
        {
            model = new Model(m, mesh, t);
            model.TwoSided = actor.GetOrDefault("bMirrored", staticMeshComp.GetOrDefault("bDisallowMeshPaintPerInstance", model.TwoSided));

            if (actor.TryGetValue(out FPackageIndex baseMaterial, "BaseMaterial") &&
                actor.TryGetAllValues(out FPackageIndex[] textureData, "TextureData"))
            {
                var material = model.Materials.FirstOrDefault(x => x.Name == baseMaterial.Name);
                if (material is { IsUsed: true })
                {
                    for (int j = 0; j < textureData.Length; j++)
                    {
                        if (textureData[j]?.Load() is not { } textureDataIdx)
                            continue;

                        if (textureDataIdx.TryGetValue(out FPackageIndex overrideMaterial, "OverrideMaterial") &&
                            overrideMaterial.TryLoad(out var oMaterial) && oMaterial is UMaterialInterface oUnrealMaterial)
                            material.SwapMaterial(oUnrealMaterial);

                        WorldTextureData(material, textureDataIdx, "Diffuse", j switch
                        {
                            0 => "Diffuse",
                            > 0 => $"Diffuse_Texture_{j + 1}",
                            _ => CMaterialParams2.FallbackDiffuse
                        });
                        WorldTextureData(material, textureDataIdx, "Normal", j switch
                        {
                            0 => "Normals",
                            > 0 => $"Normals_Texture_{j + 1}",
                            _ => CMaterialParams2.FallbackNormals
                        });
                        WorldTextureData(material, textureDataIdx, "Specular", j switch
                        {
                            0 => "SpecularMasks",
                            > 0 => $"SpecularMasks_{j + 1}",
                            _ => CMaterialParams2.FallbackNormals
                        });
                    }
                }
            }

            if (staticMeshComp.TryGetValue(out FPackageIndex[] overrideMaterials, "OverrideMaterials"))
            {
                var max = model.Sections.Length - 1;
                for (var j = 0; j < overrideMaterials.Length; j++)
                {
                    if (j > max) break;
                    if (!model.Materials[model.Sections[j].MaterialIndex].IsUsed ||
                        overrideMaterials[j].Load() is not UMaterialInterface unrealMaterial) continue;
                    model.Materials[model.Sections[j].MaterialIndex].SwapMaterial(unrealMaterial);
                }
            }

            Options.Models[guid] = model;
        }

        if (actor.TryGetValue(out FPackageIndex treasureLight, "PointLight", "TreasureLight") &&
            treasureLight.TryLoad(out var pl1) && pl1.Template.TryLoad(out var pl2))
        {
            Options.Lights.Add(new PointLight(guid, Options.Icons["pointlight"], pl1, pl2, t));
        }
        if (actor.TryGetValue(out FPackageIndex spotLight, "SpotLight") &&
            spotLight.TryLoad(out var sl1) && sl1.Template.TryLoad(out var sl2))
        {
            Options.Lights.Add(new SpotLight(guid, Options.Icons["spotlight"], sl1, sl2, t));
        }
    }

    private void WorldTextureData(Material material, UObject textureData, string name, string key)
    {
        if (textureData.TryGetValue(out FPackageIndex package, name) && package.Load() is UTexture2D texture)
            material.Parameters.Textures[key] = texture;
    }

    private void AdditionalWorlds(UObject actor, Matrix4x4 relation, CancellationToken cancellationToken)
    {
        if (!actor.TryGetValue(out FSoftObjectPath[] additionalWorlds, "AdditionalWorlds") ||
            !actor.TryGetValue(out FPackageIndex staticMeshComponent, "StaticMeshComponent", "Mesh") ||
            staticMeshComponent.Load() is not { } staticMeshComp)
            return;

        var transform = new Transform
        {
            Relation = relation,
            Position = staticMeshComp.GetOrDefault("RelativeLocation", FVector.ZeroVector) * Constants.SCALE_DOWN_RATIO,
            Rotation = staticMeshComp.GetOrDefault("RelativeRotation", FRotator.ZeroRotator).Quaternion()
        };

        for (int j = 0; j < additionalWorlds.Length; j++)
            if (Utils.TryLoadObject(additionalWorlds[j].AssetPathName.Text, out UWorld w))
                LoadWorld(cancellationToken, w, transform);
    }

    public void WindowResized(int width, int height)
    {
        CameraOp.AspectRatio = width / (float) height;
        Picking.WindowResized(width, height);
    }

    public void Save()
    {
        UserSettings.Default.CameraMode = CameraOp.Mode;
        UserSettings.Default.ShowSkybox = ShowSkybox;
        UserSettings.Default.ShowGrid = ShowGrid;
    }

    public void Dispose()
    {
        _skybox?.Dispose();
        _grid?.Dispose();
        _shader?.Dispose();
        _outline?.Dispose();
        _light?.Dispose();
        Picking?.Dispose();
        Options?.Dispose();
    }
}