// Copyright (c) 2026 SubZero Studios LLC. All rights reserved.

using System;
using Sandbox;
using Editor.TerrainEditor;

namespace Editor.SubTerrain;

public enum SubTerrainPaintLayer
{
	Base = 0,
	Overlay = 1,
	Path = 2
}

internal static class SubTerrainSplatUtil
{
	public static bool TryTexel( Terrain terrain, Vector3 hitLocal, out int x, out int y )
	{
		x = 0;
		y = 0;
		var storage = terrain?.Storage;
		if ( storage == null || storage.Resolution <= 0 || storage.TerrainSize <= 0.001f )
			return false;

		var uv = new Vector2( hitLocal.x, hitLocal.y ) / storage.TerrainSize;
		x = (int)Math.Floor( storage.Resolution * uv.x );
		y = (int)Math.Floor( storage.Resolution * uv.y );
		if ( x < 0 || y < 0 || x >= storage.Resolution || y >= storage.Resolution )
			return false;

		return storage.ControlMap != null && storage.ControlMap.Length == storage.Resolution * storage.Resolution;
	}

	public static CompactTerrainMaterial Read( Terrain terrain, int x, int y )
	{
		return new CompactTerrainMaterial( terrain.Storage.ControlMap[x + y * terrain.Storage.Resolution] );
	}

	public static string MaterialName( Terrain terrain, int index )
	{
		var materials = terrain?.Storage?.Materials;
		if ( materials == null || index < 0 || index >= materials.Count )
			return "empty";

		var material = materials[index];
		if ( material is null )
			return "empty";

		var name = material.ResourceName ?? material.ResourcePath ?? "tmat";
		var slash = Math.Max( name.LastIndexOf( '/' ), name.LastIndexOf( '\\' ) );
		if ( slash >= 0 && slash < name.Length - 1 )
			name = name[(slash + 1)..];

		if ( name.EndsWith( ".tmat", StringComparison.OrdinalIgnoreCase ) )
			name = name[..^5];

		return name.Replace( '_', ' ' );
	}

	public static int FindPathGround( Terrain terrain, SubTerrainPaintSettings settings )
	{
		var materials = terrain?.Storage?.Materials;

		if ( settings != null && settings.HasPathGround )
		{
			if ( materials != null && settings.PathGroundIndex < materials.Count )
				return settings.PathGroundIndex;
		}

		if ( settings != null && settings.HasSampledGround )
		{
			if ( materials != null && settings.SampledGroundIndex < materials.Count )
				return settings.SampledGroundIndex;
		}

		if ( materials == null || materials.Count == 0 )
			return 0;

		return 0;
	}

	/// <summary>
	/// Path wear/path tmat when the current pick is invalid (same as ground).
	/// Prefers the next slot after ground; no material-name heuristics.
	/// </summary>
	public static int FindPathWear( Terrain terrain, int pathGround )
	{
		var materials = terrain?.Storage?.Materials;
		if ( materials == null || materials.Count == 0 )
			return pathGround;

		for ( int i = 0; i < materials.Count; i++ )
		{
			if ( i != pathGround )
				return i;
		}

		return pathGround;
	}

	public static float MixEdge( float mixCap )
	{
		return Math.Clamp( mixCap * 0.25f, 0.04f, mixCap );
	}

	/// <summary>
	/// Same rise/run as cs_subterrain_splat PassesGates so the footer skip matches the stroke.
	/// </summary>
	public static bool TrySlopeDegrees( Terrain terrain, int x, int y, out float degrees )
	{
		degrees = 0f;
		var storage = terrain?.Storage;
		var map = storage?.HeightMap;
		if ( storage == null || map == null || storage.Resolution <= 0 )
			return false;

		int res = storage.Resolution;
		int i = x + y * res;
		if ( i < 0 || i >= map.Length )
			return false;

		int x1 = Math.Min( x + 1, res - 1 );
		int y1 = Math.Min( y + 1, res - 1 );
		float height = storage.TerrainHeight;
		float h0 = map[i] / 65535f * height;
		float hx = map[x1 + y * res] / 65535f * height;
		float hy = map[x + y1 * res] / 65535f * height;
		float dx = hx - h0;
		float dy = hy - h0;
		float rise = MathF.Sqrt( dx * dx + dy * dy );
		float texelWorld = storage.TerrainSize / res;
		degrees = MathF.Atan2( rise, Math.Max( texelWorld, 0.001f ) ) * (180f / MathF.PI);
		return true;
	}

