thief-mission-viewer/project/code/TMV/Mission.cs

using Godot;
using KeepersCompound.LGS;
using KeepersCompound.LGS.Database;
using KeepersCompound.LGS.Database.Chunks;
using KeepersCompound.TMV.UI;
using System;
using System.Collections.Generic;
using System.IO;

namespace KeepersCompound.TMV;

public partial class Mission : Node3D
{
	private readonly struct LightmapRectData
	{
		public readonly int cellIndex;
		public readonly int lightmapIndex;
		public readonly int textureId;
		public readonly int uvStart;
		public readonly int uvEnd;

		public LightmapRectData(int cellIndex, int lightmapIndex, int textureId, int uvStart, int uvEnd)
		{
			this.cellIndex = cellIndex;
			this.lightmapIndex = lightmapIndex;
			this.textureId = textureId;
			this.uvStart = uvStart;
			this.uvEnd = uvEnd;
		}
	}

	private const string OBJECT_MODELS_GROUP = "ObjectModels";

	[Export(PropertyHint.GlobalFile, "*.mis")]
	public string FileName { get; set; }
	[Export]
	public bool Build = false;
	[Export]
	public bool Clear = false;
	[Export]
	public bool Dump = false;
	public int LightmapLayers = 33;

	string _missionName;
	DbFile _file;
	TextureLoader _textureLoader;
	List<ShaderMaterial> _materials;
	Vector2I _lmLayerMask;

	public override void _Ready()
	{
		_materials = new List<ShaderMaterial>();
		_lmLayerMask = new Vector2I(~0, ~0);

		var lightmapToggler = GetNode<Control>("%LightmapToggler") as LightmapLayerToggler;
		lightmapToggler.Setup(this);

		var resourceSelector = GetNode<Control>("%ResourceSelector") as ResourceSelector;
		resourceSelector.ResourceSelected += (string campaign, string mission) =>
		{
			_missionName = mission;

			var context = Context.Instance;
			context.SetCampaign(campaign);
			FileName = context.CampaignResources.GetResourcePath(ResourceType.Mission, mission);
			Build = true;
		};
	}

	public override void _Process(double delta)
	{
		if (Build)
		{
			Timing.Reset();
			Timing.TimeStage("Build", () => RebuildMap());
			Timing.LogAll();
			Build = false;
		}
		if (Clear)
		{
			ClearMap();
			Clear = false;
		}
	}

	public override void _ShortcutInput(InputEvent @event)
	{
		if (@event is InputEventKey keyEvent && keyEvent.Pressed)
		{
			if (keyEvent.Keycode == Key.R)
			{
				Build = true;
			}
			if (keyEvent.Keycode == Key.L)
			{
				ToggleLightmap();
			}
			if (keyEvent.Keycode == Key.O)
			{
				ToggleObjectRendering();
			}
		}
	}

	public void ClearMap()
	{
		foreach (var node in GetChildren())
		{
			node.QueueFree();
		}

		_materials.Clear();
	}

	public void RebuildMap()
	{
		ClearMap();

		_textureLoader = new TextureLoader();
		Timing.TimeStage("DbFile Parse", () => _file = new(FileName));
		Timing.TimeStage("Register Textures", () => UseChunk<TxList>("TXLIST", RegisterTextures));
		Timing.TimeStage("Build WR", () => UseChunk<WorldRep>("WREXT", BuildWrMeshes));

		if (_file.Chunks.TryGetValue("BRLIST", out var brList))
		{
			var objHierarchy = Timing.TimeStage("Hierarchy", BuildHierarchy);
			Timing.TimeStage("Object Placement", () => PlaceObjects((BrList)brList, objHierarchy));
		}
	}

	public void ToggleObjectRendering()
	{
		// It might be "better" to use a parent node and just toggle that, but
		// the performance of this is completely fine so we'll stick with it
		foreach (var node in GetTree().GetNodesInGroup(OBJECT_MODELS_GROUP))
		{
			var node3d = node as Node3D;
			node3d.Visible = !node3d.Visible;
		}
	}

