using Editor.Prism.Compiler;
using Editor.Prism.Compiler.Ir;
using Editor.Prism.Core;
using Editor.Prism.Model;

namespace Editor.Prism.Nodes;

// ---------------------------------------------------------------------------------------------------
// WP-6 · node library A — trigonometry, rounding and screen-space derivatives.
//
// The derivative nodes are the only stage-restricted family here. They declare IStageConstrained so
// the stage planner refuses to schedule them in the vertex stage instead of letting the backend emit
// a ddx() that no vertex shader can run. Note the spelling: the real intrinsics are ddx_coarse and
// ddy_coarse — the shipped s&box ShaderGraph writes ddx_course, which is a typo, not a dialect.
// ---------------------------------------------------------------------------------------------------

/// <summary>How precisely a screen-space derivative is evaluated.</summary>
public enum PrismDerivativePrecision
{
	/// <summary>Whatever the driver picks — <c>ddx</c> / <c>ddy</c>.</summary>
	Standard,
	/// <summary>One value per 2x2 quad — <c>ddx_coarse</c> / <c>ddy_coarse</c>. Cheapest.</summary>
	Coarse,
	/// <summary>One value per pixel — <c>ddx_fine</c> / <c>ddy_fine</c>. Most accurate.</summary>
	Fine
}

// ---- trigonometry ---------------------------------------------------------------------------------

/// <summary>The sine of an angle in radians.</summary>
[NodeInfo( Id = "prism.math.sin", Title = "Sine", Category = "Math/Trigonometry", Icon = "waves",
	Keywords = ["sin", "sine", "wave", "oscillate"] )]
[NodeVersion( 1 )]
public sealed class SineNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Sin, value );
}

/// <summary>The cosine of an angle in radians.</summary>
[NodeInfo( Id = "prism.math.cos", Title = "Cosine", Category = "Math/Trigonometry", Icon = "waves",
	Keywords = ["cos", "cosine", "wave", "oscillate"] )]
[NodeVersion( 1 )]
public sealed class CosineNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Cos, value );
}

/// <summary>The tangent of an angle in radians.</summary>
[NodeInfo( Id = "prism.math.tan", Title = "Tangent", Category = "Math/Trigonometry", Icon = "show_chart",
	Keywords = ["tan", "tangent"] )]
[NodeVersion( 1 )]
public sealed class TangentMathNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Tan, value );
}

/// <summary>The angle whose sine is the input, in radians.</summary>
[NodeInfo( Id = "prism.math.asin", Title = "Arcsine", Category = "Math/Trigonometry", Icon = "waves",
	Tier = NodeTier.Advanced, Keywords = ["asin", "arcsine", "inverse sine"] )]
[NodeVersion( 1 )]
public sealed class ArcsineNode : UnaryMathNode
{
	/// <inheritdoc/>
	public override void OnValidate( ValidationContext ctx )
	{
		if ( ctx is null || ctx.IsConnected( nameof( In ) ) ) return;
		if ( DefaultIn >= -1f && DefaultIn <= 1f ) return;

		ctx.Warn( "Arcsine is only defined between -1 and 1; anything outside that produces NaN", nameof( In ) );
	}

	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Asin, value );
}

/// <summary>The angle whose cosine is the input, in radians.</summary>
[NodeInfo( Id = "prism.math.acos", Title = "Arccosine", Category = "Math/Trigonometry", Icon = "waves",
	Tier = NodeTier.Advanced, Keywords = ["acos", "arccosine", "inverse cosine"] )]
[NodeVersion( 1 )]
public sealed class ArccosineNode : UnaryMathNode
{
	/// <inheritdoc/>
	public override void OnValidate( ValidationContext ctx )
	{
		if ( ctx is null || ctx.IsConnected( nameof( In ) ) ) return;
		if ( DefaultIn >= -1f && DefaultIn <= 1f ) return;

		ctx.Warn( "Arccosine is only defined between -1 and 1; anything outside that produces NaN", nameof( In ) );
	}

	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Acos, value );
}

/// <summary>The angle whose tangent is the input, in radians.</summary>
[NodeInfo( Id = "prism.math.atan", Title = "Arctangent", Category = "Math/Trigonometry", Icon = "show_chart",
	Tier = NodeTier.Advanced, Keywords = ["atan", "arctangent", "inverse tangent"] )]