	public static bool HasSeam( Terrain terrain, int x, int y )
	{
		var center = Read( terrain, x, y );
		int res = terrain.Storage.Resolution;
		for ( int oy = -1; oy <= 1; oy++ )
		{
			for ( int ox = -1; ox <= 1; ox++ )
			{
				if ( ox == 0 && oy == 0 )
					continue;

				int nx = x + ox;
				int ny = y + oy;
				if ( nx < 0 || ny < 0 || nx >= res || ny >= res )
					continue;

				var n = Read( terrain, nx, ny );
				if ( n.BaseTextureId != center.BaseTextureId || n.OverlayTextureId != center.OverlayTextureId )
					return true;
			}
		}

		return false;
	}
}

[Title( "Paint Texture" )]
[Description( "Paint ground, wear, or a path. Alt samples. Shift smooths mix." )]
[Icon( "brush" )]
[Alias( "subterrain_paint" )]
[Group( "1" )]
[Order( 4 )]
public class SubTerrainPaintTool : EditorTool
{
	readonly SubTerrainTool _parent;
	bool _dragging;
	bool _sampling;
	RectInt _dirtyRegion;
	Vector3 _lastHitWorldPos;
	Vector3 _lastPaintLocal;
	Vector2? _cursorLockPosition;

	public SubTerrainPaintTool( SubTerrainTool parent )
	{
		_parent = parent;
	}

	public static SubTerrainPaintLayer ActiveLayer
	{
		get => _activeLayer;
		set => _activeLayer = value;
	}

	static SubTerrainPaintLayer _activeLayer = SubTerrainPaintLayer.Base;

	public override void OnEnabled()
	{
		AllowGameObjectSelection = false;
		_parent.SetBrushMode( nameof( SubTerrainPaintTool ) );
		_parent.PaintArmed = true;
		SubTerrainPaintStore.Apply( _parent.PaintSettings, SubTerrainPaintLayer.Base );
		ActiveLayer = SubTerrainPaintLayer.Base;
		_parent.TrySetPaintLayer( SubTerrainPaintLayer.Base );
		_parent.RefreshPaintChrome();
	}

	public override void OnDisabled()
	{
		_parent.PaintArmed = false;
		_parent.RefreshPaintChrome();
		_parent.DetachSidebarForReuse();
	}