	public void ToggleLightmap()
	{
		if (_lmLayerMask == Vector2I.Zero)
		{
			_lmLayerMask = new Vector2I(~0, ~0);
		}
		else
		{
			_lmLayerMask = Vector2I.Zero;
		}

		foreach (var mat in _materials)
		{
			mat.SetShaderParameter("lightmap_bitmask", _lmLayerMask);
		}
	}

	public void ToggleLmLayer(uint layer)
	{
		if (layer >= 64)
		{
			throw new ArgumentOutOfRangeException(nameof(layer));
		}

		static int ToggleBit(int bitIdx, int val)
		{
			var mask = 1 << bitIdx;
			var isEnabled = (val & mask) != 0;
			if (isEnabled)
			{
				val &= ~mask;
			}
			else
			{
				val |= mask;
			}
			return val;
		}

		if (layer < 32)
		{
			_lmLayerMask.X = ToggleBit((int)layer, _lmLayerMask.X);
		}
		else
		{
			_lmLayerMask.Y = ToggleBit((int)layer - 32, _lmLayerMask.Y);
		}

		foreach (var mat in _materials)
		{
			mat.SetShaderParameter("lightmap_bitmask", _lmLayerMask);
		}
	}

	// TODO: Make this less of a mess
	private ObjectHierarchy BuildHierarchy()
	{
		ObjectHierarchy objHierarchy;
		if (_file.Chunks.TryGetValue("GAM_FILE", out var gamFileChunk))
		{

			var options = new EnumerationOptions { MatchCasing = MatchCasing.CaseInsensitive };
			var name = ((GamFile)gamFileChunk).fileName;
			var paths = Directory.GetFiles(FileName.GetBaseDir(), name, options);
			if (!paths.IsEmpty())
			{
				objHierarchy = new ObjectHierarchy(_file, new DbFile(paths[0]));
			}
			else
			{
				objHierarchy = new ObjectHierarchy(_file);
			}
		}
		else
		{
			objHierarchy = new ObjectHierarchy(_file);
		}
		return objHierarchy;
	}

	private void UseChunk<T>(string name, Action<T> action)
	{
		if (_file.Chunks.TryGetValue(name, out var value))
		{
			action((T)value);
		}
		else
		{
			GD.Print($"No chunk found/loaded: {name}");
		}
	}

