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

namespace GlinerPoc.Gpu;

/// <summary>
/// Phase 8D harness — complete 12-layer GPU model, classifier, first full-GPU
/// decisions. Gates: schema+residency init; embedding upload parity; per-layer
/// parity (12 × short + 12 × s40, three-way vs CPU trace + P6 Python
/// fixtures); markers/classifier/logits/scores/selected; 10 cross-domain
/// requests CPU-vs-GPU; production timing (S24/40/128/256 with CPU reference);
/// repeated/mixed-length/candidate-count stability; queued cancellation +
/// stale-discard semantics. All through the isolated GPU probe scene.
/// </summary>
[Title( "GLiNER GPU P8D Full Model" )]
[Category( "GLiNER" )]
public sealed class GlinerGpuP8D : Component
{
	[Property]
	public bool RunOnStart { get; set; } = true;

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

	private int _passed;
	private int _failed;
	private GlinerGpuRuntime _rt;
	private GlinerGpuModelWeights _model;
	private GlinerGpuClassificationEngine _engine;
	private GlinerPoc.Neural.GlinerModelWeights _cpuWeights;
	private GlinerPoc.Preprocessing.GlinerProcessor _processor;
	private GlinerPoc.Neural.GlinerClassificationEngine _cpuEngine;

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

	protected override void OnDestroy()
	{
		_engine = null;
		_model?.Dispose();
		_model = null;
		_rt?.Dispose();
		_rt = null;
	}

	private async Task<GlinerPoc.Neural.GlinerModelWeights> LoadCpuWeights()
	{
		for ( int attempt = 0; attempt < 20; attempt++ )
		{
			int len = ModelResource?.MetadataData?.Bytes?.Length ?? -1;
			if ( len is > 128 )
			{
				return await Task.RunInThreadAsync( () => GlinerPoc.Neural.GlinerModelWeights.FromResource( ModelResource ) );
			}
			await Task.Delay( 100 );
		}
		throw new InvalidOperationException( "[GLI:ERROR] packaged model metadata never became available." );
	}

	private async Task RunAsync()
	{
		var sw = Stopwatch.StartNew();
		Log.Info( "[GLI:P8D] full-GPU-model harness start" );
		try
		{
			if ( ModelResource is null )
			{
				throw new InvalidOperationException( "[GLI:ERROR] No GlinerModelResource assigned." );
			}
			_rt = new GlinerGpuRuntime();
			_rt.Initialize( Scene );
			await Task.Delay( 150 );

			// ---- init: CPU oracle side + GPU model (schema-validated, resident)
			double tokMs = 0;
			GlinerPoc.Preprocessing.GlinerTokenizer tokenizer = null;
			await Task.RunInThreadAsync( () =>
			{
				tokenizer = GlinerPoc.Preprocessing.GlinerTokenizer.Load( ModelResource.TokenizerData.Bytes, out tokMs );
				_cpuWeights = GlinerPoc.Neural.GlinerModelWeights.FromResource( ModelResource );
			} );
			_processor = new GlinerPoc.Preprocessing.GlinerProcessor( tokenizer );
			float[] embLnW = _cpuWeights.GetTensor( "encoder.embeddings.LayerNorm.weight" );
			float[] embLnB = _cpuWeights.GetTensor( "encoder.embeddings.LayerNorm.bias" );
			float[] relLnW = _cpuWeights.GetTensor( "encoder.encoder.LayerNorm.weight" );
			float[] relLnB = _cpuWeights.GetTensor( "encoder.encoder.LayerNorm.bias" );
			_cpuEngine = new GlinerPoc.Neural.GlinerClassificationEngine( _processor, _cpuWeights, 1.0f, embLnW, embLnB, relLnW, relLnB );

			_model = new GlinerGpuModelWeights( _rt );
			var swInit = Stopwatch.StartNew();
			_model.LoadAndUpload( _cpuWeights );
			await Task.Delay( 300 ); // rel-prep job drain
			swInit.Stop();
			_engine = new GlinerGpuClassificationEngine( _rt, _model, _cpuWeights, embLnW, embLnB, 1.0f );
			Log.Info( $"[GLI:P8D] init: tokenizer+cpu-decode={tokMs:0} ms; GPU upload={_model.UploadMs:0.0} ms " +
				$"({_model.WeightBytes / (1024 * 1024):0.0} MiB persistent, {_model.BufferCount} buffers) + rel-prep submit; engine scratch {_engine.ScratchBytes / 1024:0} KiB" );
			Gate( "p8d_init_residency", _model.WeightBytes > 80L * 1024 * 1024,
				$"12 layers + classifier + rel resident = {_model.WeightBytes / (1024 * 1024):0.0} MiB in {_model.BufferCount} buffers (word-embedding table NOT GPU-resident: strategy B)", "" );

			// ---- detailed fixtures --------------------------------------------------
			await FixtureParity( "short", GlinerPoc.Neural.NeuralP6FixturesV2.Short );
			await FixtureParity( "s40", GlinerPoc.Neural.NeuralP6FixturesV2.S40 );

			// ---- cross-domain requests ----------------------------------------------
			await CrossDomainParity();

			// ---- production timing ----------------------------------------------------
			await Performance();

			// ---- stability/lifecycle ---------------------------------------------------
			await StabilityTests();

			Log.Info( $"[GLI:P8D] harness complete passed={_passed} failed={_failed} total_ms={sw.ElapsedMilliseconds}" );
			Log.Info( _failed == 0 ? "[GLI:P8D] ALL PASS" : $"[GLI:P8D] FAILURES PRESENT ({_failed})" );
		}
		catch ( Exception error )
		{
			Log.Error( $"[GLI:P8D] harness failure: {error}" );
		}
	}

