using System;
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Threading.Tasks;
using Sandbox;

namespace GlinerPoc.Gpu;

/// <summary>
/// Phase 8A.1 — render-context compute dispatch probe.
///
/// P8A proved the DIRECT out-of-context ComputeShader.Dispatch path broken on
/// 26.09.22 (project shaders silently no-op "bound pipeline does not have a
/// compute shader"; library shaders crash the editor). This probe tests
/// compute execution INSIDE a real render context, using the public API:
///
///   1. SceneCustomObject.RenderSceneObject() override — the official
///      in-render hook (exactly the pattern of Facepunch's gpu-voxels
///      library: a scene object whose render callback drains an action
///      queue on the render thread). Inside it, Graphics.Attributes is a
///      live render-context attribute set and ComputeShader.Dispatch runs
///      "async" per its own docs.
///   2. Rendering.CommandList.DispatchCompute recorded on a command list
///      attached to the camera via CameraComponent.AddCommandList (the
///      decompiled entry executes computeShader.DispatchWithAttributes(
///      Graphics.Attributes, ...) during the camera render).
///   3. GpuBuffer.GetDataAsync issued from inside the render hook — the
///      place P8A proved a render context exists.
///
/// Canary policy (P8A.1 §6): outputs poisoned to 0xdeadbeef before dispatch;
/// success = every element equals the expected value, all elements checked,
/// expected-vs-actual mismatches reported. "No exception" is never success.
/// </summary>
[Title( "GLiNER GPU CommandList Probe" )]
[Category( "GLiNER" )]
public sealed class GlinerGpuCommandListProbe : Component
{
	[Property]
	public bool RunOnStart { get; set; } = true;

	[Property]
	public GlinerPoc.Packaging.GlinerModelResource ModelResource { get; set; }

	private int _passed;
	private int _failed;

	internal readonly ConcurrentQueue<Action> RenderQueue = new();
	private GlinerClRenderHook _hook;

	private static float Lcg( ref uint state )
	{
		state = state * 1664525u + 1013904223u;
		return ((state >> 8) & 0xFFFF) / 65535f - 0.5f;
	}

	private static float[] LcgArray( int count, uint seed )
	{
		uint s = seed;
		var a = new float[count];
		for ( int i = 0; i < count; i++ )
		{
			a[i] = Lcg( ref s );
		}
		return a;
	}

	protected override void OnStart()
	{
		if ( RunOnStart )
		{
			_ = RunAsync();
		}
	}

	protected override void OnDestroy()
	{
		_hook?.Delete();
		_hook = null;
	}

	private async Task RunAsync()
	{
		var sw = Stopwatch.StartNew();
		Log.Info( "[GLI:CL] render-context compute probe start (SceneCustomObject.RenderSceneObject + Rendering.CommandList)" );
		try
		{
			var camera = FindMainCamera();
			if ( camera is null )
			{
				throw new InvalidOperationException( "[GLI:ERROR] No CameraComponent found in the scene." );
			}
			Log.Info( $"[GLI:CL] camera resolved: {camera.GameObject.Name}" );

			_hook = new GlinerClRenderHook( Scene.SceneWorld, this );
			await Task.Delay( 200 ); // let the hook render at least once

			// 1: render-hook dispatch with the project (main-Assets) shader
			bool projectOk = await GateHookVec( "Shaders/gliner/gliner_cl_vec_add.shader", "hook_vec_project" );

			// 2: single controlled differential with the library shader
			//    (registers CS modules; DIRECT out-of-context dispatch
			//    crashed the editor in P8A)
			bool libraryOk = false;
			if ( !projectOk )
			{
				libraryOk = await GateHookVec( "Shaders/gliner/gliner_vec_add.shader", "hook_vec_library" );
			}
			if ( !projectOk && !libraryOk )
			{
				Log.Info( "[GLI:CL] RENDER-CONTEXT PATH ALSO FAILS — documenting BLOCK (P8A.1 §9)" );
			}

			string goodPath = projectOk ? "Shaders/gliner/gliner_cl_vec_add.shader"
				: (libraryOk ? "Shaders/gliner/gliner_vec_add.shader" : null);

			// 3: camera-attached CommandList.DispatchCompute variant (uses the
			//    bindings the hook places into Graphics.Attributes)
			if ( goodPath is not null )
			{
				await GateCameraCommandListVec( camera, goodPath );
				await GateHookAsyncReadback( goodPath );
				await GateGemmSynthetic();
				await GateGemmRealLayer0Q();
			}

			Log.Info( $"[GLI:CL] probe complete passed={_passed} failed={_failed} total_ms={sw.ElapsedMilliseconds}" );
			Log.Info( _failed == 0
				? "[GLI:CL] ALL PASS — render-context GPU compute executes"
				: $"[GLI:CL] FAILURES PRESENT ({_failed})" );
		}
		catch ( Exception error )
		{
			Log.Error( $"[GLI:CL] harness failure: {error}" );
		}
		finally
		{
			_hook?.Delete();
			_hook = null;
		}
	}