	public override void OnUpdate()
	{
		var terrain = GetSelectedComponent<Terrain>();
		if ( !terrain.IsValid() )
		{
			_parent.ClearStatus();
			return;
		}

		if ( !terrain.RayIntersects( Gizmo.CurrentRay, Gizmo.RayDepth, out var hitPosition ) )
		{
			_parent.ClearStatus();
			return;
		}

		var tx = terrain.WorldTransform;
		_parent.UpdateStatus( terrain, hitPosition );

		if ( Application.MouseButtons.HasFlag( MouseButtons.Middle ) )
		{
			_cursorLockPosition ??= Application.UnscaledCursorPosition;
			var d = Application.UnscaledCursorPosition - _cursorLockPosition.Value;
			var brush = _parent.BrushSettings;

			if ( Gizmo.IsShiftPressed )
			{
				brush.Size = SubTerrainBrushSize.Clamp( (int)(brush.Size + d.x * 0.25f) );
				brush.Opacity = Math.Clamp( brush.Opacity - d.y * 0.002f, 0f, 1f );
				_parent.SyncBrushSizeWidget();
			}
			else if ( Gizmo.IsCtrlPressed && !brush.RandomRotation )
			{
				brush.Rotation = ((brush.Rotation + d.x * 0.5f) % 360f + 360f) % 360f;
				_parent.SyncBrushRotationRows();
			}

			Application.UnscaledCursorPosition = _cursorLockPosition.Value;
			SceneOverlay.Parent.Cursor = CursorShape.Blank;
			_parent.DrawBrushPreview( new Transform( _lastHitWorldPos, tx.Rotation ), terrain );
			return;
		}

		if ( _cursorLockPosition.HasValue )
		{
			SceneOverlay.Parent.Cursor = CursorShape.None;
			_cursorLockPosition = null;
		}

		_lastHitWorldPos = tx.PointToWorld( hitPosition );
		_parent.DrawBrushPreview( new Transform( _lastHitWorldPos, tx.Rotation ), terrain );

		bool alt = Application.KeyboardModifiers.HasFlag( KeyboardModifiers.Alt );
		bool left = Gizmo.IsLeftMouseDown;

		if ( left && alt )
		{
			if ( !_sampling )
			{
				_sampling = true;
				Sample( terrain, hitPosition );
			}

			return;
		}

		if ( _sampling )
		{
			if ( !left )
				_sampling = false;
			return;
		}

		if ( left )
		{
			bool shouldPaint = !_dragging || !Application.CursorDelta.IsNearZeroLength;

			if ( !_dragging )
			{
				if ( _parent.BrushSettings.RandomRotation )
					_parent.BrushSettings.Rotation = Random.Shared.NextSingle() * 360f;

				_dragging = true;

				var uv = new Vector2( hitPosition.x, hitPosition.y ) / terrain.Storage.TerrainSize;
				var x = (int)Math.Floor( terrain.Storage.Resolution * uv.x );
				var y = (int)Math.Floor( terrain.Storage.Resolution * uv.y );
				_dirtyRegion = new( new Vector2Int( x, y ) );
				_lastPaintLocal = hitPosition;
				OnPaint( terrain, hitPosition );
			}
			else if ( shouldPaint )
			{
				StrokePaint( terrain, hitPosition );
			}
		}
		else if ( _dragging )
		{
			_dragging = false;
			OnPaintEnded( terrain );
		}
	}

	/// <summary>
	/// Fill stamps along the drag in world space. Pack center Path uses a dense step so
	/// the solid core continues; apron-only Path can step wider.
	/// </summary>
	void StrokePaint( Terrain terrain, Vector3 hitLocal )
	{
		var from = _lastPaintLocal;
		var to = hitLocal;
		float dist = (to - from).Length;
		bool path = _parent.PaintSettings.Layer == SubTerrainPaintLayer.Path;
		bool pack = path && _parent.PaintSettings.PackCenter;
		// Pack core is only the soft-brush center. Step must stay under that diameter or
		// you get apron streaks with a solid disc only at the stroke tip.
		float stepScale = pack ? 0.18f : path ? 0.35f : 0.2f;
		float step = Math.Max( _parent.BrushSettings.Size * stepScale, 8f );

		if ( dist < 0.001f )
			return;

		if ( path && !_parent.BrushSettings.RandomRotation )
			OrientPathBrush( from, to );

		if ( dist <= step )
		{
			OnPaint( terrain, to );
			_lastPaintLocal = to;
			return;
		}

		const int maxStamps = 64;
		int stamps = (int)Math.Ceiling( dist / step );
		if ( stamps > maxStamps )
		{
			stamps = maxStamps;
			step = dist / stamps;
		}

		for ( int i = 1; i <= stamps; i++ )
		{
			float t = i / (float)stamps;
			var p = Vector3.Lerp( from, to, t );
			if ( path && _parent.BrushSettings.RandomRotation )
				_parent.BrushSettings.Rotation = Random.Shared.NextSingle() * 360f;

			OnPaint( terrain, p );
		}

		_lastPaintLocal = to;
	}

