using System;
using System.Collections.Generic;
using Sandbox;

namespace GlinerPoc.Gpu;

/// <summary>
/// Phase 8C — ONE complete DeBERTa encoder layer executing entirely on GPU
/// through the render-context executor. Exact port of the validated CPU
/// oracle (GlinerDebertaLayer / Phase 5): scale √192 divides K and the
/// c2p/p2c contributions (never Q, never the summed score); c2p/p2c gather
/// index = clamp(q−k+256, 0, 511); encoder padding mask = outer product with
/// float32-minimum fill; post-norm residuals; erf GELU.
///
/// Head layout (P8C.4, option B — NO physical transposes): Q/K/V stay
/// row-major [S,384] straight from the GEMM; the attention kernels index them
/// head-aware (Q[h,q,d] = Q[q·384 + h·64 + d]) and the context kernel writes
/// the merged [S,384] layout the output projection consumes directly.
///
/// Buffer set (S ≤ 256, per the P8B plan; persistent allocations reused
/// across requests): hiddenIn/out ping-pong, Q/K/V, posQ/posK [512,384],
/// c2c/c2p/p2c [6S²] (c2p/p2c share one scratch in production — kept separate
/// here for stage parity), scores/probs [6S²], merged, attn sum/ln buffers,
/// FFN mid/act/out.
///
/// Production = ONE executor job containing every dispatch with a UAV
/// barrier after each producer whose output is consumed later; NO intermediate
/// CPU readback. Diagnostics read stage buffers AFTER the job completes
/// (buffers persist), so parity needs no second implementation.
/// </summary>
public sealed class GlinerGpuLayer : IDisposable
{
	private const int Heads = 6;
	private const int HeadDim = 64;
	private const int Hidden = 384;
	private const int RelRows = 512;
	private const float LnEps = 1e-7f;
	private static readonly float Scale = MathF.Sqrt( HeadDim * 3f );

	public sealed class LayerGpuWeights
	{
		public GlinerGpuBufferDesc QueryW, QueryB, KeyW, KeyB, ValueW, ValueB;
		public GlinerGpuBufferDesc AttnDenseW, AttnDenseB, AttnLnW, AttnLnB;
		public GlinerGpuBufferDesc FfnUpW, FfnUpB, FfnDownW, FfnDownB;
		public GlinerGpuBufferDesc OutLnW, OutLnB;
		/// <summary>encoder.encoder.LayerNorm.* — the REL-EMBEDDING norm (not the block LNs).</summary>
		public GlinerGpuBufferDesc RelLnW, RelLnB;
	}

	private readonly GlinerGpuRuntime _rt;
	private readonly int _maxSeq;

	// scratch buffers (persistent allocations, sized for _maxSeq)
	private GpuBuffer<float> _q, _k, _v, _posQ, _posK, _normRel;
	private GpuBuffer<float> _c2c, _c2p, _p2c, _scores, _probs;
	private GpuBuffer<float> _merged, _attnSum, _postAttn, _ffnMid, _ffnAct, _ffnSum;
	private GpuBuffer<float> _residA, _residB;
	private GpuBuffer<float> _hiddenOut;
	private GpuBuffer<uint> _maskBuf;

	/// <summary>Number of dispatches issued by the last Run (diagnostic counter).</summary>
	public int LastDispatchCount { get; private set; }
	/// <summary>Number of UAV barriers issued by the last Run.</summary>
	public int LastBarrierCount { get; private set; }
	/// <summary>Zero — production Run never reads GPU buffers back to CPU.</summary>
	public int ProductionReadbacks => 0;

	public GlinerGpuLayer( GlinerGpuRuntime runtime, int maxSeq = 256 )
	{
		_rt = runtime;
		_maxSeq = maxSeq;
		Allocate();
	}

	private void Allocate()
	{
		int s = _maxSeq;
		int att = Heads * s * s;
		_normRel = new GpuBuffer<float>( RelRows * Hidden );
		_posQ = new GpuBuffer<float>( RelRows * Hidden );
		_posK = new GpuBuffer<float>( RelRows * Hidden );
		_q = new GpuBuffer<float>( s * Hidden );
		_k = new GpuBuffer<float>( s * Hidden );
		_v = new GpuBuffer<float>( s * Hidden );
		_c2c = new GpuBuffer<float>( att );
		_c2p = new GpuBuffer<float>( att );
		_p2c = new GpuBuffer<float>( att );
		_scores = new GpuBuffer<float>( att );
		_probs = new GpuBuffer<float>( att );
		_merged = new GpuBuffer<float>( s * Hidden );
		_attnSum = new GpuBuffer<float>( s * Hidden );
		_postAttn = new GpuBuffer<float>( s * Hidden );
		_ffnMid = new GpuBuffer<float>( s * 1536 );
		_ffnAct = new GpuBuffer<float>( s * 1536 );
		_ffnSum = new GpuBuffer<float>( s * Hidden );
		_residA = new GpuBuffer<float>( s * Hidden );
		_residB = new GpuBuffer<float>( s * Hidden );
		_hiddenOut = new GpuBuffer<float>( s * Hidden );
		_maskBuf = new GpuBuffer<uint>( s );
	}