	private CameraComponent FindMainCamera()
	{
		CameraComponent fallback = null;
		foreach ( var go in Scene.GetAllObjects( false ) )
		{
			var cam = go.GetComponent<CameraComponent>();
			if ( cam is null || !cam.Enabled )
			{
				continue;
			}
			if ( cam.IsMainCamera )
			{
				return cam;
			}
			fallback ??= cam;
		}
		return fallback;
	}

	// ---- canary: out[i] = i*0.25f + constant, bit-exact -------------------------

	private async Task<bool> GateHookVec( string shaderPath, string gateName )
	{
		const int N = 64;
		var input = new float[N];
		for ( int i = 0; i < N; i++ )
		{
			input[i] = i * 0.25f;
		}
		var expected = new float[N];
		for ( int i = 0; i < N; i++ )
		{
			expected[i] = input[i] + 1.0f;
		}

		try
		{
			var shader = new ComputeShader( shaderPath );
			using var inBuf = new GpuBuffer<float>( N );
			using var outBuf = new GpuBuffer<float>( N );
			inBuf.SetData( input );
			outBuf.Clear( 0xdeadbeefu );

			int dispatchCount = 0;
			// executed INSIDE RenderSceneObject on the render thread:
			RenderQueue.Enqueue( () =>
			{
				var attrs = Graphics.Attributes;
				attrs.Set( "InputValues", inBuf );
				attrs.Set( "OutputValues", outBuf );
				attrs.Set( "ValueCount", N );
				attrs.Set( "AddConstant", 1f );
				shader.Dispatch( N, 1, 1 ); // in-graphics-context: queued async per docs
				dispatchCount++;
			} );

			var sw = Stopwatch.StartNew();
			await Task.Delay( 500 ); // frames render; hook drains the queue in-context
			var result = new float[N];
			outBuf.GetData( result, 0, N );
			sw.Stop();
			await Task.Delay( 100 );

			int mismatches = 0;
			int stillPoison = 0;
			var firstBad = new System.Text.StringBuilder();
			for ( int i = 0; i < N; i++ )
			{
				if ( BitConverter.SingleToInt32Bits( result[i] ) != BitConverter.SingleToInt32Bits( expected[i] ) )
				{
					mismatches++;
					if ( BitConverter.SingleToInt32Bits( result[i] ) == unchecked((int)0xdeadbeef) )
					{
						stillPoison++;
					}
					if ( firstBad.Length < 100 )
					{
						_ = firstBad.Append( $"[{i}] exp={expected[i]:0.###e-00} act={result[i]:0.###e-00} " );
					}
				}
			}
			bool ok = mismatches == 0 && dispatchCount >= 1;
			Log.Info( $"[GLI:CL] {gateName} ('{shaderPath}'): dispatchesInContext={dispatchCount} " +
				$"expected-vs-actual mismatches={mismatches}/{N} stillPoison={stillPoison} | firstBad: {firstBad} | wallMs={sw.Elapsed.TotalMilliseconds:N0}" );
			Gate( gateName, ok,
				ok
					? $"ALL {N} elements bit-exact — KERNEL EXECUTED inside the render context"
					: (stillPoison == mismatches
						? "output untouched (poison) — in-context dispatch still no-oped"
						: "output changed but wrong values"),
				"" );
			return ok;
		}
		catch ( Exception e )
		{
			Log.Error( $"[GLI:CL] {gateName} ('{shaderPath}') threw {e.GetType().Name}: {e.Message}" );
			Gate( gateName, false, $"exception {e.GetType().Name}: {e.Message}", "" );
			return false;
		}
	}

	// ---- camera-attached CommandList.DispatchCompute ------------------------------

