using System;
using System.Collections.Generic;
using Sandbox;
using Sandbox.Rendering;

namespace RedSnail.WaterTool2;

/// <summary>
/// A lake's surface drawn the way the sea's is: rings of one grid round the camera, each twice as coarse as the one
/// inside it, rather than one mesh built once for the whole lake.
/// <para>
/// A built mesh has to pick its detail before anyone looks at it, so a lake a kilometre across either costs millions of
/// points or carries waves too coarse to see. Rings put the fine water wherever the camera is and coarsen with
/// distance, at a fixed cost whatever the lake's size.
/// </para>
/// <para>
/// The rings' cell grid is anchored in the world, not to the camera - each ring's centre is snapped to its own cell
/// size - so a mesh built on that grid lines up with it however the camera moves. That is what lets a river's mouth
/// stay an exact built mesh while the lake around it is drawn by rings.
/// </para>
/// </summary>
public sealed partial class WaterNetwork
{
	private const int MaxLakeRings = 10;

	/// <summary>
	/// Draws lakes in rings round the camera instead of as one built mesh. Where a river comes in is still built, so
	/// the two meet vertex for vertex; the rings draw the rest.
	/// </summary>
	[Property(Title = "Ring Lakes"), Group("Lakes"), Order(20)]
	public bool LakeRings { get; set; } = false;

	/// <summary>
	/// How many rings a lake is drawn in: each reaches twice as far as the one inside it.
	/// </summary>
	[Property(Title = "Lake Rings"), Group("Lakes"), Order(21), Range(2, 10)]
	public int LakeRingCount { get; set { field = value.Clamp(1, MaxLakeRings); } } = 6;

	/// <summary>
	/// How far out from a river mouth a lake's built mesh keeps before the rings take over. Wide enough to hold the
	/// river's rows and meet them vertex for vertex, with room to fade into the rings' surface.
	/// <para>
	/// This is the only handover in open water, so it is the only one that has to look right. The rest of the lake is
	/// the rings' and they stop at its outline, which is under the terrain.
	/// </para>
	/// </summary>
	/// <para>
	/// Keep it as small as the join allows. The cap round each mouth is meshed at the full cell size, so its cost grows
	/// with the square of this: at a cell size of 8, a band of 1024 is tens of thousands of points per river.
	/// </para>
	[Property(Title = "Lake Mouth Band"), Group("Lakes"), Order(22), Range(64, 4096)]
	public float LakeBandWidth { get; set { field = MathF.Max(value, 1.0f); } } = 256.0f;

	/// <summary>
	/// One lake's rings this frame: where they are, and the attributes each is drawn with.
	/// </summary>
	private sealed class LakeRingSet
	{
		public GpuBuffer<Vector4> Buffer;
		public readonly Vector4[] Rings = new Vector4[MaxLakeRings];
		public RenderAttributes[] Attributes = Array.Empty<RenderAttributes>();
		public int Count;

		/// <summary>
		/// Which of them have any of this lake to draw. The rings are placed round the camera, not round the lake, so
		/// most of a lake's rings usually reach none of it.
		/// </summary>
		public readonly bool[] Draw = new bool[MaxLakeRings];

		/// <summary>
		/// The lake's still surface. The built mesh needs it too: it is the plane both it and the rings refer the mask
		/// to, so the two agree on where one ends and the other begins.
		/// </summary>
		public float Height;

		public void Release()
		{
			Buffer?.Dispose();
			Buffer = null;
			Count = 0;
		}
	}

	private readonly Dictionary<int, LakeRingSet> m_LakeRings = new();

	/// <summary>
	/// Which of each lake's water its rings may draw, and the same on the GPU. The built mesh throws away the pixels
	/// the rings own and the rings throw away the rest, so the two cover the lake between them without either being cut
	/// to the other's shape.
	/// </summary>
	private readonly Dictionary<int, LakeMask> m_LakeMasks = new();
	private readonly Dictionary<int, GpuBuffer<Vector2>> m_LakeMaskBuffers = new();