	private void PlaceObjects(BrList brList, ObjectHierarchy objHierarchy)
	{
		foreach (var brush in brList.Brushes)
		{
			if (brush.media != BrList.Brush.Media.Object)
			{
				continue;
			}

			var id = (int)brush.brushInfo;
			var modelNameProp = objHierarchy.GetProperty<PropLabel>(id, "P$ModelName");
			var scaleProp = objHierarchy.GetProperty<PropVector>(id, "P$Scale");
			var renderTypeProp = objHierarchy.GetProperty<PropRenderType>(id, "P$RenderTyp");
			var jointPosProp = objHierarchy.GetProperty<PropJointPos>(id, "P$JointPos");
			var txtRepl0 = objHierarchy.GetProperty<PropString>(id, "P$OTxtRepr0");
			var txtRepl1 = objHierarchy.GetProperty<PropString>(id, "P$OTxtRepr1");
			var txtRepl2 = objHierarchy.GetProperty<PropString>(id, "P$OTxtRepr2");
			var txtRepl3 = objHierarchy.GetProperty<PropString>(id, "P$OTxtRepr3");
			var renderAlpha = objHierarchy.GetProperty<PropFloat>(id, "P$RenderAlp");
			var renderMode = renderTypeProp == null ? PropRenderType.Mode.Normal : renderTypeProp.mode;

			if (modelNameProp == null || renderMode == PropRenderType.Mode.NotRendered)
			{
				continue;
			}

			// Let's try and place an object :)
			var modelName = modelNameProp.value + ".bin";
			var pos = brush.position.ToGodotVec3();
			var rot = brush.angle.ToGodotVec3(false);
			var scale = scaleProp == null ? Vector3.One : scaleProp.value.ToGodotVec3(false);
			var meshDetails = Timing.TimeStage("Get Models", () => Context.Instance.ModelLoader.Load(modelName));
			if (meshDetails.Length != 0)
			{
				var model = new Node3D();
				model.Position = pos;
				model.RotationDegrees = rot;
				model.Scale = scale;
				
				var joints = jointPosProp != null ? jointPosProp.Positions : [0, 0, 0, 0, 0, 0];
				var meshes = ModelLoader.TransformMeshes(joints, meshDetails);
				
				bool GetTextReplPath(PropString prop, out string path)
				{
					path = "";
					if (prop == null)
					{
						return false;
					}

					path = prop.value;

					var campaignResources = Context.Instance.CampaignResources;
					if (path.StartsWith("fam", StringComparison.OrdinalIgnoreCase))
					{
						var resType = ResourceType.Texture;
						path = campaignResources.GetResourcePath(resType, prop.value);
					}
					else
					{
						var resType = ResourceType.ObjectTexture;
						var convertedValue = ResourcePathManager.ConvertSeparator(prop.value);
						var resName = Path.GetFileNameWithoutExtension(convertedValue);
						path = campaignResources.GetResourcePath(resType, resName);
					}

					return path != null;
				}

				var repls = new PropString[] { txtRepl0, txtRepl1, txtRepl2, txtRepl3 };
				foreach (var meshInstance in meshes)
				{
					for (var i = 0; i < 4; i++)
					{
						if (GetTextReplPath(repls[i], out var path))
						{
							var overrideMat = new StandardMaterial3D
							{
								AlbedoTexture = TextureLoader.LoadTexture(path),
								Transparency = BaseMaterial3D.TransparencyEnum.AlphaDepthPrePass,
							};

							var surfaceCount = meshInstance.Mesh.GetSurfaceCount();
							for (var idx = 0; idx < surfaceCount; idx++)
							{
								var surfaceMat = meshInstance.Mesh.SurfaceGetMaterial(idx);
								if (surfaceMat.HasMeta($"TxtRepl{i}"))
								{
									meshInstance.SetSurfaceOverrideMaterial(idx, overrideMat);
								}
							}
						}
					}

					if (renderAlpha != null)
					{
						meshInstance.Transparency = 1.0f - renderAlpha.value;
					}
					
					model.AddChild(meshInstance);
				}
				
				model.AddToGroup(OBJECT_MODELS_GROUP);
				
				AddChild(model);
			}
		}
	}

	private void BuildWrMeshes(WorldRep worldRep)
	{
		var cells = worldRep.Cells;
		var lmHdr = worldRep.DataHeader.LightmapFormat == 2;
		GD.Print($"HDR Lightmap: {lmHdr}");

		var packingRects = new List<RectPacker.Rect>();
		var surfaceDataMap = new Dictionary<int, MeshSurfaceData>();
		var rectDataMap = new Dictionary<int, LightmapRectData>();

		for (var cellIdx = 0; cellIdx < cells.Length; cellIdx++)
		{
			var cell = cells[cellIdx];
			var numPolys = cell.PolyCount;
			var numRenderPolys = cell.RenderPolyCount;
			var numPortalPolys = cell.PortalPolyCount;

			// There's nothing to render
			if (numRenderPolys == 0 || numPortalPolys >= numPolys)
			{
				continue;
			}

			// You'd think these would be the same number, but apparently not
			// I think it's because water counts as a render poly and a portal poly
			var maxPolyIdx = Math.Min(numRenderPolys, numPolys - numPortalPolys);
			var cellIdxOffset = 0;
			for (int polyIdx = 0; polyIdx < maxPolyIdx; polyIdx++)
			{
				var lightmapRectData = ProcessCellSurfaceData(surfaceDataMap, worldRep, cellIdx, polyIdx, cellIdxOffset);
				rectDataMap.Add(packingRects.Count, lightmapRectData);

				var light = cell.LightList[polyIdx];
				var rect = new RectPacker.Rect
				{
					Width = light.Width,
					Height = light.Height,
					Id = packingRects.Count,
				};
				packingRects.Add(rect);

				cellIdxOffset += cell.Polys[polyIdx].VertexCount;
			}
		}

		var lightmapTexture = Timing.TimeStage("Build Lightmap", () =>
		{
			return BuildLightmapTexture(cells, packingRects.ToArray(), rectDataMap, surfaceDataMap);
		});

		foreach (var (textureId, surface) in surfaceDataMap)
		{
			if (surface.Empty)
			{
				continue;
			}

			var albedoTexture = _textureLoader.Get(textureId);
			var mat = BuildMaterial(albedoTexture, lightmapTexture, lmHdr, _lmLayerMask);
			_materials.Add(mat);

			var mesh = new ArrayMesh();
			mesh.AddSurfaceFromArrays(Mesh.PrimitiveType.Triangles, surface.BuildSurfaceArray());
			mesh.SurfaceSetMaterial(0, mat);

			var meshInstance = new MeshInstance3D { Mesh = mesh };
			AddChild(meshInstance);
		}
	}

