using System;
using System.Collections.Generic;

namespace Effigy;

public enum BrushKind
{
	Smooth,
	Draw,
	Inflate,
	Grab,
	Flatten,
	Pinch
}

public enum BrushFalloff
{
	Smooth,
	Linear,
	Sharp,
	Constant
}

/// <summary>
/// Which plane a brush mirrors its samples across, or none.
///
/// ONE ENUM, SHARED BY SCULPT, PAINT AND WEIGHT PAINT. Mirroring was a `bool MirrorX` in three
/// separate sessions, which made "a model facing +X gets symmetry everywhere" true and "a model
/// facing +Y gets it nowhere" true at once, with no setting that could help. The plane is always
/// the origin plane — nothing here mirrors about an arbitrary plane, and this must not grow that.
/// </summary>
public enum MirrorAxis
{
	None,
	X,
	Y,
	Z
}

/// <summary>One sample on a stroke. The editor produces these; the kernel never learns what a mouse is.</summary>
public readonly struct BrushSample
{
	public readonly Vec3 Position;
	public readonly Vec3 Normal;
	public readonly Vec3 Direction;
	public readonly float Radius;
	public readonly float Strength;

	public BrushSample( Vec3 position, Vec3 normal, float radius, float strength, Vec3 direction = default )
	{
		Position = position;
		Normal = normal;
		Direction = direction;
		Radius = radius;
		Strength = strength;
	}
}

/// <summary>A list of samples plus the brush that consumes them.</summary>
public sealed class BrushStroke
{
	public BrushKind Kind;
	public BrushFalloff Falloff = BrushFalloff.Smooth;
	public MirrorAxis Mirror;
	public readonly List<BrushSample> Samples = new();
}

/// <summary>
/// Per-stroke undo: the original position of every vertex the stroke actually moved.
/// A naive undo snapshots the whole mesh; this stores only the working set.
/// </summary>
public sealed class BrushUndo
{
	readonly Dictionary<int, Vec3> _previous = new();

	public int Count => _previous.Count;

	/// <summary>Every vertex this stroke moved, with the position it had BEFORE.</summary>
	public IReadOnlyDictionary<int, Vec3> Previous => _previous;

	internal void Remember( int vertex, Vec3 position ) => _previous.TryAdd( vertex, position );

	/// <summary>
	/// Fold a later undo into this one, so several brush applications read as a single stroke.
	///
	/// EARLIEST WINS, which is what TryAdd gives: a vertex moved by three dabs in one stroke has to
	/// go back to where it was before the FIRST of them, not before the last. Absorbing in the other
	/// order would leave the model two thirds sculpted after an undo, which looks like a brush bug
	/// rather than an undo one.
	/// </summary>
	public void Absorb( BrushUndo later )
	{
		if ( later is null )
			throw new ArgumentNullException( nameof( later ) );

		foreach ( var (vertex, position) in later._previous )
			_previous.TryAdd( vertex, position );
	}

	public void Restore( PolyMesh mesh )
	{
		if ( mesh is null )
			throw new ArgumentNullException( nameof( mesh ) );

		foreach ( var (vertex, position) in _previous )
			mesh.Positions[vertex] = position;
	}
}

/// <summary>
/// Brushes as pure functions over a mesh and a stroke. Spatial queries go through
/// <see cref="MeshBVH"/>; the kernel does not know about a cursor.
/// </summary>
public static class Brush
{
	/// <param name="neighbors">
	/// Vertex adjacency, if the caller already has it. Optional, and worth passing whenever this
	/// is called more than once on the same mesh: building it walks every face, so an interactive
	/// sculpt that calls Apply per mouse sample rebuilt the whole mesh's adjacency for every dab -
	/// on a dense body that is most of the frame and tens of megabytes of garbage, thrown away and
	/// recomputed identically a few milliseconds later. Topology does not change under a stroke,
	/// which is exactly what makes it cacheable; see SculptSession, which holds it beside the BVH
	/// and drops both together.
	/// </param>
	public static BrushUndo Apply( PolyMesh mesh, BrushStroke stroke, SculptFrames frames = null, float[] mask = null, MeshBVH bvh = null, List<EdgeKey>[] neighbors = null )
	{
		if ( mesh is null )
			throw new ArgumentNullException( nameof( mesh ) );

		if ( stroke is null )
			throw new ArgumentNullException( nameof( stroke ) );

		if ( frames is not null && frames.Count != mesh.VertexCount )
			throw new ArgumentException( $"frames ({frames.Count}) and mesh ({mesh.VertexCount}) disagree" );

		if ( mask is not null && mask.Length != mesh.VertexCount )
			throw new ArgumentException( $"mask ({mask.Length}) and mesh ({mesh.VertexCount}) disagree" );

		bvh ??= MeshBVH.Build( mesh );

		// Frames and adjacency are each O(vertices) to build and each is read by exactly one brush —
		// frames by Inflate, adjacency by Smooth. Every other brush leaves them null, so a caller
		// that does not need them (a Flatten stroke on a dense import) skips two full walks of the
		// mesh per stroke. Built here when absent rather than silently dropped, so a bare call still
		// works.
		if ( frames is null && stroke.Kind == BrushKind.Inflate )
			frames = SculptFrames.Build( mesh );

		if ( neighbors is null && stroke.Kind == BrushKind.Smooth )
			neighbors = mesh.BuildVertexEdges();

		var found = new List<int>();
		var undo = new BrushUndo();

		foreach ( var sample in stroke.Samples )
		{
			ApplySample( mesh, stroke, frames, mask, bvh, neighbors, found, undo, sample );

			if ( stroke.Mirror == MirrorAxis.None )
				continue;

			ApplySample( mesh, stroke, frames, mask, bvh, neighbors, found, undo, Mirror( sample, stroke.Mirror ) );
		}

		return undo;
	}