	/// <summary>
	/// A lake's built tiles are drawn with these rather than the shared set, so each carries its own mask.
	/// </summary>
	private readonly Dictionary<int, RenderAttributes> m_LakeMeshAttributes = new();



	/// <summary>
	/// Takes a lake's mask from its build, or drops it when the lake could not be built.
	/// </summary>
	private void SetLakeMask(int _LakeId, LakeMask _Mask)
	{
		SetLakeShape(_LakeId, _Mask);

		if (m_LakeMaskBuffers.Remove(_LakeId, out var old))
			old?.Dispose();

		if (_Mask is null)
		{
			m_LakeMasks.Remove(_LakeId);
			return;
		}

		m_LakeMasks[_LakeId] = _Mask;

		// A dedicated server keeps the mask for gameplay to read, but has no GPU to put it on and nothing to draw it with.
		if (Application.IsHeadless)
			return;

		var buffer = new GpuBuffer<Vector2>(_Mask.Cells.Length, GpuBuffer.UsageFlags.Structured);
		buffer.SetData(_Mask.Cells);

		m_LakeMaskBuffers[_LakeId] = buffer;
	}



	/// <summary>
	/// Puts a lake's mask on a set of attributes. Without one the whole lake is the built mesh's, which is what should
	/// happen while a lake is still building or could not be built at all.
	/// </summary>
	private void ApplyLakeMask(RenderAttributes _Attributes, int _LakeId)
	{
		if (!m_LakeMasks.TryGetValue(_LakeId, out var mask) ||
			!m_LakeMaskBuffers.TryGetValue(_LakeId, out var buffer) || !buffer.IsValid())
		{
			_Attributes.Set("LakeMaskWidth", 0);
			_Attributes.Set("LakeMaskHeight", 0);
			return;
		}

		_Attributes.Set("LakeMaskWidth", mask.Width);
		_Attributes.Set("LakeMaskHeight", mask.Height);
		_Attributes.Set("LakeMaskOrigin", mask.Origin);
		_Attributes.Set("LakeMaskCell", mask.CellSize);
		_Attributes.Set("LakeMask", buffer);

		// How far round a mouth the built mesh keeps, and over how much of that it fades its surface into the rings'.
		_Attributes.Set("LakeMaskBand", LakeBandWidth);
		_Attributes.Set("LakeMaskBlend", LakeBandWidth * 0.5f);
	}



	/// <summary>
	/// The attributes a lake's built tiles are drawn with: the shared set plus that lake's own mask.
	/// </summary>
	private RenderAttributes LakeMeshAttributes(int _LakeId)
	{
		if (!m_LakeMeshAttributes.TryGetValue(_LakeId, out var attributes))
		{
			attributes = new RenderAttributes();
			m_LakeMeshAttributes[_LakeId] = attributes;
		}

		ApplySharedAttributes(attributes);

		attributes.Set("OceanRing", false);
		attributes.Set("LakeRing", false);

		// Only the rings give the built mesh anything to give up; without them it draws the whole lake as it always has.
		if (LakeRings)
		{
			ApplyLakeMask(attributes, _LakeId);
			ApplyLakeRingGrid(attributes, _LakeId);
		}
		else
		{
			attributes.Set("LakeMaskWidth", 0);
			attributes.Set("LakeRingCount", 0);
		}

		return attributes;
	}



	/// <summary>
	/// Puts a lake's rings on its built mesh's attributes. The built mesh never draws as a ring, but it has to be able
	/// to ask what the rings draw at a point, so it can fade its own surface into theirs before handing over.
	/// <para>
	/// Without this it asks an unbound buffer and is handed a ring at the world origin one unit across, whose surface
	/// is all but the wave itself - so it fades from the wave into the wave, and the handover it was meant to close
	/// stays exactly as open as it was.
	/// </para>
	/// </summary>
	private void ApplyLakeRingGrid(RenderAttributes _Attributes, int _LakeId)
	{
		if (!m_LakeRings.TryGetValue(_LakeId, out var set) || set.Count <= 0 || !set.Buffer.IsValid())
		{
			_Attributes.Set("LakeRingCount", 0);
			return;
		}

		_Attributes.Set("LakeRingCount", set.Count);
		_Attributes.Set("LakeRings", set.Buffer);

		// The lake's still surface, which the built mesh needs as the plane both it and the rings refer the mask to.
		_Attributes.Set("LakeRingHeight", set.Height);
	}