	private async Task GateCameraCommandListVec( CameraComponent camera, string shaderPath )
	{
		const int N = 64;
		var input = new float[N];
		for ( int i = 0; i < N; i++ )
		{
			input[i] = i * 0.25f;
		}
		try
		{
			var shader = new ComputeShader( shaderPath );
			using var inBuf = new GpuBuffer<float>( N );
			using var outBuf = new GpuBuffer<float>( N );
			inBuf.SetData( input );
			outBuf.Clear( 0xdeadbeefu );

			// CommandList entries cannot bind structured buffers (no public
			// attribute command; the decompiled DispatchCompute binds
			// Graphics.Attributes), so the render hook supplies the bindings
			// each frame and the camera stage executes the recorded dispatch.
			int bindFrames = 0;
			void Bind()
			{
				var attrs = Graphics.Attributes;
				attrs.Set( "InputValues", inBuf );
				attrs.Set( "OutputValues", outBuf );
				attrs.Set( "ValueCount", N );
				attrs.Set( "AddConstant", 3f );
				bindFrames++;
			}

			var cl = new Sandbox.Rendering.CommandList( "gliner_cl_camera_vec" );
			cl.Clear( outBuf, 0xdeadbeefu );
			cl.DispatchCompute( shader, N, 1, 1 );
			// AfterDepthPrepass (1000) executes BEFORE scene objects render, so
			// the hook's Graphics.Attributes bindings would not exist yet;
			// AfterOpaque (2000) runs after the scene pass that drains the hook.
			camera.AddCommandList( cl, Sandbox.Rendering.Stage.AfterOpaque, 0 );

			var sw = Stopwatch.StartNew();
			var result = new float[N];
			int mismatches = -1;
			int stillPoison = -1;
			for ( int frame = 0; frame < 12; frame++ )
			{
				RenderQueue.Enqueue( Bind );
				await Task.Delay( 42 );
				outBuf.GetData( result, 0, N );
				mismatches = 0;
				stillPoison = 0;
				for ( int i = 0; i < N; i++ )
				{
					float e = input[i] + 3f;
					if ( BitConverter.SingleToInt32Bits( result[i] ) != BitConverter.SingleToInt32Bits( e ) )
					{
						mismatches++;
						if ( BitConverter.SingleToInt32Bits( result[i] ) == unchecked((int)0xdeadbeef) )
						{
							stillPoison++;
						}
					}
				}
				if ( mismatches == 0 )
				{
					break;
				}
			}
			sw.Stop();
			camera.RemoveCommandList( cl );
			await Task.Delay( 100 );

			Log.Info( $"[GLI:CL] camera command list: bindFrames={bindFrames} mismatches={mismatches}/{N} stillPoison={stillPoison} wallMs={sw.Elapsed.TotalMilliseconds:N0} " +
				$"sampleActual[0..3]={result[0]:0.###e-00},{result[1]:0.###e-00},{result[2]:0.###e-00},{result[3]:0.###e-00} " +
				$"expected[0..3]={input[0] + 3f:0.###e-00},{input[1] + 3f:0.###e-00},{input[2] + 3f:0.###e-00},{input[3] + 3f:0.###e-00}" );
			Gate( "cl_camera_commandlist", mismatches == 0,
				mismatches == 0
					? $"ALL {N} elements bit-exact via CameraComponent.AddCommandList + CommandList.DispatchCompute"
					: $"not verified through the camera command list (mismatches={mismatches}, stillPoison={stillPoison}) — bindings may not reach the stage; the render-hook path remains the proven one",
				"" );
		}
		catch ( Exception e )
		{
			Log.Error( $"[GLI:CL] camera command list gate threw {e.GetType().Name}: {e.Message}" );
			Gate( "cl_camera_commandlist", false, $"exception {e.GetType().Name}: {e.Message}", "" );
		}
	}

	// ---- async readback issued inside the render hook ------------------------------

