using System;
using Sandbox;

namespace GlinerPoc.Gpu;

/// <summary>
/// Phase 8B.7+ — the FP32 numerical operations of the GLiNER GPU backend.
///
/// CONTRACT: every method here MUST be called from inside a render-context
/// job (GlinerGpuExecutor work action) — they bind Graphics.Attributes and
/// dispatch ComputeShaders, which only execute correctly with a live
/// Graphics.Context (P8A/P8A.1). Calling them outside the executor is a bug.
///
/// All ops are FP32. Ops never allocate buffers; callers pass pre-created
/// runtime buffers. Dispatch arguments are THREAD counts (the engine divides
/// by the shader's numthreads).
/// </summary>
public static class GlinerGpuOps
{
	// ---- canary ---------------------------------------------------------------

	public static void VecAddConstant( ComputeShader sh, GpuBuffer<float> src, GpuBuffer<float> dst, int count, float constant )
	{
		var a = Graphics.Attributes;
		a.Set( "InputValues", src );
		a.Set( "OutputValues", dst );
		a.Set( "ValueCount", count );
		a.Set( "AddConstant", constant );
		sh.Dispatch( count, 1, 1 );
	}

	// ---- GEMM family: Y = X*W^T (+ bias); X [rows,inDim], W [outDim,inDim] ------

	private static void BindGemm( GpuBuffer<float> x, GpuBuffer<float> w, GpuBuffer<float> bias, GpuBuffer<float> y,
		int rows, int inDim, int outDim, bool withBias )
	{
		var a = Graphics.Attributes;
		a.Set( "X", x );
		a.Set( "W", w );
		a.Set( "Bias", bias );
		a.Set( "Y", y );
		a.Set( "RowCount", rows );
		a.Set( "InDim", inDim );
		a.Set( "OutDim", outDim );
		a.Set( "HasBias", withBias ? 1 : 0 );
	}

	/// <summary>Naïve one-thread-per-output GEMM (P8A.1 kernel — permanent parity reference).</summary>
	public static void LinearNaive( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> w, GpuBuffer<float> bias, GpuBuffer<float> y,
		int rows, int inDim, int outDim, bool withBias )
	{
		BindGemm( x, w, bias, y, rows, inDim, outDim, withBias );
		sh.Dispatch( rows, outDim, 1 );
	}

	/// <summary>Tiled 16×16 groupshared GEMM (P8B.8 production kernel).</summary>
	public static void LinearTiled16( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> w, GpuBuffer<float> bias, GpuBuffer<float> y,
		int rows, int inDim, int outDim, bool withBias )
	{
		BindGemm( x, w, bias, y, rows, inDim, outDim, withBias );
		sh.Dispatch( rows, outDim, 1 );
	}

	/// <summary>Tiled 8×8 groupshared GEMM (P8B.8 variant).</summary>
	public static void LinearTiled8( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> w, GpuBuffer<float> bias, GpuBuffer<float> y,
		int rows, int inDim, int outDim, bool withBias )
	{
		BindGemm( x, w, bias, y, rows, inDim, outDim, withBias );
		sh.Dispatch( rows, outDim, 1 );
	}

	// ---- elementwise ------------------------------------------------------------

	public static void Add( ComputeShader sh, GpuBuffer<float> a, GpuBuffer<float> b, GpuBuffer<float> c, int count )
	{
		var at = Graphics.Attributes;
		at.Set( "A", a );
		at.Set( "B", b );
		at.Set( "C", c );
		at.Set( "ValueCount", count );
		sh.Dispatch( count, 1, 1 );
	}

	public static void Relu( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> y, int count )
	{
		var a = Graphics.Attributes;
		a.Set( "X", x );
		a.Set( "Y", y );
		a.Set( "ValueCount", count );
		sh.Dispatch( count, 1, 1 );
	}

	/// <summary>Exact erf GELU (A&S 7.1.26, same polynomial as the CPU oracle).</summary>
	public static void Gelu( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> y, int count )
	{
		var a = Graphics.Attributes;
		a.Set( "X", x );
		a.Set( "Y", y );
		a.Set( "ValueCount", count );
		sh.Dispatch( count, 1, 1 );
	}

	// ---- row reductions (sequential order = CPU oracle order) --------------------

	/// <summary>LayerNorm over the last dim; one thread per row; eps inside sqrt.</summary>
	public static void LayerNorm( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> weight, GpuBuffer<float> bias, GpuBuffer<float> y,
		int rows, int dim, float eps )
	{
		var a = Graphics.Attributes;
		a.Set( "X", x );
		a.Set( "LnWeight", weight );
		a.Set( "LnBias", bias );
		a.Set( "Y", y );
		a.Set( "RowCount", rows );
		a.Set( "Dim", dim );
		a.Set( "Eps", eps );
		sh.Dispatch( rows, 1, 1 );
	}