	/// <summary>
	/// Places every lake's rings round this frame's camera. Each ring's centre is snapped to twice its cell size, so
	/// the ring inside it lands exactly on its grid and a mesh built on that grid lines up with both.
	/// </summary>
	private void UpdateLakeRings(Vector3 _Viewer)
	{
		if (!LakeRings)
		{
			ReleaseLakeRings();
			return;
		}

		// The rings are drawn from the grid the sea builds, which it only builds when there is a sea to draw.
		EnsureRingMesh();

		int count = Math.Clamp(LakeRingCount, 1, MaxLakeRings);
		float half = SeaRingCells * 0.5f;

		foreach (var lake in Lakes)
		{
			if (lake is null || !lake.IsValid || lake.IsCoast)
				continue;

			if (!m_LakeRings.TryGetValue(lake.Id, out var set))
			{
				set = new LakeRingSet();
				m_LakeRings[lake.Id] = set;
			}

			// Each ring's centre is snapped to its own cell size, so on most frames every one of them lands exactly
			// where it was and there is nothing to send. Only an actual move is uploaded.
			bool moved = !set.Buffer.IsValid() || set.Count != count;

			set.Count = count;

			// Where the lake is, so the rings that reach none of it are never recorded. Without its mask there is
			// nothing to test against, and every ring is drawn as before.
			bool bounded = m_LakeMasks.TryGetValue(lake.Id, out var mask) && mask is not null;

			Vector2 mins = bounded ? mask.Origin : Vector2.Zero;
			Vector2 maxs = bounded ? mask.Origin + new Vector2(mask.Width * mask.CellSize, mask.Height * mask.CellSize) : Vector2.Zero;

			for (int k = 0; k < count; k++)
			{
				float cell = CellSize * (1 << k);
				float snap = cell * 2.0f;

				var ring = new Vector4(
					MathF.Round(_Viewer.x / snap) * snap,
					MathF.Round(_Viewer.y / snap) * snap,
					cell,
					half * cell);

				if (set.Rings[k] != ring)
				{
					set.Rings[k] = ring;
					moved = true;
				}

				set.Draw[k] = !bounded || RingReachesLake(set.Rings[k], set.Rings[Math.Max(k - 1, 0)], k > 0, mins, maxs);
			}

			if (!set.Buffer.IsValid())
				set.Buffer = new GpuBuffer<Vector4>(MaxLakeRings, GpuBuffer.UsageFlags.Structured);

			if (moved)
				set.Buffer.SetData(set.Rings.AsSpan(0, MaxLakeRings));

			UpdateLakeRingAttributes(lake, set);
		}

		// Lakes that have gone take their rings with them.
		if (m_LakeRings.Count <= Lakes.Count)
			return;

		var stale = new List<int>();

		foreach (var (lakeId, _) in m_LakeRings)
		{
			if (FindLake(lakeId) is null)
				stale.Add(lakeId);
		}

		foreach (int lakeId in stale)
		{
			if (m_LakeRings.Remove(lakeId, out var set))
				set.Release();
		}
	}