	private async Task GateHookAsyncReadback( string shaderPath )
	{
		const int N = 64;
		var input = new float[N];
		for ( int i = 0; i < N; i++ )
		{
			input[i] = i * 0.25f;
		}
		try
		{
			var shader = new ComputeShader( shaderPath );
			using var inBuf = new GpuBuffer<float>( N );
			using var outBuf = new GpuBuffer<float>( N );
			inBuf.SetData( input );

			var result = new float[N];
			bool completed = false;
			long completedAtMs = 0;
			var swAll = Stopwatch.StartNew();

			RenderQueue.Enqueue( () =>
			{
				var attrs = Graphics.Attributes;
				attrs.Set( "InputValues", inBuf );
				attrs.Set( "OutputValues", outBuf );
				attrs.Set( "ValueCount", N );
				attrs.Set( "AddConstant", 2f );
				shader.Dispatch( N, 1, 1 );
				// issued AFTER the dispatch in the same in-context callback:
				// P8A proved GetDataAsync needs exactly this context
				outBuf.GetDataAsync( ( ReadOnlySpan<float> data ) =>
				{
					for ( int i = 0; i < N; i++ )
					{
						result[i] = data[i];
					}
					completedAtMs = swAll.ElapsedMilliseconds;
					completed = true;
				} );
			} );

			int waits = 0;
			while ( !completed && waits < 120 )
			{
				await Task.Delay( 16 );
				waits++;
			}
			await Task.Delay( 100 );

			int mismatches = 0;
			for ( int i = 0; i < N; i++ )
			{
				if ( BitConverter.SingleToInt32Bits( result[i] ) != BitConverter.SingleToInt32Bits( input[i] + 2f ) )
				{
					mismatches++;
				}
			}
			Gate( "hook_async_readback", completed && mismatches == 0,
				$"completed={completed} bitExact={(mismatches == 0)} completionAfterMs={completedAtMs} waits={waits} (GetDataAsync inside RenderSceneObject — legal render context)", "" );
		}
		catch ( Exception e )
		{
			Log.Error( $"[GLI:CL] async readback gate threw {e.GetType().Name}: {e.Message}" );
			Gate( "hook_async_readback", false, $"exception {e.GetType().Name}: {e.Message}", "" );
		}
	}

	// ---- GEMM through the render hook ----------------------------------------------

	private async Task<(float[] Gpu, double WallMs)> RunHookGemm( string shaderPath,
		float[] x, float[] w, float[] bias, int rows, int inDim, int outDim )
	{
		var shader = new ComputeShader( shaderPath );
		using var xBuf = new GpuBuffer<float>( x.Length );
		using var wBuf = new GpuBuffer<float>( w.Length );
		using var bBuf = new GpuBuffer<float>( bias is null ? 1 : bias.Length );
		using var yBuf = new GpuBuffer<float>( rows * outDim );
		xBuf.SetData( x );
		wBuf.SetData( w );
		if ( bias is not null )
		{
			bBuf.SetData( bias );
		}
		yBuf.Clear( 0xdeadbeefu );

		RenderQueue.Enqueue( () =>
		{
			var attrs = Graphics.Attributes;
			attrs.Set( "X", xBuf );
			attrs.Set( "W", wBuf );
			attrs.Set( "Bias", bBuf );
			attrs.Set( "Y", yBuf );
			attrs.Set( "RowCount", rows );
			attrs.Set( "InDim", inDim );
			attrs.Set( "OutDim", outDim );
			attrs.Set( "HasBias", bias is not null ? 1 : 0 );
			shader.Dispatch( rows, outDim, 1 );
		} );

		var sw = Stopwatch.StartNew();
		await Task.Delay( 400 );
		var y = new float[rows * outDim];
		yBuf.GetData( y, 0, y.Length );
		sw.Stop();
		await Task.Delay( 100 );
		return (y, sw.Elapsed.TotalMilliseconds);
	}

	private async Task GateGemmSynthetic()
	{
		const string shaderPath = "Shaders/gliner/gliner_cl_gemm.shader";
		bool all = true;
		all &= await HookGemmCase( shaderPath, 4, 3, 2, true, 7001, "cl_gemm_syn_bias" );
		all &= await HookGemmCase( shaderPath, 4, 3, 2, false, 7002, "cl_gemm_syn_nobias" );
		all &= await HookGemmCase( shaderPath, 5, 7, 3, true, 7003, "cl_gemm_syn_odd" );
		Gate( "cl_gemm_synthetic", all, "3 synthetic shapes (incl. odd/group-remainder dims, ±bias) via render-context dispatch vs CPU oracle", "" );
	}