	private LightmapRectData ProcessCellSurfaceData(Dictionary<int, MeshSurfaceData> surfaceDataMap, WorldRep worldRep, int cellIdx, int polyIdx, int indicesOffset)
	{
		var cell = worldRep.Cells[cellIdx];
		var poly = cell.Polys[polyIdx];
		var normal = cell.Planes[poly.PlaneId].Normal.ToGodotVec3();
		var vertices = new List<Vector3>();
		var textureUvs = new List<Vector2>();
		var lightmapUvs = new List<Vector2>();

		var numPolyVertices = poly.VertexCount;
		for (var j = 0; j < numPolyVertices; j++)
		{
			var vertex = cell.Vertices[cell.Indices[indicesOffset + j]];
			vertices.Add(vertex.ToGodotVec3());
		}

		var renderPoly = cell.RenderPolys[polyIdx];
		var light = cell.LightList[polyIdx];
		var lightmapScale = worldRep.DataHeader.LightmapScaleMultiplier();
		var textureId = CalcBaseUV(cell, poly, renderPoly, light, textureUvs, lightmapUvs, lightmapScale, indicesOffset);

		if (!surfaceDataMap.ContainsKey(textureId))
		{
			surfaceDataMap.Add(textureId, new MeshSurfaceData());
		}
		var surfaceData = surfaceDataMap[textureId];
		var (start, end) = surfaceData.AddPolygon(vertices, normal, textureUvs, lightmapUvs);

		if (polyIdx >= cell.Lightmaps.Length) GD.Print("HUH");

		return new LightmapRectData(cellIdx, polyIdx, textureId, start, end);
	}