	void OrientPathBrush( Vector3 from, Vector3 to )
	{
		var d = to - from;
		float lenSq = d.x * d.x + d.y * d.y;
		if ( lenSq < 0.0001f )
			return;

		float deg = MathF.Atan2( d.y, d.x ) * (180f / MathF.PI);
		_parent.BrushSettings.Rotation = ((deg % 360f) + 360f) % 360f;
	}

	void Sample( Terrain terrain, Vector3 hitPosition )
	{
		if ( !SubTerrainSplatUtil.TryTexel( terrain, hitPosition, out var x, out var y ) )
			return;

		var mat = SubTerrainSplatUtil.Read( terrain, x, y );
		var settings = _parent.PaintSettings;
		settings.SampledGroundIndex = mat.BaseTextureId;

		if ( settings.Layer == SubTerrainPaintLayer.Overlay )
			settings.LimitToSampledGround = true;

		if ( settings.Layer == SubTerrainPaintLayer.Path )
			return;

		int pick = mat.BaseTextureId;
		if ( settings.Layer == SubTerrainPaintLayer.Overlay && mat.BlendFactor > 20 )
			pick = mat.OverlayTextureId;

		settings.MaterialIndex = pick;
		TrySetSplatChannel( pick );
	}

	public static void TrySetSplatChannel( int index )
	{
		PaintTextureTool.SplatChannel = index;
	}

	void OnPaint( Terrain terrain, Vector3 hitPosition )
	{
		var brushTexture = _parent.SelectedBrushTexture;
		if ( brushTexture is null )
			return;

		var settings = _parent.PaintSettings;
		ActiveLayer = settings.Layer;

		var hitUV = new Vector2( hitPosition.x, hitPosition.y ) / terrain.Storage.TerrainSize;
		int size = (int)Math.Floor( _parent.BrushSettings.Size * 2.0f / terrain.Storage.TerrainSize * terrain.Storage.Resolution );
		size = Math.Max( size, 1 );

		bool shift = Application.KeyboardModifiers.HasFlag( KeyboardModifiers.Shift );
		bool ctrl = Gizmo.IsCtrlPressed;
		float strength = shift
			? _parent.BrushSettings.Opacity
			: _parent.BrushSettings.Opacity * (ctrl ? -1.0f : 1.0f);

		if ( settings.Layer == SubTerrainPaintLayer.Path && !shift )
			strength = ctrl ? -1.0f : 1.0f;

		int mode = 0;
		if ( shift )
			mode = 2;
		else if ( settings.Layer == SubTerrainPaintLayer.Path )
			mode = 3;
		else if ( settings.Layer == SubTerrainPaintLayer.Overlay && settings.MixOnly )
			mode = 1;

		int pathGround = SubTerrainSplatUtil.FindPathGround( terrain, settings );
		if ( mode == 3 && pathGround == settings.MaterialIndex )
			return;

		// Path packing uses PaintMode 3 in cs_subterrain_splat.
		int limit = -1;
		if ( settings.Layer != SubTerrainPaintLayer.Path
			&& settings.LimitToSampledGround
			&& settings.HasSampledGround )
		{
			limit = settings.SampledGroundIndex;
		}

		float mixEdge = SubTerrainSplatUtil.MixEdge( settings.MixCap );
		float mixMax = SubTerrainPaintSettings.MixMaxFor( settings.Layer );

		var cs = new ComputeShader( "terrain/cs_subterrain_splat" );
		cs.Attributes.Set( "ControlMap", terrain.ControlMap );
		cs.Attributes.Set( "Heightmap", terrain.HeightMap );
		cs.Attributes.Set( "ControlUV", hitUV );
		cs.Attributes.Set( "BrushStrength", strength );
		cs.Attributes.Set( "BrushSize", size );
		cs.Attributes.Set( "BrushRotation", _parent.BrushSettings.Rotation * MathF.PI / 180f );
		cs.Attributes.Set( "Brush", brushTexture );
		cs.Attributes.Set( "SplatChannel", settings.MaterialIndex );
		cs.Attributes.Set( "PaintLayer", (int)settings.Layer );
		cs.Attributes.Set( "PaintMode", mode );
		cs.Attributes.Set( "MixCap", Math.Clamp( settings.MixCap, 0.05f, mixMax ) );
		cs.Attributes.Set( "MixEdge", mixEdge );
		cs.Attributes.Set( "PathGround", pathGround );
		cs.Attributes.Set( "LimitGround", limit );
		cs.Attributes.Set( "PackCenter", settings.PackCenter ? 1 : 0 );
		cs.Attributes.Set( "PackThreshold", settings.PackThreshold );
		cs.Attributes.Set( "SlopeEnable", settings.OnlyOnSlope ? 1 : 0 );
		cs.Attributes.Set( "SlopeMinDegrees", settings.SlopeDegrees );
		cs.Attributes.Set( "TexelWorld", terrain.Storage.TerrainSize / Math.Max( terrain.Storage.Resolution, 1 ) );
		cs.Attributes.Set( "TerrainHeight", terrain.Storage.TerrainHeight );

		cs.Dispatch( size, size, 1 );

		var x = (int)Math.Floor( terrain.Storage.Resolution * hitUV.x ) - size / 2;
		var y = (int)Math.Floor( terrain.Storage.Resolution * hitUV.y ) - size / 2;
		_dirtyRegion.Add( new RectInt( x, y, size + 1, size + 1 ) );
	}