[NodeVersion( 1 )]
public sealed class ArctangentNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Atan, value );
}

/// <summary>The quadrant-aware angle of a two-dimensional direction, in radians.</summary>
[NodeInfo( Id = "prism.math.atan2", Title = "Arctangent2", Category = "Math/Trigonometry",
	Icon = "explore", Keywords = ["atan2", "angle", "direction", "heading"] )]
[NodeVersion( 1 )]
public sealed class Arctangent2Node : PrismNode
{
	/// <summary>The vertical component.</summary>
	[In( "T", Name = "Y" )] public PortRef Y { get; set; }

	/// <summary>The horizontal component.</summary>
	[In( "T", Name = "X" )] public PortRef X { get; set; }

	/// <summary>The angle in radians, between -pi and pi.</summary>
	[Out( "T", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>Y</c> is unconnected.</summary>
	[InlineValue( nameof( Y ) ), Title( "Y" )] public float DefaultY { get; set; }

	/// <summary>The literal used when <c>X</c> is unconnected.</summary>
	[InlineValue( nameof( X ) ), Title( "X" )] public float DefaultX { get; set; } = 1f;

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var (y, x) = ctx.InPair( nameof( Y ), nameof( X ) );

		if ( !y.IsValid || !x.IsValid ) return;

		ctx.Out( nameof( Out ), ctx.Call( Intrinsic.Atan2, y, x ) );
	}
}

/// <summary>The hyperbolic sine of a value.</summary>
[NodeInfo( Id = "prism.math.sinh", Title = "Hyperbolic Sine", Category = "Math/Trigonometry",
	Icon = "show_chart", Tier = NodeTier.Advanced, Keywords = ["sinh", "hyperbolic"] )]
[NodeVersion( 1 )]
public sealed class HyperbolicSineNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Sinh, value );
}

/// <summary>The hyperbolic cosine of a value.</summary>
[NodeInfo( Id = "prism.math.cosh", Title = "Hyperbolic Cosine", Category = "Math/Trigonometry",
	Icon = "show_chart", Tier = NodeTier.Advanced, Keywords = ["cosh", "hyperbolic", "catenary"] )]
[NodeVersion( 1 )]
public sealed class HyperbolicCosineNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Cosh, value );
}

/// <summary>The hyperbolic tangent of a value — a smooth soft clip into -1..1.</summary>
[NodeInfo( Id = "prism.math.tanh", Title = "Hyperbolic Tangent", Category = "Math/Trigonometry",
	Icon = "show_chart", Tier = NodeTier.Advanced, Keywords = ["tanh", "hyperbolic", "soft clip", "sigmoid"] )]
[NodeVersion( 1 )]
public sealed class HyperbolicTangentNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Tanh, value );
}

/// <summary>Convert radians to degrees.</summary>
[NodeInfo( Id = "prism.math.degrees", Title = "Degrees", Category = "Math/Trigonometry",
	Icon = "architecture", Keywords = ["degrees", "radians to degrees", "angle", "convert"] )]
[NodeVersion( 1 )]
public sealed class DegreesNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Degrees, value );
}

/// <summary>Convert degrees to radians.</summary>
[NodeInfo( Id = "prism.math.radians", Title = "Radians", Category = "Math/Trigonometry",
	Icon = "architecture", Keywords = ["radians", "degrees to radians", "angle", "convert"] )]
[NodeVersion( 1 )]
public sealed class RadiansNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Radians, value );
}

// ---- rounding -------------------------------------------------------------------------------------

/// <summary>The largest whole number no greater than the input.</summary>
[NodeInfo( Id = "prism.math.floor", Title = "Floor", Category = "Math/Round", Icon = "vertical_align_bottom",
	Keywords = ["floor", "round down", "truncate"] )]
[NodeVersion( 1 )]
public sealed class FloorNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Floor, value );
}

/// <summary>The smallest whole number no less than the input.</summary>
[NodeInfo( Id = "prism.math.ceil", Title = "Ceiling", Category = "Math/Round", Icon = "vertical_align_top",
	Keywords = ["ceil", "ceiling", "round up"] )]
[NodeVersion( 1 )]
public sealed class CeilingNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Ceil, value );
}

