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 brush-stroke scaffolding for SubTerrain sculpt sub-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 Plane StrokePlane;

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

	public override void OnEnabled()
	{
		AllowGameObjectSelection = false;
	}

	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 ( Gizmo.IsLeftMouseDown )
		{
			bool shouldSculpt = !_dragging || !Application.CursorDelta.IsNearZeroLength;

			if ( !_dragging )
			{
				StrokePlane = new Plane( tx.PointToWorld( hitPosition ), tx.Rotation.Up );
				_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 ) );

				OnStrokeStart( terrain, hitPosition );
			}

			if ( shouldSculpt )
			{
				var parameters = new SubTerrainPaintParameters
				{
					HitPosition = hitPosition,
					HitUV = new Vector2( hitPosition.x, hitPosition.y ) / terrain.Storage.TerrainSize,
					FlattenHeight = hitPosition.z / terrain.Storage.TerrainHeight,
					BrushSettings = _parent.BrushSettings
				};

				OnPaint( terrain, parameters );
			}
		}
		else if ( _dragging )
		{
			_dragging = false;
			OnPaintEnded( terrain );
		}

		var previewTransform = new Transform( tx.PointToWorld( hitPosition ), tx.Rotation );
		_parent.DrawBrushPreview( previewTransform );
		DrawToolOverlay( terrain, tx.PointToWorld( hitPosition ) );
	}

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

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

	protected int BrushTexelSize( Terrain terrain, float worldSize )
	{
		int size = (int)Math.Floor( worldSize / 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 strength = paint.BrushSettings.Opacity * (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 )
	{
		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 );

		cs.Dispatch( size, size, 1 );

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

	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 ) => () =>
	{
		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();
	};

	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 ),
				CreateUndoAction( terrain, terrain.Storage.HeightMap, regionAfter, _dirtyRegion ) );
		}
		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 ),
				CreateUndoAction( terrain, terrain.Storage.ControlMap, regionAfter, _dirtyRegion ) );
		}
	}
}