	private async Task<bool> HookGemmCase( string shaderPath,
		int rows, int inDim, int outDim, bool withBias, uint seed, string name )
	{
		float[] x = LcgArray( rows * inDim, seed );
		float[] w = LcgArray( outDim * inDim, seed + 1 );
		float[] b = LcgArray( outDim, seed + 2 );

		float[] cpu = await Task.RunInThreadAsync( () =>
			GlinerPoc.Neural.GlinerMath.Linear( x, rows, w, withBias ? b : null, inDim, outDim ) );

		var run = await RunHookGemm( shaderPath, x, w, withBias ? b : null, rows, inDim, outDim );
		var m = GlinerPoc.Neural.GlinerMath.Compare( run.Gpu, cpu );
		Gate( name, m.Pass,
			$"[{rows},{inDim}]x[{inDim},{outDim}] bias={withBias} maxAbs={m.MaxAbs:0.###e-00} meanAbs={m.MeanAbs:0.###e-00} " +
			$"maxRel={m.MaxRel:0.###e-00} worst={m.WorstIndex} wallMs={run.WallMs:N0}", "" );
		return m.Pass;
	}

	private async Task GateGemmRealLayer0Q()
	{
		if ( ModelResource is null )
		{
			Gate( "cl_gemm_real_layer0_query", false, "no ModelResource assigned", "" );
			return;
		}
		GlinerPoc.Neural.GlinerModelWeights weights =
			await Task.RunInThreadAsync( () => GlinerPoc.Neural.GlinerModelWeights.FromResource( ModelResource ) );

		int seq = GlinerPoc.Neural.NeuralP4FixturesV2.EmbeddingCaseSeqLen; // 40
		float[] embOut = DecodeB64( GlinerPoc.Neural.NeuralP4FixturesV2.EmbeddingStageOutput );
		float[] wq = weights.GetTensor( "encoder.encoder.layer.0.attention.self.query_proj.weight" );
		float[] bq = weights.GetTensor( "encoder.encoder.layer.0.attention.self.query_proj.bias" );

		var cpuRun = await Task.RunInThreadAsync( () =>
		{
			var sw = Stopwatch.StartNew();
			float[] y = GlinerPoc.Neural.GlinerMath.Linear( embOut, seq, wq, bq, 384, 384 );
			return (y, sw.Elapsed.TotalMilliseconds);
		} );

		var run = await RunHookGemm( "Shaders/gliner/gliner_cl_gemm.shader", embOut, wq, bq, seq, 384, 384 );
		var mGpuVsCpu = GlinerPoc.Neural.GlinerMath.Compare( run.Gpu, cpuRun.y );
		float[] pythonRef = DecodeB64( GlinerPoc.Neural.NeuralP4FixturesV2.Layer0QueryProjOutput );
		var mGpuVsPython = GlinerPoc.Neural.GlinerMath.Compare( run.Gpu, pythonRef );

		Gate( "cl_gemm_real_layer0_query", mGpuVsCpu.Pass && mGpuVsPython.Pass,
			$"packaged query_proj W[384,384]+b, X=real embedding output [{seq},384] | " +
			$"GPUvsCPU maxAbs={mGpuVsCpu.MaxAbs:0.###e-00} meanAbs={mGpuVsCpu.MeanAbs:0.###e-00} | " +
			$"GPUvsPython maxAbs={mGpuVsPython.MaxAbs:0.###e-00} | cpuMs={cpuRun.Item2:N2} attachToResultMs={run.WallMs:N0}", "" );
	}

	private static float[] DecodeB64( string b64 )
	{
		byte[] bytes = Convert.FromBase64String( b64 );
		return GlinerPoc.Neural.GlinerMath.DecodeF32( bytes, 0, bytes.Length / 4 );
	}

	private void Gate( string name, bool pass, string detail, string extra )
	{
		if ( pass )
		{
			_passed++;
			Log.Info( $"[GLI:CL] PASS {name} {detail} {extra}" );
		}
		else
		{
			_failed++;
			Log.Error( $"[GLI:CL] FAIL {name} {detail} {extra}" );
		}
	}
}

/// <summary>
/// Official in-render hook (pattern copied from Facepunch's gpu-voxels
/// SceneDummyObject): a SceneCustomObject whose render callback drains the
/// probe's action queue on the render thread, inside a live graphics context.
/// It renders nothing itself.
/// </summary>
internal sealed class GlinerClRenderHook : SceneCustomObject
{
	private readonly GlinerGpuCommandListProbe _probe;

	public GlinerClRenderHook( SceneWorld sceneWorld, GlinerGpuCommandListProbe probe )
		: base( sceneWorld )
	{
		_probe = probe;
		// always-visible bounds so the object is never frustum-culled
		Bounds = new BBox( Vector3.One * -100000f, Vector3.One * 100000f );
		RenderingEnabled = true;
	}

	public override void RenderSceneObject()
	{
		while ( _probe.RenderQueue.TryDequeue( out var action ) )
		{
			action();
		}
	}
}