/// <summary>The nearest whole number.</summary>
[NodeInfo( Id = "prism.math.round", Title = "Round", Category = "Math/Round", Icon = "adjust",
	Keywords = ["round", "nearest", "whole"] )]
[NodeVersion( 1 )]
public sealed class RoundNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Round, value );
}

/// <summary>Drop the fractional part, rounding toward zero.</summary>
[NodeInfo( Id = "prism.math.trunc", Title = "Truncate", Category = "Math/Round", Icon = "content_cut",
	Keywords = ["trunc", "truncate", "toward zero", "integer part"] )]
[NodeVersion( 1 )]
public sealed class TruncateNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Trunc, value );
}

/// <summary>The fractional part of a value.</summary>
[NodeInfo( Id = "prism.math.frac", Title = "Fraction", Category = "Math/Round", Icon = "pie_chart",
	Keywords = ["frac", "fract", "fractional", "repeat", "wrap 0..1"] )]
[NodeVersion( 1 )]
public sealed class FractionNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) => ctx.Call( Intrinsic.Frac, value );
}

/// <summary>One minus the input — the standard way to invert a mask.</summary>
[NodeInfo( Id = "prism.math.oneMinus", Title = "One Minus", Category = "Math/Basic", Icon = "flip",
	Keywords = ["one minus", "invert", "complement", "1-x"] )]
[NodeVersion( 1 )]
public sealed class OneMinusNode : UnaryMathNode
{
	/// <inheritdoc/>
	protected override IrValue Apply( EmitContext ctx, IrValue value ) =>
		ctx.Bin( BinaryOp.Sub, ctx.Const( 1f ), value );
}

// ---- quantisation ---------------------------------------------------------------------------------

/// <summary>Snap a value to a fixed number of evenly spaced steps between zero and one.</summary>
[NodeInfo( Id = "prism.math.posterize", Title = "Posterize", Category = "Math/Round", Icon = "filter_b_and_w",
	Keywords = ["posterize", "quantize", "steps", "banding", "toon"] )]
[NodeVersion( 1 )]
public sealed class PosterizeNode : PrismNode
{
	/// <summary>The value to band.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>How many bands to produce.</summary>
	[In( "T", Name = "Steps" )] public PortRef Steps { get; set; }

	/// <summary>The banded value.</summary>
	[Out( "T", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public float DefaultIn { get; set; } = 0.5f;

	/// <summary>The literal used when <c>Steps</c> is unconnected.</summary>
	[InlineValue( nameof( Steps ) ), Title( "Steps" ), Range( 1f, 32f )] public float DefaultSteps { get; set; } = 4f;

	/// <inheritdoc/>
	public override void OnValidate( ValidationContext ctx )
	{
		if ( ctx is null || ctx.IsConnected( nameof( Steps ) ) ) return;
		if ( DefaultSteps >= 1f ) return;

		ctx.Warn( "Fewer than one step leaves the value untouched", nameof( Steps ) );
	}

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var (value, steps) = ctx.InPair( nameof( In ), nameof( Steps ) );

		if ( !value.IsValid || !steps.IsValid ) return;

		// floor( in * steps ) / steps, passing the value through unchanged where steps is zero.
		var banded = ctx.Call( Intrinsic.Floor, ctx.Bin( BinaryOp.Mul, value, steps ) );

		ctx.Out( nameof( Out ), PrismMathHelpers.SafeDivide( ctx, banded, steps, value ) );
	}
}

/// <summary>Snap a value to the nearest multiple of a step size.</summary>
[NodeInfo( Id = "prism.math.quantize", Title = "Quantize", Category = "Math/Round", Icon = "grid_4x4",
	Keywords = ["quantize", "snap", "grid", "step", "round to"] )]
[NodeVersion( 1 )]
public sealed class QuantizeNode : PrismNode
{
	/// <summary>The value to snap.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The spacing between snapped values.</summary>
	[In( "T", Name = "Step" )] public PortRef Step { get; set; }