	private static Texture2DArray BuildLightmapTexture(
		WorldRep.Cell[] cells,
		RectPacker.Rect[] packingRects,
		Dictionary<int, LightmapRectData> rectDataMap,
		Dictionary<int, MeshSurfaceData> surfaceDataMap)
	{
		var bounds = Timing.TimeStage("RectPack", () =>
		{
			return RectPacker.Pack(packingRects);
		});
		GD.Print($"Packed {packingRects.Length} rects in ({bounds.Width}, {bounds.Height})");

		var lightmapFormat = Image.Format.Rgba8;
		var lmLayerCount = 33; // TODO: Use LightmapLayers
		var lmImages = new Godot.Collections.Array<Image>();
		lmImages.Resize(lmLayerCount);
		for (var i = 0; i < lmLayerCount; i++)
		{
			lmImages[i] = Image.CreateEmpty(bounds.Width, bounds.Height, false, lightmapFormat);
		}

		foreach (var rect in packingRects)
		{
			if (!rectDataMap.ContainsKey(rect.Id)) GD.Print("Invalid rectDataMap key");
			var info = rectDataMap[rect.Id];

			if (info.cellIndex >= cells.Length) GD.Print($"CellIndex too big: {info.cellIndex}/{cells.Length}");
			if (info.lightmapIndex >= cells[info.cellIndex].Lightmaps.Length) GD.Print($"LightmapIndex too big: {info.lightmapIndex}/{cells[info.cellIndex].Lightmaps.Length}");

			var lightmap = cells[info.cellIndex].Lightmaps[info.lightmapIndex];
			var width = lightmap.Width;
			var height = lightmap.Height;
			var layerCount = lightmap.Layers;
			var srcRect = new Rect2I(0, 0, width, height);
			var dst = new Vector2I(rect.X, rect.Y);
			Timing.TimeStage("Blit LM Data", () =>
			{
				for (var i = 0; i < layerCount; i++)
				{
					var bytes = Timing.TimeStage("Get LM Data", () => { return lightmap.AsBytesRgba(i); });
					var cellLm = Image.CreateFromData(width, height, false, lightmapFormat, bytes);
					lmImages[i].BlitRect(cellLm, srcRect, dst);
				}
			});

			Timing.TimeStage("Transform LM UVs", () =>
			{
				if (!surfaceDataMap.ContainsKey(info.textureId)) GD.Print("Invalid SurfaceDataMap key");
				surfaceDataMap[info.textureId].TransformUv2s(info.uvStart, info.uvEnd, (uv) =>
				{
					var u = uv.X;
					var v = uv.Y;

					// Clamp uv range to [0..1]
					u %= 1;
					v %= 1;
					if (u < 0) u = Math.Abs(u);
					if (v < 0) v = Math.Abs(v);

					// Transform!
					u = (rect.X + rect.Width * u) / bounds.Width;
					v = (rect.Y + rect.Height * v) / bounds.Height;
					return new Vector2(u, v);
				});
			});
		}

		for (var i = 0; i < lmLayerCount; i++)
		{
			lmImages[i].GenerateMipmaps();
		}

		var lightmapTexture = new Texture2DArray();
		lightmapTexture.CreateFromImages(lmImages);
		return lightmapTexture;
	}

	private int CalcBaseUV(
		WorldRep.Cell cell,
		WorldRep.Cell.Poly poly,
		WorldRep.Cell.RenderPoly renderPoly,
		WorldRep.Cell.LightmapInfo light,
		List<Vector2> textureUvs,
		List<Vector2> lightmapUvs,
		float lightmapScale,
		int cellIdxOffset)
	{
		// TODO: This is slightly hardcoded for ND. Check other stuff at some point. Should be handled in LG side imo
		// TODO: This is a mess lol
		var textureId = renderPoly.TextureId;

		var texture = Timing.TimeStage("Load Textures", () => { return _textureLoader.Get(textureId); });
		var texU = renderPoly.TextureVectors.Item1.ToGodotVec3();
		var texV = renderPoly.TextureVectors.Item2.ToGodotVec3();
		var baseU = renderPoly.TextureBases.Item1;
		var baseV = renderPoly.TextureBases.Item2;

		var txUScale = 64.0f / lightmapScale / texture.GetWidth();
		var txVScale = 64.0f / lightmapScale / texture.GetHeight();
		var lmUScale = 4.0f / light.Width;
		var lmVScale = 4.0f / light.Height;

		var txUBase = baseU * txUScale;
		var txVBase = baseV * txVScale;
		var lmUBase = lmUScale * (baseU + (0.5f - light.Bases.Item1) / 4.0f);
		var lmVBase = lmVScale * (baseV + (0.5f - light.Bases.Item2) / 4.0f);

		var uu = texU.Dot(texU);
		var vv = texV.Dot(texV);
		var uv = texU.Dot(texV);
		var anchor = cell.Vertices[cell.Indices[cellIdxOffset + 0]].ToGodotVec3(); // TODO: This probably shouldn't be hardcoded idx 0
		if (uv == 0.0)
		{
			var txUVec = texU * txUScale / uu;
			var txVVec = texV * txVScale / vv;
			var lmUVec = texU * lmUScale / uu;
			var lmVVec = texV * lmVScale / vv;
			for (var i = 0; i < poly.VertexCount; i++)
			{
				var v = cell.Vertices[cell.Indices[cellIdxOffset + i]].ToGodotVec3();
				var delta = new Vector3(v.X - anchor.X, v.Y - anchor.Y, v.Z - anchor.Z);
				var txUV = new Vector2(delta.Dot(txUVec) + txUBase, delta.Dot(txVVec) + txVBase);
				var lmUV = new Vector2(delta.Dot(lmUVec) + lmUBase, delta.Dot(lmVVec) + lmVBase);
				textureUvs.Add(txUV);
				lightmapUvs.Add(lmUV);
			}
		}
		else
		{
			var denom = 1.0f / (uu * vv - uv * uv);

			var txUu = uu * txVScale * denom;
			var txVv = vv * txUScale * denom;
			var txUvu = txUScale * denom * uv;
			var txUvv = txVScale * denom * uv;
			var lmUu = uu * lmVScale * denom;
			var lmVv = vv * lmUScale * denom;
			var lmUvu = lmUScale * denom * uv;
			var lmUvv = lmVScale * denom * uv;
			for (var i = 0; i < poly.VertexCount; i++)
			{
				var v = cell.Vertices[cell.Indices[cellIdxOffset + i]].ToGodotVec3();
				var delta = new Vector3(v.X - anchor.X, v.Y - anchor.Y, v.Z - anchor.Z);
				var du = delta.Dot(texU);
				var dv = delta.Dot(texV);

				var txUV = new Vector2(txUBase + txVv * du - txUvu * dv, txVBase + txUu * dv - txUvv * du);
				var lmUV = new Vector2(lmUBase + lmVv * du - lmUvu * dv, lmVBase + lmUu * dv - lmUvv * du);
				textureUvs.Add(txUV);
				lightmapUvs.Add(lmUV);
			}
		}

		return textureId;
	}