	static void ApplySample(
		PolyMesh mesh, BrushStroke stroke, SculptFrames frames, float[] mask,
		MeshBVH bvh, List<EdgeKey>[] neighbors, List<int> found, BrushUndo undo, BrushSample sample )
	{
		if ( sample.Radius <= 0f )
			return;

		bvh.VerticesInRadius( mesh, sample.Position, sample.Radius, found );

		var n = sample.Normal.LengthSquared >= 0.5f ? sample.Normal.Normal : new Vec3( 0, 0, 1 );
		var planePoint = Vec3.Zero;
		var planeCount = 0;

		if ( stroke.Kind == BrushKind.Flatten )
		{
			foreach ( var vi in found )
			{
				planePoint += mesh.Positions[vi];
				planeCount++;
			}

			if ( planeCount > 0 )
				planePoint /= planeCount;
			else
				planePoint = sample.Position;
		}

		foreach ( var vi in found )
		{
			var pos = mesh.Positions[vi];
			var dist = (pos - sample.Position).Length;
			var t = dist / sample.Radius;
			var weight = Falloff( t, stroke.Falloff ) * sample.Strength * (mask is null ? 1f : mask[vi]);

			if ( MathF.Abs( weight ) < 1e-8f )
				continue;

			Vec3 next;

			switch ( stroke.Kind )
			{
				case BrushKind.Smooth:
					next = Vec3.Lerp( pos, NeighbourAverage( mesh, neighbors[vi], vi ), Math.Clamp( weight, 0f, 1f ) );
					break;

				case BrushKind.Draw:
					next = pos + n * weight;
					break;

				case BrushKind.Inflate:
					next = pos + frames.At[vi].Normal * weight;
					break;

				case BrushKind.Grab:
					next = pos + sample.Direction * weight;
					break;

				case BrushKind.Flatten:
					var d = Vec3.Dot( pos - planePoint, n );
					next = pos - n * (d * Math.Clamp( weight, 0f, 1f ));
					break;

				case BrushKind.Pinch:
					var along = Vec3.Dot( pos - sample.Position, n );
					var closest = sample.Position + n * along;
					next = Vec3.Lerp( pos, closest, Math.Clamp( weight, 0f, 1f ) );
					break;

				default:
					continue;
			}

			if ( next.AlmostEquals( pos, 1e-8f ) )
				continue;

			undo.Remember( vi, pos );
			mesh.Positions[vi] = next;
		}

		// ONLY WHERE THE BRUSH WAS. A full refit rebuilds every box in the tree, which on a dense
		// sculpt is most of the frame spent on the 99% of the mesh the sample could not reach.
		// The vertices that moved are the ones VerticesInRadius returned, so the sphere that found
		// them is exactly the region that can need new bounds.
		bvh.RefitRegion( mesh, sample.Position, sample.Radius );
	}

	public static float Falloff( float t, BrushFalloff kind )
	{
		t = Math.Clamp( t, 0f, 1f );

		return kind switch
		{
			BrushFalloff.Constant => 1f,
			BrushFalloff.Linear => 1f - t,
			BrushFalloff.Sharp => (1f - t) * (1f - t),
			_ => 1f - t * t * (3f - 2f * t)
		};
	}

	static Vec3 NeighbourAverage( PolyMesh mesh, List<EdgeKey> edges, int vi )
	{
		if ( edges.Count == 0 )
			return mesh.Positions[vi];

		var sum = Vec3.Zero;

		foreach ( var key in edges )
			sum += mesh.Positions[key.A == vi ? key.B : key.A];

		return sum / edges.Count;
	}

	/// <summary>Flip a vector across the chosen origin plane. The one place the "which component"
	/// decision lives, shared by sculpt, paint and weight paint so they cannot drift apart.</summary>
	public static Vec3 Mirror( Vec3 v, MirrorAxis axis ) => axis switch
	{
		MirrorAxis.X => new Vec3( -v.x, v.y, v.z ),
		MirrorAxis.Y => new Vec3( v.x, -v.y, v.z ),
		MirrorAxis.Z => new Vec3( v.x, v.y, -v.z ),
		_ => v
	};

	static BrushSample Mirror( BrushSample s, MirrorAxis axis ) =>
		new(
			Mirror( s.Position, axis ),
			Mirror( s.Normal, axis ),
			s.Radius,
			s.Strength,
			Mirror( s.Direction, axis ) );
}