	/// <summary>
	/// Whether a ring has any of a lake to draw: its square has to reach the lake at all, and has to reach some of it
	/// the finer ring inside it is not already covering - that part is a hole in this one.
	/// </summary>
	private static bool RingReachesLake(Vector4 _Ring, Vector4 _Inner, bool _HasInner, Vector2 _Mins, Vector2 _Maxs)
	{
		float lowX = MathF.Max(_Mins.x, _Ring.x - _Ring.w);
		float lowY = MathF.Max(_Mins.y, _Ring.y - _Ring.w);
		float highX = MathF.Min(_Maxs.x, _Ring.x + _Ring.w);
		float highY = MathF.Min(_Maxs.y, _Ring.y + _Ring.w);

		if (lowX > highX || lowY > highY)
			return false;

		if (!_HasInner)
			return true;

		// What is left of the lake within this ring, against the hole the finer one punches in it.
		return lowX < _Inner.x - _Inner.w || highX > _Inner.x + _Inner.w ||
			   lowY < _Inner.y - _Inner.w || highY > _Inner.y + _Inner.w;
	}



	private void UpdateLakeRingAttributes(WaterLake _Lake, LakeRingSet _Set)
	{
		if (_Set.Attributes.Length < _Set.Count)
		{
			var grown = new RenderAttributes[_Set.Count];

			for (int i = 0; i < grown.Length; i++)
				grown[i] = i < _Set.Attributes.Length ? _Set.Attributes[i] : new RenderAttributes();

			_Set.Attributes = grown;
		}

		m_LakeJetSlots.TryGetValue(_Lake.Id, out int slot);

		_Set.Height = _Lake.Height;

		for (int k = 0; k < _Set.Count; k++)
		{
			var attributes = _Set.Attributes[k];

			ApplySharedAttributes(attributes);

			attributes.Set("OceanRing", false);
			attributes.Set("LakeRing", true);
			attributes.Set("LakeRingIndex", k);
			attributes.Set("LakeRingCount", _Set.Count);
			attributes.Set("LakeRings", _Set.Buffer);
			attributes.Set("LakeRingHeight", _Lake.Height);
			attributes.Set("LakeRingWaves", _Lake.Waves);

			// The same slot the lake's built vertices carry, so the rings are stirred by the same jets.
			attributes.Set("LakeRingSlot", slot + 1.0f);

			// Nothing outside what the outline encloses, and nothing in the disc round a river mouth the built mesh
			// keeps.
			ApplyLakeMask(attributes, _Lake.Id);
		}
	}






	/// <summary>
	/// Records every lake's rings, finest first, the way the sea's are.
	/// </summary>
	private void DrawLakeRings(CommandList _CommandList, Material _Material)
	{
		if (!LakeRings || !m_RingVertices.IsValid() || !m_RingIndices.IsValid())
			return;

		foreach (var (_, set) in m_LakeRings)
		{
			for (int k = 0; k < set.Count; k++)
			{
				if (!set.Draw[k])
					continue;

				_CommandList.DrawIndexed(m_RingVertices, m_RingIndices, _Material, 0, m_RingIndexCount, set.Attributes[k]);
			}
		}
	}



	private void ReleaseLakeRings()
	{
		foreach (var (_, set) in m_LakeRings)
			set.Release();

		m_LakeRings.Clear();

		foreach (var (_, buffer) in m_LakeMaskBuffers)
			buffer?.Dispose();

		m_LakeMaskBuffers.Clear();
		m_LakeMasks.Clear();
		m_LakeMeshAttributes.Clear();
	}



	// ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────
	// Lake shapes
	//
	// Every lake's mask and the height of its water, packed together for another system's shaders - one that has to
	// keep out from under the water, as snow does. The water writes no depth, so nothing drawn before or after it can
	// tell what it covers; this says, from the same field the water is drawn to.
	//
	// Kept apart from the rings' masks, which go whenever the rings are switched off: a lake is the same shape either way.
	// ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────

	private readonly Dictionary<int, LakeMask> m_LakeShapes = new();
	private readonly Dictionary<int, float> m_LakeShapeHeights = new();
	private GpuBuffer<Vector4> m_LakeShapeRows;
	private GpuBuffer<float> m_LakeShapeShores;
	private int m_LakeShapeCount;
	private bool m_LakeShapesDirty = true;



	/// <summary>
	/// How many lakes <see cref="LakeShapeRows"/> describes.
	/// </summary>
	public int LakeShapeCount
	{
		get
		{
			RefreshLakeShapes();
			return m_LakeShapeCount;
		}
	}