	void OnPaintEnded( Terrain terrain )
	{
		_dirtyRegion.Left = Math.Clamp( _dirtyRegion.Left, 0, terrain.Storage.Resolution - 1 );
		_dirtyRegion.Right = Math.Clamp( _dirtyRegion.Right, 0, terrain.Storage.Resolution - 1 );
		_dirtyRegion.Top = Math.Clamp( _dirtyRegion.Top, 0, terrain.Storage.Resolution - 1 );
		_dirtyRegion.Bottom = Math.Clamp( _dirtyRegion.Bottom, 0, terrain.Storage.Resolution - 1 );

		var dirtyRegion = _dirtyRegion;

		static uint[] CopyRegion( uint[] data, int stride, RectInt rect )
		{
			uint[] region = new uint[rect.Width * rect.Height];

			for ( int y = 0; y < rect.Height; y++ )
			{
				for ( int x = 0; x < rect.Width; x++ )
				{
					region[x + y * rect.Width] = data[rect.Left + x + (rect.Top + y) * stride];
				}
			}

			return region;
		}

		var regionBefore = CopyRegion( terrain.Storage.ControlMap, terrain.Storage.Resolution, dirtyRegion );
		terrain.SyncCPUTexture( Terrain.SyncFlags.Control, dirtyRegion );
		var regionAfter = CopyRegion( terrain.Storage.ControlMap, terrain.Storage.Resolution, dirtyRegion );

		Action CreateUndoAction( uint[] region ) => () =>
		{
			if ( !terrain.IsValid() )
				return;

			for ( int y = 0; y < dirtyRegion.Height; y++ )
			{
				for ( int x = 0; x < dirtyRegion.Width; x++ )
				{
					terrain.Storage.ControlMap[dirtyRegion.Left + x + (dirtyRegion.Top + y) * terrain.Storage.Resolution] = region[x + y * dirtyRegion.Width];
				}
			}

			terrain.SyncGPUTexture();
			terrain.UpdateCollision( Terrain.SyncFlags.Control, dirtyRegion );
		};

		SceneEditorSession.Active.UndoSystem.Insert( $"Terrain {DisplayInfo.For( this ).Name}",
			CreateUndoAction( regionBefore ),
			CreateUndoAction( regionAfter ) );
	}
}