	/// <summary>Row-wise stable softmax (max-shift). Row length ≤ 256 (V1 limit).</summary>
	public static void SoftmaxRows( ComputeShader sh, GpuBuffer<float> x, GpuBuffer<float> y, int rows, int dim )
	{
		var a = Graphics.Attributes;
		a.Set( "X", x );
		a.Set( "Y", y );
		a.Set( "RowCount", rows );
		a.Set( "Dim", dim );
		sh.Dispatch( rows, 1, 1 );
	}

	// ---- index/copy ----------------------------------------------------------------

	public static void Copy( ComputeShader sh, GpuBuffer<float> src, GpuBuffer<float> dst, int count, int srcOffset = 0 )
	{
		var a = Graphics.Attributes;
		a.Set( "Src", src );
		a.Set( "Dst", dst );
		a.Set( "ValueCount", count );
		a.Set( "SrcOffset", srcOffset );
		sh.Dispatch( count, 1, 1 );
	}

	/// <summary>Row gather: dst[r,j] = src[indices[r], j]. Exact index parity required.</summary>
	public static void GatherRows( ComputeShader sh, GpuBuffer<float> src, GpuBuffer<int> indices, GpuBuffer<float> dst, int rows, int width )
	{
		var a = Graphics.Attributes;
		a.Set( "Src", src );
		a.Set( "Indices", indices );
		a.Set( "Dst", dst );
		a.Set( "RowCount", rows );
		a.Set( "Width", width );
		sh.Dispatch( rows * width, 1, 1 );
	}

	// ---- P8C attention kernels (virtual head layout, no transposes) -----------

	/// <summary>c2c: out[h,q,k] = Σ_d Q[h,q,d]·(K[h,k,d]/scale). Buffers row-major [rows,384].</summary>
	public static void C2C( ComputeShader sh, GpuBuffer<float> q, GpuBuffer<float> k, GpuBuffer<float> o, int rows, float scale )
	{
		var a = Graphics.Attributes;
		a.Set( "Q", q );
		a.Set( "K", k );
		a.Set( "Out", o );
		a.Set( "Rows", rows );
		a.Set( "Scale", scale );
		sh.Dispatch( 6 * rows * rows, 1, 1 );
	}

	/// <summary>c2p: out[h,q,k] = (Σ_d Q[h,q,d]·posK[h,clamp(q-k+256),d])/scale. posK [512,384].</summary>
	public static void C2P( ComputeShader sh, GpuBuffer<float> q, GpuBuffer<float> posK, GpuBuffer<float> o, int rows, float scale )
	{
		var a = Graphics.Attributes;
		a.Set( "Q", q );
		a.Set( "PosK", posK );
		a.Set( "Out", o );
		a.Set( "Rows", rows );
		a.Set( "Scale", scale );
		sh.Dispatch( 6 * rows * rows, 1, 1 );
	}

	/// <summary>p2c: out[h,q,k] = (Σ_d K[h,k,d]·posQ[h,clamp(q-k+256),d])/scale (pre-transpose indexing). posQ [512,384].</summary>
	public static void P2C( ComputeShader sh, GpuBuffer<float> k, GpuBuffer<float> posQ, GpuBuffer<float> o, int rows, float scale )
	{
		var a = Graphics.Attributes;
		a.Set( "K", k );
		a.Set( "PosQ", posQ );
		a.Set( "Out", o );
		a.Set( "Rows", rows );
		a.Set( "Scale", scale );
		sh.Dispatch( 6 * rows * rows, 1, 1 );
	}

	/// <summary>Elementwise combine of three tensors.</summary>
	public static void Combine3( ComputeShader sh, GpuBuffer<float> a, GpuBuffer<float> b, GpuBuffer<float> c, GpuBuffer<float> o, int count )
	{
		var at = Graphics.Attributes;
		at.Set( "A", a );
		at.Set( "B", b );
		at.Set( "C", c );
		at.Set( "Out", o );
		at.Set( "ValueCount", count );
		sh.Dispatch( count, 1, 1 );
	}

	/// <summary>Encoder padding mask (outer product); invalid scores ← float.MinValue.</summary>
	public static void MaskScores( ComputeShader sh, GpuBuffer<float> scores, GpuBuffer<uint> mask, int rows, float maskFill )
	{
		var a = Graphics.Attributes;
		a.Set( "Scores", scores );
		a.Set( "Mask", mask );
		a.Set( "Rows", rows );
		a.Set( "MaskFill", maskFill );
		sh.Dispatch( 6 * rows * rows, 1, 1 );
	}

	/// <summary>Attention context + head merge: merged[q*384+hd] = Σ_k probs[h,q,k]·V[k*384+hd].</summary>
	public static void ContextMerged( ComputeShader sh, GpuBuffer<float> probs, GpuBuffer<float> v, GpuBuffer<float> merged, int rows )
	{
		var a = Graphics.Attributes;
		a.Set( "Probs", probs );
		a.Set( "V", v );
		a.Set( "Merged", merged );
		a.Set( "Rows", rows );
		sh.Dispatch( rows * 384, 1, 1 );
	}
}