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

using System;
using Sandbox;
using Editor.TerrainEditor;

namespace Editor.SubTerrain;

public struct SubTerrainPaintParameters
{
	public Vector3 HitPosition { get; set; }
	public Vector2 HitUV { get; set; }
	public float FlattenHeight { get; set; }
	public SubTerrainBrushSettings BrushSettings { get; set; }
}

/// <summary>
/// Shared stroke logic for SubTerrain sculpt tools.
/// </summary>
public abstract class SubTerrainBaseBrushTool : EditorTool
{
	protected SubTerrainTool _parent;
	protected bool _dragging;
	protected RectInt _dirtyRegion;

	protected SubTerrainSculptMode Mode { get; set; }
	protected Vector2 SlopeStartUV;
	protected float SlopeStartHeight01;
	protected bool AllowBrushInvert { get; set; }
	protected bool AllowHeightSample { get; set; }
	protected Plane StrokePlane;
	bool _samplingHeight;

	GpuBuffer<BrushData> _engineBrushBuffer;

	Vector3 _lastHitWorldPos;
	Vector3 _lastPaintLocal;
	Transform _lastHitTx;
	Vector2? _cursorLockPosition;
	float _raiseStampAccum;

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

	public override void OnEnabled()
	{
		AllowGameObjectSelection = false;
		_parent.SetBrushMode( GetType().Name );
	}

	public override void OnDisabled()
	{
		_parent.DetachSidebarForReuse();
	}

	public virtual bool GetHitPosition( Terrain terrain, out Vector3 position )
	{
		return terrain.RayIntersects( Gizmo.CurrentRay, Gizmo.RayDepth, out position );
	}

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

		if ( !GetHitPosition( terrain, 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;

			DrawBrushAdjustText();
			_parent.DrawBrushPreview( new Transform( _lastHitWorldPos, _lastHitTx.Rotation ), terrain );
			return;
		}

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

		_lastHitWorldPos = tx.PointToWorld( hitPosition );
		_lastHitTx = tx;

		if ( AllowHeightSample && Application.KeyboardModifiers.HasFlag( KeyboardModifiers.Alt ) )
		{
			if ( Gizmo.IsLeftMouseDown )
			{
				if ( !_samplingHeight )
				{
					_samplingHeight = true;
					_parent.BrushSettings.TargetHeight = hitPosition.z;
				}
			}
			else
			{
				_samplingHeight = false;
			}

			_parent.DrawBrushPreview( new Transform( _lastHitWorldPos, tx.Rotation ), terrain );
			DrawToolOverlay( terrain, _lastHitWorldPos );
			return;
		}

		_samplingHeight = false;

		if ( Gizmo.IsLeftMouseDown )
		{
			bool moved = !Application.CursorDelta.IsNearZeroLength;
			bool shouldSculpt = !_dragging || moved;

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

				StrokePlane = new Plane( _lastHitWorldPos, tx.Rotation.Up );
				_dragging = true;
				_raiseStampAccum = 0f;

				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;

				OnStrokeStart( terrain, hitPosition );
				OnPaint( terrain, MakePaintParams( hitPosition, terrain ) );
			}
			else if ( Mode == SubTerrainSculptMode.RaiseLower )
			{
				if ( moved )
				{
					StrokeSculpt( terrain, hitPosition );
					_raiseStampAccum = 0f;
				}
				else
				{
					// Hold in place: keep raising/lowering at Speed stamps per second.
					float t = Math.Clamp( (_parent.BrushSettings.Speed - 1) / 99f, 0f, 1f );
					float stampsPerSec = 2f + t * 58f;
					_raiseStampAccum += Math.Clamp( RealTime.Delta, 0f, 0.1f ) * stampsPerSec;
					while ( _raiseStampAccum >= 1f )
					{
						_raiseStampAccum -= 1f;
						OnPaint( terrain, MakePaintParams( hitPosition, terrain ) );
						_lastPaintLocal = hitPosition;
					}
				}
			}
			else if ( shouldSculpt )
			{
				StrokeSculpt( terrain, hitPosition );
			}
		}
		else if ( _dragging )
		{
			_dragging = false;
			_raiseStampAccum = 0f;
			OnPaintEnded( terrain );
		}