	private void RegisterTextures(TxList textureList)
	{
		// TODO: Use PathJoin
		var count = textureList.ItemCount;
		for (var i = 0; i < count; i++)
		{
			var item = textureList.Items[i];
			var path = "";
			for (var j = 0; j < item.Tokens.Length; j++)
			{
				var token = item.Tokens[j];
				if (token == 0)
				{
					break;
				}

				path += $"{textureList.Tokens[token - 1]}/";
			}
			path += item.Name;

			if (!_textureLoader.Register(i, path))
			{
				GD.Print($"Failed to register texture: {path}");
			}
		}

		if (Dump) DumpTextureList(textureList);
	}

	private static void DumpTextureList(TxList textureList)
	{
		GD.Print($"TXLIST:\n  BlockSize: {textureList.BlockSize}\n  ItemCount: {textureList.ItemCount}\n  TokenCount: {textureList.TokenCount}\n  Tokens:");
		for (var i = 0; i < textureList.TokenCount; i++)
		{
			GD.Print($"    {i}: {textureList.Tokens[i]}");
		}
		GD.Print($"  Items:");
		for (var i = 0; i < textureList.ItemCount; i++)
		{
			var item = textureList.Items[i];
			GD.Print($"    {i}:\n      Tokens: [{item.Tokens[0]}, {item.Tokens[1]}, {item.Tokens[2]}, {item.Tokens[3]}]\n      Name: {item.Name}");
		}
	}

	const string MATERIAL_PATH = "res://project/materials/base.tres";
	private static ShaderMaterial BuildMaterial(
		Texture2D albedoTexture,
		Texture2DArray lightmapTexture,
		bool lmHdr,
		Vector2I layerMask)
	{
		var material = ResourceLoader.Load<ShaderMaterial>(MATERIAL_PATH).Duplicate() as ShaderMaterial;
		material.SetShaderParameter("texture_albedo", albedoTexture);
		material.SetShaderParameter("lightmap_albedo", lightmapTexture);
		material.SetShaderParameter("lightmap_2x", lmHdr);
		material.SetShaderParameter("lightmap_bitmask", layerMask);
		return material;
	}
}