	/// <summary>The snapped value.</summary>
	[Out( "T", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public float DefaultIn { get; set; } = 0.5f;

	/// <summary>The literal used when <c>Step</c> is unconnected.</summary>
	[InlineValue( nameof( Step ) ), Title( "Step" )] public float DefaultStep { get; set; } = 0.25f;

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var (value, step) = ctx.InPair( nameof( In ), nameof( Step ) );

		if ( !value.IsValid || !step.IsValid ) return;

		// round( in / step ) * step, passing the value through unchanged where the step is zero.
		var scaled = PrismMathHelpers.SafeDivide( ctx, value, step, value );
		var snapped = ctx.Call( Intrinsic.Round, scaled );
		var result = ctx.Bin( BinaryOp.Mul, snapped, step );

		var isZero = ctx.Bin( BinaryOp.Equal, step, ctx.Const( 0f ) );

		ctx.Out( nameof( Out ), ctx.Select( isZero, value, result ) );
	}
}

// ---- screen-space derivatives ---------------------------------------------------------------------

/// <summary>
/// The shared shape of a screen-space derivative node. Derivatives only exist where a 2x2 pixel quad
/// does, so every node deriving from this is pixel-stage only and says so through
/// <see cref="IStageConstrained"/> rather than failing later inside the backend.
/// </summary>
public abstract class DerivativeNode : PrismNode, IStageConstrained
{
	/// <summary>The value to differentiate.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The rate of change of the input across the screen.</summary>
	[Out( "T", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected. A constant has a derivative of zero.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public float DefaultIn { get; set; }

	/// <inheritdoc/>
	public StageMask RequiredStages => StageMask.Pixel;

	/// <inheritdoc/>
	public ShaderStage PreferredStage => ShaderStage.Pixel;

	/// <summary>The intrinsic this node calls, chosen from its precision setting.</summary>
	protected abstract Intrinsic Operation { get; }

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var value = ctx.In( nameof( In ) );

		if ( !value.IsValid ) return;

		ctx.Require( Capability.PixelDerivatives );
		ctx.Out( nameof( Out ), ctx.Call( Operation, value ) );
	}
}

/// <summary>How fast a value changes from one pixel to the next horizontally.</summary>
[NodeInfo( Id = "prism.derivative.ddx", Title = "DDX", Category = "Derivative", Icon = "swap_horiz",
	Keywords = ["ddx", "derivative", "screen space", "gradient", "dfdx"] )]
[NodeVersion( 1 )]
public sealed class DdxNode : DerivativeNode
{
	/// <summary>Whether to take the coarse per-quad derivative, the fine per-pixel one, or let the driver choose.</summary>
	public PrismDerivativePrecision Precision { get; set; } = PrismDerivativePrecision.Standard;

	/// <inheritdoc/>
	protected override Intrinsic Operation => Precision switch
	{
		PrismDerivativePrecision.Coarse => Intrinsic.DdxCoarse,
		PrismDerivativePrecision.Fine => Intrinsic.DdxFine,
		_ => Intrinsic.Ddx
	};
}

/// <summary>How fast a value changes from one pixel to the next vertically.</summary>
[NodeInfo( Id = "prism.derivative.ddy", Title = "DDY", Category = "Derivative", Icon = "swap_vert",
	Keywords = ["ddy", "derivative", "screen space", "gradient", "dfdy"] )]
[NodeVersion( 1 )]
public sealed class DdyNode : DerivativeNode
{
	/// <summary>Whether to take the coarse per-quad derivative, the fine per-pixel one, or let the driver choose.</summary>
	public PrismDerivativePrecision Precision { get; set; } = PrismDerivativePrecision.Standard;

	/// <inheritdoc/>
	protected override Intrinsic Operation => Precision switch
	{
		PrismDerivativePrecision.Coarse => Intrinsic.DdyCoarse,
		PrismDerivativePrecision.Fine => Intrinsic.DdyFine,
		_ => Intrinsic.Ddy
	};
}

/// <summary>The total screen-space rate of change, <c>abs( ddx ) + abs( ddy )</c>.</summary>
[NodeInfo( Id = "prism.derivative.fwidth", Title = "Screen Width", Category = "Derivative",
	Icon = "zoom_out_map", Keywords = ["fwidth", "ddxy", "derivative", "antialias", "filter width"] )]
[NodeVersion( 1 )]
public sealed class ScreenWidthNode : DerivativeNode
{
	/// <inheritdoc/>
	protected override Intrinsic Operation => Intrinsic.Fwidth;
}