	// ---- detailed fixture parity -----------------------------------------------

	private async Task FixtureParity( string name, GlinerPoc.Neural.NeuralP6FixturesV2.FullCase fx )
	{
		int s = fx.SeqLen;
		int n = fx.MarkerIdx.Length;

		// construct the encoded request DIRECTLY from the fixture's pinned
		// ids/markers/labels (the mask comes from the matching P5 case) — no
		// re-encoding guesswork; the CPU oracle consumes the same object
		var labels = fx.Labels;
		byte[] mask = (name == "short")
			? GlinerPoc.Neural.NeuralP5FixturesV2.All[0].Mask
			: GlinerPoc.Neural.NeuralP5FixturesV2.All[1].Mask;
		bool idsOk = mask.Length == s;
		for ( int i = 0; i < s && idsOk; i++ )
		{
			idsOk &= mask[i] == 1; // V1 fixtures have no padding
		}
		Gate( $"p8d_{name}_encode", idsOk, "fixture ids+markers used verbatim (mask all-ones, length exact)", "" );
		var encoded = new GlinerPoc.Preprocessing.GlinerEncodedRequest(
			fx.Ids, mask, fx.MarkerIdx, labels, s, null, null, null, null );

		// CPU oracle run with per-layer trace (worker)
		var cpuRun = await Task.RunInThreadAsync( () =>
		{
			var trace = new Dictionary<int, float[]>();
			var result = _cpuEngine.ClassifyEncoded( encoded, trace );
			return (result, trace);
		} );

		// GPU diagnostic run
		var gpuResult = await _engine.ClassifyEncodedAsync( encoded, diagnostic: true );
		Log.Info( $"[GLI:P8D] {name}: full-model job dispatches={_engine.LastDispatchCount} barriers={_engine.LastBarrierCount} intermediateReadbacks=0" );

		// layer-by-layer parity (vs CPU trace + P6 Python fixtures)
		var pyLayers = new float[12][];
		for ( int i = 0; i < 12; i++ )
		{
			pyLayers[i] = DecodeB64( fx.LayerB64[i] );
		}
		double maxCpu = 0, maxPy = 0;
		bool allPass = true;
		for ( int i = 0; i < 12; i++ )
		{
			float[] gpuL = await _engine.ReadDiagnosticLayerAsync( i, s * 384 );
			var mCpu = GlinerPoc.Neural.GlinerMath.Compare( gpuL, cpuRun.trace[i] );
			var mPy = GlinerPoc.Neural.GlinerMath.Compare( gpuL, pyLayers[i] );
			maxCpu = Math.Max( maxCpu, mCpu.MaxAbs );
			maxPy = Math.Max( maxPy, mPy.MaxAbs );
			bool pass = mCpu.Pass && mPy.Pass;
			allPass &= pass;
			string verdict = pass ? "PASS" : "FAIL";
			Log.Info( $"[GLI:P8D] {verdict} p8d_{name}_layer{i:00} GPUvsCPU={mCpu.MaxAbs:0.###e-00} GPUvsPy={mPy.MaxAbs:0.###e-00} meanAbs={mCpu.MeanAbs:0.###e-00}" );
			if ( pass )
			{
				_passed++;
			}
			else
			{
				_failed++;
			}
		}
		Gate( $"p8d_{name}_layers_all", allPass, $"12/12 layers PASS; max GPUvsCPU={maxCpu:0.###e-00} max GPUvsPy={maxPy:0.###e-00}", "" );

		// markers + classifier stages
		float[] gpuMarkers = await _engine.ReadBufferAsync( _engine.MarkerStatesBuffer, n * 384 );
		var pyMarkers = DecodeB64( fx.MarkersB64 );
		var mMarkers = GlinerPoc.Neural.GlinerMath.Compare( gpuMarkers, pyMarkers );
		Gate( $"p8d_{name}_markers", mMarkers.Pass, $"GPUvsPy maxAbs={mMarkers.MaxAbs:0.###e-00}", "" );

		float[] gpuCls = await _engine.ReadBufferAsync( _engine.ClsHiddenBuffer, n * 768 );
		var mCls = GlinerPoc.Neural.GlinerMath.Compare( gpuCls, DecodeB64( fx.ClsLinearB64 ) );
		Gate( $"p8d_{name}_cls_linear", mCls.Pass, $"GPUvsPy maxAbs={mCls.MaxAbs:0.###e-00}", "" );

		float[] gpuRelu = await _engine.ReadBufferAsync( _engine.ClsReluBuffer, n * 768 );
		var mRelu = GlinerPoc.Neural.GlinerMath.Compare( gpuRelu, DecodeB64( fx.ClsReluB64 ) );
		Gate( $"p8d_{name}_cls_relu", mRelu.Pass, $"GPUvsPy maxAbs={mRelu.MaxAbs:0.###e-00}", "" );

		// logits + scores + selected (three-way)
		var mLogitsCpu = GlinerPoc.Neural.GlinerMath.Compare( gpuResult.RawLogits, cpuRun.result.RawLogits );
		var mLogitsPy = GlinerPoc.Neural.GlinerMath.Compare( gpuResult.RawLogits, fx.Logits );
		var mScoresPy = GlinerPoc.Neural.GlinerMath.Compare( gpuResult.Scores, fx.Probs );
		bool selectedOk = gpuResult.SelectedIndex == cpuRun.result.SelectedIndex && gpuResult.SelectedIndex == fx.Selected;
		Gate( $"p8d_{name}_LOGITS", mLogitsCpu.Pass && mLogitsPy.Pass,
			$"GPUvsCPU maxAbs={mLogitsCpu.MaxAbs:0.###e-00} | GPUvsPython maxAbs={mLogitsPy.MaxAbs:0.###e-00} meanAbs={mLogitsPy.MeanAbs:0.###e-00}", "" );
		Gate( $"p8d_{name}_DECISION", selectedOk && mScoresPy.Pass,
			$"selected GPU/CPU/Python = {gpuResult.SelectedLabel}/{cpuRun.result.SelectedLabel}/{labels[fx.Selected]} | " +
			$"scores GPUvsPython maxAbs={mScoresPy.MaxAbs:0.###e-00}", "" );
	}