		_parent.DrawBrushPreview( new Transform( _lastHitWorldPos, tx.Rotation ), terrain );
		DrawToolOverlay( terrain, _lastHitWorldPos );
	}

	SubTerrainPaintParameters MakePaintParams( Vector3 hitPosition, Terrain terrain )
	{
		return new SubTerrainPaintParameters
		{
			HitPosition = hitPosition,
			HitUV = new Vector2( hitPosition.x, hitPosition.y ) / terrain.Storage.TerrainSize,
			FlattenHeight = hitPosition.z / terrain.Storage.TerrainHeight,
			BrushSettings = _parent.BrushSettings
		};
	}

	/// <summary>
	/// World-space stamp fill so fast / zoomed-out drags do not leave dashed strokes.
	/// </summary>
	void StrokeSculpt( Terrain terrain, Vector3 hitLocal )
	{
		var from = _lastPaintLocal;
		var to = hitLocal;
		float dist = (to - from).Length;
		float step = Math.Max( _parent.BrushSettings.Size * 0.2f, 8f );

		if ( dist < 0.001f )
			return;

		if ( dist <= step )
		{
			OnPaint( terrain, MakePaintParams( to, terrain ) );
			_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;
			OnPaint( terrain, MakePaintParams( Vector3.Lerp( from, to, t ), terrain ) );
		}

		_lastPaintLocal = to;
	}

	void DrawBrushAdjustText()
	{
		var textScope = new TextRendering.Scope
		{
			TextColor = Color.White,
			FontSize = 16 * Gizmo.Settings.GizmoScale * Application.DpiScale,
			FontName = "Roboto Mono",
			FontWeight = 600,
			LineHeight = 1,
			Outline = new TextRendering.Outline() { Color = Color.Black, Enabled = true, Size = 3 }
		};

		var offset = Vector2.Up * 24;
		var brush = _parent.BrushSettings;

		if ( Gizmo.IsShiftPressed )
		{
			textScope.Text = $"Size: {brush.Size}";
			Gizmo.Draw.ScreenText( textScope, _lastHitWorldPos, offset );

			textScope.Text = $"Opacity: {brush.Opacity:0.##}";
			Gizmo.Draw.ScreenText( textScope, _lastHitWorldPos, offset * 2 );

			if ( Mode == SubTerrainSculptMode.RaiseLower )
			{
				textScope.Text = $"Speed: {brush.Speed}";
				Gizmo.Draw.ScreenText( textScope, _lastHitWorldPos, offset * 3 );
			}
		}
		else if ( Gizmo.IsCtrlPressed && !brush.RandomRotation )
		{
			textScope.Text = $"Rotation: {brush.Rotation:0.#}";
			Gizmo.Draw.ScreenText( textScope, _lastHitWorldPos, offset );
		}
	}

	protected virtual void DrawToolOverlay( Terrain terrain, Vector3 worldCenter )
	{
	}

	protected virtual void OnStrokeStart( Terrain terrain, Vector3 hitPosition )
	{
	}

	/// <summary>
	/// World Size is brush radius; dispatch covers the diameter (Size * 2 in world units).
	/// </summary>
	protected int BrushTexelSize( Terrain terrain, float worldSize )
	{
		int size = (int)Math.Floor( worldSize * 2.0f / terrain.Storage.TerrainSize * terrain.Storage.Resolution );
		return Math.Max( size, 1 );
	}

	protected void GrowDirtyRegion( Terrain terrain, Vector2 hitUV, int size )
	{
		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 ) );
	}

	protected virtual void OnPaint( Terrain terrain, SubTerrainPaintParameters paint )
	{
		var brushTexture = _parent.SelectedBrushTexture;
		if ( brushTexture is null )
			return;

		int size = BrushTexelSize( terrain, paint.BrushSettings.Size );
		var amount = paint.BrushSettings.Opacity;
		if ( Mode == SubTerrainSculptMode.RaiseLower )
		{
			// Speed 50 = historical 1x height per stamp; 100 ≈ 4x. Hold rate is separate.
			amount *= Math.Clamp( paint.BrushSettings.Speed / 50f, 0.02f, 4f );
		}

		var strength = amount * (AllowBrushInvert && Gizmo.IsCtrlPressed ? -1.0f : 1.0f);

		DispatchSculpt( terrain, paint, brushTexture, size, strength );
	}

	protected void DispatchSculpt( Terrain terrain, SubTerrainPaintParameters paint, Texture brushTexture, int size, float strength )
	{
		if ( (int)Mode <= (int)SubTerrainSculptMode.Noise )
		{
			DispatchEngineSculpt( terrain, paint, brushTexture, size, strength );
			return;
		}

		var cs = new ComputeShader( "terrain/cs_subterrain_sculpt" );

		cs.Attributes.SetComboEnum( "D_SCULPT_MODE", Mode );

		cs.Attributes.Set( "Heightmap", terrain.HeightMap );
		cs.Attributes.Set( "ControlMap", terrain.ControlMap );

		cs.Attributes.Set( "HeightUV", paint.HitUV );
		cs.Attributes.Set( "FlattenHeight", paint.FlattenHeight );
		cs.Attributes.Set( "SlopeStartUV", SlopeStartUV );
		cs.Attributes.Set( "SlopeEndUV", paint.HitUV );
		cs.Attributes.Set( "SlopeStartHeight", SlopeStartHeight01 );
		cs.Attributes.Set( "SlopeEndHeight", paint.FlattenHeight );
		cs.Attributes.Set( "BrushStrength", strength );
		cs.Attributes.Set( "BrushSize", size );
		cs.Attributes.Set( "Brush", brushTexture );

		ApplyExtraAttributes( cs, paint );

		cs.Dispatch( size, size, 1 );

		GrowDirtyRegion( terrain, paint.HitUV, size );
	}

	/// <summary>
	/// Modes 0-4 match the engine sculpt shader, which already supports brush rotation.
	/// Avoids recompiling cs_subterrain_sculpt (that compile has been crashing the editor).
	/// </summary>
	void DispatchEngineSculpt( Terrain terrain, SubTerrainPaintParameters paint, Texture brushTexture, int size, float strength )
	{
		var cs = new ComputeShader( "terrain/cs_terrain_sculpt" );

		cs.Attributes.SetComboEnum( "D_SCULPT_MODE", (SculptMode)(int)Mode );

		cs.Attributes.Set( "Heightmap", terrain.HeightMap );
		cs.Attributes.Set( "ControlMap", terrain.ControlMap );

		_engineBrushBuffer ??= new GpuBuffer<BrushData>( 1 );
		_engineBrushBuffer.SetData( new[]
		{
			new BrushData
			{
				UV = paint.HitUV,
				Strength = strength,
				Size = size,
				Rotation = paint.BrushSettings.Rotation * MathF.PI / 180f,
				FlattenHeight = paint.FlattenHeight,
			}
		} );
		cs.Attributes.Set( "BrushSettings", _engineBrushBuffer );
		cs.Attributes.Set( "Brush", brushTexture );

		cs.Dispatch( size, size, 1 );

		GrowDirtyRegion( terrain, paint.HitUV, size );
	}

	/// <summary>
	/// Set any attributes a sub-tool needs on top of the shared sculpt ones.
	/// Runs after the shared attributes, so it can override them.
	/// </summary>
	protected virtual void ApplyExtraAttributes( ComputeShader cs, SubTerrainPaintParameters paint )
	{
	}

	private static T[] CopyRegion<T>( T[] data, int stride, RectInt rect ) where T : unmanaged
	{
		T[] region = new T[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;
	}

	private static Action CreateUndoAction<T>( Terrain terrain, T[] dest, T[] region, RectInt dirtyRegion, Terrain.SyncFlags flags ) => () =>
	{
		if ( !terrain.IsValid() )
			return;

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

		terrain.SyncGPUTexture();
		terrain.UpdateCollision( flags, dirtyRegion );
	};

	protected virtual 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 name = $"Terrain {DisplayInfo.For( this ).Name}";

		if ( Mode != SubTerrainSculptMode.Hole )
		{
			var regionBefore = CopyRegion( terrain.Storage.HeightMap, terrain.Storage.Resolution, _dirtyRegion );
			terrain.SyncCPUTexture( Terrain.SyncFlags.Height, _dirtyRegion );
			var regionAfter = CopyRegion( terrain.Storage.HeightMap, terrain.Storage.Resolution, _dirtyRegion );

			SceneEditorSession.Active.UndoSystem.Insert( name,
				CreateUndoAction( terrain, terrain.Storage.HeightMap, regionBefore, _dirtyRegion, Terrain.SyncFlags.Height ),
				CreateUndoAction( terrain, terrain.Storage.HeightMap, regionAfter, _dirtyRegion, Terrain.SyncFlags.Height ) );
		}
		else
		{
			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 );

			SceneEditorSession.Active.UndoSystem.Insert( name,
				CreateUndoAction( terrain, terrain.Storage.ControlMap, regionBefore, _dirtyRegion, Terrain.SyncFlags.Control ),
				CreateUndoAction( terrain, terrain.Storage.ControlMap, regionAfter, _dirtyRegion, Terrain.SyncFlags.Control ) );
		}
	}
}