	/// <summary>
	/// Every lake's shape, for another system's shaders: two rows a lake - (the origin of its cells x and y, the size of a
	/// cell, the height of its water at rest), then (cells across, cells down, where its cells start in
	/// <see cref="LakeShapeShores"/>, 0). Null while there are no lakes, and on a dedicated server.
	/// </summary>
	public GpuBuffer<Vector4> LakeShapeRows
	{
		get
		{
			RefreshLakeShapes();
			return m_LakeShapeRows;
		}
	}



	/// <summary>
	/// The lakes' cells, one lake after another, row after row: how far each cell's centre is inside its lake's shore, in
	/// world units, negative outside it. Read between the cells, at (point - origin) / cell size - 0.5, its zero is the
	/// shore the water is drawn to, as <see cref="LakeMask.ShoreDistance"/> reads it. Null while there are no lakes, and on a
	/// dedicated server.
	/// </summary>
	public GpuBuffer<float> LakeShapeShores
	{
		get
		{
			RefreshLakeShapes();
			return m_LakeShapeShores;
		}
	}



	/// <summary>
	/// Takes a lake's shape from its build, or drops it when the lake could not be built.
	/// </summary>
	private void SetLakeShape(int _LakeId, LakeMask _Mask)
	{
		if (_Mask is null)
			m_LakeShapes.Remove(_LakeId);
		else
			m_LakeShapes[_LakeId] = _Mask;

		m_LakeShapesDirty = true;
	}



	/// <summary>
	/// Packs the shapes again when a lake has been built or dropped since, or moved up or down: its height is no part of
	/// its shape, and changes without a build.
	/// </summary>
	private void RefreshLakeShapes()
	{
		// The shapes are only for shaders: a dedicated server has none, nor a GPU to put them on.
		if (Application.IsHeadless)
			return;

		foreach (var (lakeId, height) in m_LakeShapeHeights)
		{
			var lake = FindLake(lakeId);

			if (lake is null || lake.Height != height)
			{
				m_LakeShapesDirty = true;
				break;
			}
		}

		if (!m_LakeShapesDirty)
			return;

		m_LakeShapesDirty = false;
		m_LakeShapeHeights.Clear();

		var rows = new List<Vector4>();
		var shores = new List<float>();

		foreach (var (lakeId, mask) in m_LakeShapes)
		{
			var lake = FindLake(lakeId);

			// The sea's water round the land is the coast zone's, not a lake's.
			if (lake is null || !lake.IsValid || lake.IsCoast)
				continue;

			rows.Add(new Vector4(mask.Origin.x, mask.Origin.y, mask.CellSize, lake.Height));
			rows.Add(new Vector4(mask.Width, mask.Height, shores.Count, 0.0f));

			foreach (var cell in mask.Cells)
				shores.Add(cell.x);

			m_LakeShapeHeights[lakeId] = lake.Height;
		}

		ReleaseLakeShapeBuffers();

		m_LakeShapeCount = rows.Count / 2;

		if (m_LakeShapeCount == 0)
			return;

		m_LakeShapeRows = new GpuBuffer<Vector4>(rows.Count, GpuBuffer.UsageFlags.Structured);
		m_LakeShapeRows.SetData(rows.ToArray());

		m_LakeShapeShores = new GpuBuffer<float>(shores.Count, GpuBuffer.UsageFlags.Structured);
		m_LakeShapeShores.SetData(shores.ToArray());
	}



	private void ReleaseLakeShapeBuffers()
	{
		m_LakeShapeRows?.Dispose();
		m_LakeShapeRows = null;

		m_LakeShapeShores?.Dispose();
		m_LakeShapeShores = null;

		m_LakeShapeCount = 0;
	}



	/// <summary>
	/// Gives up the shapes' buffers. The shapes themselves stay, to be packed again the next time they are asked for.
	/// </summary>
	private void ReleaseLakeShapes()
	{
		ReleaseLakeShapeBuffers();
		m_LakeShapesDirty = true;
	}
}