	public void Dispose()
	{
		_posQ?.Dispose(); _posK?.Dispose(); _normRel?.Dispose();
		_q?.Dispose(); _k?.Dispose(); _v?.Dispose();
		_c2c?.Dispose(); _c2p?.Dispose(); _p2c?.Dispose();
		_scores?.Dispose(); _probs?.Dispose();
		_merged?.Dispose(); _attnSum?.Dispose(); _postAttn?.Dispose();
		_ffnMid?.Dispose(); _ffnAct?.Dispose(); _ffnSum?.Dispose();
		_residA?.Dispose(); _residB?.Dispose();
		_hiddenOut?.Dispose(); _maskBuf?.Dispose();
	}

	/// <summary>Total scratch bytes allocated by this layer instance.</summary>
	public long ScratchBytes =>
		( ( 2L * RelRows + 3L * _maxSeq + 3L * _maxSeq + 2L * _maxSeq + 3L * _maxSeq ) * Hidden
		+ 5L * Heads * _maxSeq * _maxSeq + 2L * _maxSeq * 1536 ) * 4;

	// ---- setup ------------------------------------------------------------------

	/// <summary>
	/// Build the positional states ONCE on GPU (own small executor job):
	/// raw rel_embeddings [512,384] → encoder.LayerNorm (eps 1e-7, the
	/// rel-embedding norm) → project with the CONTENT query/key weights
	/// (share_att_key=true) into resident posQ/posK [512,384].
	/// </summary>
	public void PrepareRelative( GpuBuffer<float> rawRel, LayerGpuWeights w )
	{
		_rt.Executor.Submit( "rel_prepare", () =>
		{
			int d = 0, b = 0;
			GlinerGpuOps.LayerNorm( _rt.LayerNorm, rawRel, w.RelLnW.Buffer, w.RelLnB.Buffer, _normRel, RelRows, Hidden, LnEps ); d++;
			Graphics.UavBarrier( _normRel ); b++;
			GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _normRel, w.QueryW.Buffer, w.QueryB.Buffer, _posQ, RelRows, Hidden, Hidden, true ); d++;
			Graphics.UavBarrier( _posQ ); b++;
			GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _normRel, w.KeyW.Buffer, w.KeyB.Buffer, _posK, RelRows, Hidden, Hidden, true ); d++;
			Graphics.UavBarrier( _posK ); b++;
			LastDispatchCount = d;
			LastBarrierCount = b;
		} );
	}

	/// <summary>Upload the per-request mask (bytes 0/1 → uint).</summary>
	public void SetMask( byte[] mask, int seq )
	{
		var m = new uint[seq];
		for ( int i = 0; i < seq; i++ )
		{
			m[i] = mask[i] != 0 ? 1u : 0u;
		}
		_maskBuf.SetData( m );
	}

	// ---- production layer ----------------------------------------------------------

	/// <summary>
	/// Issue the complete layer as ONE executor job body. Call from inside a
	/// submitted job action (the harness wraps this with a diagnostic readback
	/// of the final buffer; production reads only hiddenOut).
	/// hiddenIn must already be resident and uploaded. All dispatches +
	/// barriers documented in P08C.md; counters update per call.
	/// </summary>
	public void Run( GpuBuffer<float> hiddenIn, int seq, LayerGpuWeights w, GpuBuffer<float> layerOut,
		GpuBuffer<float> externalPosQ = null, GpuBuffer<float> externalPosK = null )
	{
		// P8D: posQ/posK are LAYER-SPECIFIC (share_att_key means each layer
		// projects the shared normalized rel table with its OWN Q/K weights);
		// the model owner passes per-layer buffers. Standalone P8C use keeps
		// the internally prepared pair.
		GpuBuffer<float> posQ = externalPosQ ?? _posQ;
		GpuBuffer<float> posK = externalPosK ?? _posK;
		int d = 0;
		int b = 0;
		void Barrier( GpuBuffer<float> buf )
		{
			Graphics.UavBarrier( buf );
			b++;
		}

		// Q/K/V projections (resident GEMM)
		GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, hiddenIn, w.QueryW.Buffer, w.QueryB.Buffer, _q, seq, Hidden, Hidden, true ); d++;
		Barrier( _q );
		GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, hiddenIn, w.KeyW.Buffer, w.KeyB.Buffer, _k, seq, Hidden, Hidden, true ); d++;
		Barrier( _k );
		GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, hiddenIn, w.ValueW.Buffer, w.ValueB.Buffer, _v, seq, Hidden, Hidden, true ); d++;
		Barrier( _v );

		// attention score terms
		int att = Heads * seq * seq;
		GlinerGpuOps.C2C( _rt.C2C, _q, _k, _c2c, seq, Scale ); d++;
		Barrier( _c2c );
		GlinerGpuOps.C2P( _rt.C2P, _q, posK, _c2p, seq, Scale ); d++;
		Barrier( _c2p );
		GlinerGpuOps.P2C( _rt.P2C, _k, posQ, _p2c, seq, Scale ); d++;
		Barrier( _p2c );

		// combine + mask + softmax (rows = 6*S)
		GlinerGpuOps.Combine3( _rt.Combine3, _c2c, _c2p, _p2c, _scores, att ); d++;
		Barrier( _scores );
		GlinerGpuOps.MaskScores( _rt.MaskKernel, _scores, _maskBuf, seq, float.MinValue ); d++;
		Barrier( _scores );
		GlinerGpuOps.SoftmaxRows( _rt.Softmax, _scores, _probs, Heads * seq, seq ); d++;
		Barrier( _probs );

		// context + merge (single kernel) + output projection
		GlinerGpuOps.ContextMerged( _rt.Context, _probs, _v, _merged, seq ); d++;
		Barrier( _merged );
		GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _merged, w.AttnDenseW.Buffer, w.AttnDenseB.Buffer, _attnSum, seq, Hidden, Hidden, true ); d++;
		Barrier( _attnSum );

		// residual(hiddenIn) + LN → post-attention (residual written to its own
		// buffer so the attn_dense stage stays intact for diagnostics)
		GlinerGpuOps.Add( _rt.Add, _attnSum, hiddenIn, _residA, seq * Hidden ); d++;
		Barrier( _residA );
		GlinerGpuOps.LayerNorm( _rt.LayerNorm, _residA, w.AttnLnW.Buffer, w.AttnLnB.Buffer, _postAttn, seq, Hidden, LnEps ); d++;
		Barrier( _postAttn );

		// FFN (P8B-proven chain) on resident weights
		GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _postAttn, w.FfnUpW.Buffer, w.FfnUpB.Buffer, _ffnMid, seq, Hidden, 1536, true ); d++;
		Barrier( _ffnMid );
		GlinerGpuOps.Gelu( _rt.Gelu, _ffnMid, _ffnAct, seq * 1536 ); d++;
		Barrier( _ffnAct );
		GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _ffnAct, w.FfnDownW.Buffer, w.FfnDownB.Buffer, _ffnSum, seq, 1536, Hidden, true ); d++;
		Barrier( _ffnSum );

		// residual(postAttn) + LN → layer output (own buffer; ffn_out stage stays pure)
		GlinerGpuOps.Add( _rt.Add, _ffnSum, _postAttn, _residB, seq * Hidden ); d++;
		Barrier( _residB );
		GlinerGpuOps.LayerNorm( _rt.LayerNorm, _residB, w.OutLnW.Buffer, w.OutLnB.Buffer, layerOut, seq, Hidden, LnEps ); d++;

		LastDispatchCount = d;
		LastBarrierCount = b;
	}

	// ---- diagnostic stage access (post-job readback only) ----------------------------

	/// <summary>Diagnostic accessor for stage buffers — NEVER called during production dispatch.</summary>
	public IReadOnlyDictionary<string, GpuBuffer<float>> StageBuffers() => new Dictionary<string, GpuBuffer<float>>
	{
		["q_proj"] = _q,
		["k_proj"] = _k,
		["v_proj"] = _v,
		["norm_rel"] = _normRel,
		["pos_q"] = _posQ,
		["pos_k"] = _posK,
		["c2c"] = _c2c,
		["c2p_contrib"] = _c2p,
		["p2c_contrib"] = _p2c,
		["scores_premask"] = _scores,
		["probs"] = _probs,
		["context"] = _merged,
		["attn_dense"] = _attnSum,
		["attn_ln"] = _postAttn,
		["ffn_pre"] = _ffnMid,
		["gelu"] = _ffnAct,
		["ffn_out"] = _ffnSum,
		["attn_residual"] = _residA,
		["out_residual"] = _residB,
	};

	/// <summary>The mask buffer (diagnostic/timing use).</summary>
	public GpuBuffer<uint> MaskBuffer => _maskBuf;

	/// <summary>Buffer receiving the final layer output (read back via executor diagnostic readback).</summary>
	public GpuBuffer<float> OutputBuffer => _hiddenOut;
}