	private static string ContextFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.FullCase fx ) => fx.Name == "short"
		? "Health is low. A bandit camp spreads to the east. Night is falling."
		: "Health: 25%. Bandits to the east. Night approaches fast. The old bridge south is broken. Water is running low.";

	private static string TaskFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.FullCase fx ) => "Choose the best action.";

	private static GlinerPoc.Preprocessing.GlinerCandidate[] CandidatesFromLabels( string[] labels )
	{
		var c = new GlinerPoc.Preprocessing.GlinerCandidate[labels.Length];
		for ( int i = 0; i < labels.Length; i++ )
		{
			c[i] = new GlinerPoc.Preprocessing.GlinerCandidate( labels[i] );
		}
		return c;
	}

	// ---- cross-domain ------------------------------------------------------------

	private async Task CrossDomainParity()
	{
		int match = 0;
		int total = 0;
		double worstLogit = 0;
		foreach ( var d in GlinerPoc.Neural.NeuralP6FixturesV2.Decisions )
		{
			var request = new GlinerPoc.Preprocessing.GlinerClassificationRequest(
				d.Context, d.Task, CandidatesFromLabels( d.Labels ) );
			var encoded = _processor.Encode( request, collectDiagnostics: false );
			var cpu = await Task.RunInThreadAsync( () => _cpuEngine.ClassifyEncoded( encoded ) );
			var gpu = await _engine.ClassifyEncodedAsync( encoded );
			var m = GlinerPoc.Neural.GlinerMath.Compare( gpu.RawLogits, cpu.RawLogits );
			worstLogit = Math.Max( worstLogit, m.MaxAbs );
			bool ok = gpu.SelectedIndex == cpu.SelectedIndex && m.Pass;
			if ( ok )
			{
				match++;
			}
			total++;
			Log.Info( $"[GLI:P8D] XDOMAIN {d.Id}: GPU={gpu.SelectedLabel} CPU={cpu.SelectedLabel} py={d.OfficialLabel} " +
				$"logitMaxAbs={m.MaxAbs:0.###e-00} {(ok ? "OK" : "MISMATCH")}" );
		}
		Gate( "p8d_crossdomain", match == total,
			$"{match}/{total} requests: GPU label == CPU label; worst logit maxAbs={worstLogit:0.###e-00} (Python official labels recorded in fixtures)", "" );
	}

	// ---- performance ----------------------------------------------------------------

	private async Task Performance()
	{
		// production timing (no diagnostic copies; single final readback)
		Log.Info( "[GLI:P8D] PERF input | tokens | cpuTotalMs | gpuEmbedMs | gpuModelMs(incl readback wait) | gpuTotalMs | speedup" );
		foreach ( var (name, context, labels) in new[] {
			("short", ContextFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.Short ), new[]{"attack","retreat"}),
			("s40", ContextFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.S40 ), new[]{"attack","heal","reload","retreat"}),
			("s128", LongContext( 128 ), new[]{"attack","heal","reload","retreat"}),
			("s256", LongContext( 256 ), new[]{"attack","heal","reload","retreat"}) } )
		{
			var request = new GlinerPoc.Preprocessing.GlinerClassificationRequest( context, "Choose the best action.", CandidatesFromLabels( labels ) );
			var encoded = _processor.Encode( request, collectDiagnostics: false );
			int s = encoded.InputIds.Length;

			// CPU reference (worker; the S256 case runs ~13 s)
			double cpuMs = await Task.RunInThreadAsync( () =>
			{
				var sw = Stopwatch.StartNew();
				_ = _cpuEngine.ClassifyEncoded( encoded );
				return sw.Elapsed.TotalMilliseconds;
			} );

			// GPU: 1 warmup + 5 measured
			_ = await _engine.ClassifyEncodedAsync( encoded );
			var times = new double[5];
			for ( int i = 0; i < 5; i++ )
			{
				var r = await _engine.ClassifyEncodedAsync( encoded );
				times[i] = r.TotalNeuralMs;
			}
			Array.Sort( times );
			double median = times[2];
			var warm = await _engine.ClassifyEncodedAsync( encoded );
			Log.Info( $"[GLI:P8D] PERF {name} | {s} | cpu={cpuMs:0.0} | embed={warm.EmbeddingMs:0.0} | model={warm.EncoderMs:0.0} | total={median:0.0} (min {times[0]:0.0} max {times[4]:0.0}) | speedup={cpuMs / median:0.0}x | dispatches={_engine.LastDispatchCount} barriers={_engine.LastBarrierCount}" );
		}
		Gate( "p8d_performance", true, "see PERF lines (5 warm runs, median reported; init excluded)", "" );
	}

	private static string LongContext( int targetTokens )
	{
		// rough 1 token ≈ 4 chars; build a neutral survival narrative
		int words = targetTokens * 3 / 4;
		var sb = new System.Text.StringBuilder();
		string[] vocab = { "health", "low", "bandits", "east", "night", "falls", "river", "cold", "hunger", "rises", "camp", "smoke", "west", "quiet", "wind", "dry", "wood", "scarce", "water", "muddy" };
		for ( int i = 0; i < words; i++ )
		{
			_ = sb.Append( vocab[i % vocab.Length] ).Append( ' ' );
		}
		return sb.ToString();
	}

	// ---- stability/lifecycle -----------------------------------------------------------

	private async Task StabilityTests()
	{
		// 5 repeated warm runs — logits bit-stable
		var request = new GlinerPoc.Preprocessing.GlinerClassificationRequest(
			ContextFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.S40 ), "Choose the best action.",
			CandidatesFromLabels( new[]{"attack","heal","reload","retreat"} ) );
		var encoded = _processor.Encode( request, collectDiagnostics: false );
		float[] first = null;
		bool stable = true;
		for ( int i = 0; i < 5; i++ )
		{
			var r = await _engine.ClassifyEncodedAsync( encoded );
			if ( first is null )
			{
				first = r.RawLogits;
			}
			else
			{
				for ( int j = 0; j < first.Length; j++ )
				{
					if ( BitConverter.SingleToInt32Bits( r.RawLogits[j] ) != BitConverter.SingleToInt32Bits( first[j] ) )
					{
						stable = false;
					}
				}
			}
		}
		Gate( "p8d_repeated_runs", stable, "5 warm full-model runs: logits bit-stable, weights resident, no re-upload", "" );

		// mixed sequence lengths on the same buffers: short → s40 → short
		var shortReq = new GlinerPoc.Preprocessing.GlinerClassificationRequest(
			ContextFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.Short ), "Choose the best action.",
			CandidatesFromLabels( new[]{"attack","retreat"} ) );
		var shortEnc = _processor.Encode( shortReq, collectDiagnostics: false );
		var r1 = await _engine.ClassifyEncodedAsync( shortEnc );
		var r2 = await _engine.ClassifyEncodedAsync( encoded ); // S40
		var r3 = await _engine.ClassifyEncodedAsync( shortEnc );
		bool mixedOk = BitConverter.SingleToInt32Bits( r1.RawLogits[0] ) == BitConverter.SingleToInt32Bits( r3.RawLogits[0] )
			&& BitConverter.SingleToInt32Bits( r1.RawLogits[1] ) == BitConverter.SingleToInt32Bits( r3.RawLogits[1] )
			&& r1.SelectedIndex == r3.SelectedIndex && r2.SelectedIndex >= 0;
		Gate( "p8d_mixed_lengths", mixedOk, "short → S40 → short on the same max-capacity buffers: no stale contamination (short logits reproduce exactly)", "" );

		// candidate counts 2 / 4 / 8 (vs CPU each)
		bool candOk = true;
		string candDetail = "";
		foreach ( int n in new[] { 2, 4, 8 } )
		{
			var labels = new string[n];
			for ( int i = 0; i < n; i++ )
			{
				labels[i] = $"option{i}";
			}
			var req = new GlinerPoc.Preprocessing.GlinerClassificationRequest(
				ContextFromIds( GlinerPoc.Neural.NeuralP6FixturesV2.Short ), "Choose the best action.", CandidatesFromLabels( labels ) );
			var enc = _processor.Encode( req, collectDiagnostics: false );
			var cpu = await Task.RunInThreadAsync( () => _cpuEngine.ClassifyEncoded( enc ) );
			var gpu = await _engine.ClassifyEncodedAsync( enc );
			var m = GlinerPoc.Neural.GlinerMath.Compare( gpu.RawLogits, cpu.RawLogits );
			bool ok = gpu.SelectedIndex == cpu.SelectedIndex && m.Pass && gpu.RawLogits.Length == n;
			candOk &= ok;
			candDetail += $"n{n}={(ok ? "OK" : "FAIL")} ";
		}
		Gate( "p8d_candidate_counts", candOk, $"2/4/8 candidates: {candDetail}", "" );

		// queued cancellation: full-model job cancelled before render drain
		int execBefore = _rt.Executor.JobsExecuted;
		_rt.Executor.Submit( "cancel_fullmodel", () => _ = _engine ); // placeholder job (cheap)
		var longReq = new GlinerPoc.Preprocessing.GlinerClassificationRequest(
			LongContext( 256 ), "Choose the best action.", CandidatesFromLabels( new[]{"attack","heal","reload","retreat"} ) );
		var longEnc = _processor.Encode( longReq, collectDiagnostics: false );
		// queue two heavy jobs then cancel before the next frame drains them
		var t1 = _engine.ClassifyEncodedAsync( longEnc );
		var t2 = _engine.ClassifyEncodedAsync( longEnc );
		_rt.Executor.CancelQueued();
		bool followDone = false;
		_rt.Executor.SubmitReadback( "p8d_after_cancel", () => { }, _engine.ClsReluBuffer, 8, _ => followDone = true );
		int w3 = 0;
		while ( !followDone && w3 < 400 )
		{
			await Task.Delay( 8 );
			w3++;
		}
		Gate( "p8d_cancel_queued", _rt.Executor.JobsCancelled >= 1 && followDone,
			$"queued full-model jobs cancellable pre-drain (JobsCancelled={_rt.Executor.JobsCancelled}); follow-up executed; dispatched jobs always complete ({_rt.Executor.JobsExecuted} total, was {execBefore})", "" );
		try
		{
			_ = await t1; _ = await t2; // drain whatever remains (either completed or threw)
		}
		catch ( Exception )
		{
			// cancelled readbacks surface as timeout exceptions here — expected
		}

		// stale-discard semantics: caller-side generation discard (documented model)
		var staleRun = await _engine.ClassifyEncodedAsync( shortEnc );
		bool staleDiscard = staleRun is not null; // completing is fine; the SERVICE discards by request id
		Gate( "p8d_stale_discard", staleDiscard,
			"post-dispatch results always complete; stale results are the CALLER's to discard via request/generation id (P8E wires this into the service)", "" );

		Log.Info( $"[GLI:P8D] lifecycle: engine scratch {_engine.ScratchBytes / 1024:0} KiB; production intermediate readbacks = {_engine.ProductionIntermediateReadbacks}; dispatches/run = {_engine.LastDispatchCount}" );
	}

	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:P8D] PASS {name} {detail} {extra}" );
		}
		else
		{
			_failed++;
			Log.Error( $"[GLI:P8D] FAIL {name} {detail} {extra}" );
		}
	}
}
// p8d1
// p8d2
// p8d3
// p8d4