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

namespace GlinerPoc.Preprocessing;

/// <summary>
/// Phase 3.18/3.19 parity harness: loads the packaged GlinerModelResource,
/// builds the native tokenizer + processor from the packaged tokenizer blob,
/// and compares native preprocessing against every embedded Python-oracle
/// fixture (Phase 1 originals + Phase 3 edge cases). No neural inference.
///
/// Comparison per fixture: context words, combined tokens, piece stream,
/// input IDs, attention mask, marker positions, marker token IDs, length.
/// Reports per-stage first mismatch and a deterministic summary.
/// </summary>
[Title( "GLiNER Preprocessing Parity" )]
[Category( "GLiNER" )]
public sealed class GlinerPreprocessingParity : Component
{
	[Property]
	public bool RunOnStart { get; set; } = true;

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

	private int _passed;
	private int _failed;

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

	private async Task RunAsync()
	{
		var sw = Stopwatch.StartNew();
		Log.Info( "[GLI:PRE] parity harness start" );

		try
		{
			if ( ModelResource is null )
			{
				throw new InvalidOperationException( "[GLI:ERROR] No GlinerModelResource assigned." );
			}
			if ( ModelResource.FormatMagic != GlinerPoc.Packaging.GlinerModelResource.ExpectedMagic )
			{
				throw new InvalidOperationException( "[GLI:ERROR] Root resource magic mismatch." );
			}

			double loadMs = 0.0;
			GlinerTokenizer tokenizer = await Task.RunInThreadAsync( () =>
				GlinerTokenizer.Load( ModelResource.TokenizerData.Bytes, out loadMs ) );
			Log.Info( $"[GLI:PRE] tokenizer loaded bytes={ModelResource.TokenizerData.Bytes.Length:N0} " +
				$"vocab={tokenizer.VocabularySize:N0} init_ms={loadMs:N1}" );

			var processor = new GlinerProcessor( tokenizer );

			// Ground truth: confirms which fixture build is loaded at runtime.
			Log.Info( $"[GLI:PRE] fixture_build first='{PreprocessingFixturesV2.All[0].Name}' " +
				$"count={PreprocessingFixturesV2.All.Length} " +
				$"nbsp_words={PreprocessingFixturesV2.All[^1].ContextWords.Length}" );

			// cold pass + correctness
			var swEncode = Stopwatch.StartNew();
			foreach ( var fixture in PreprocessingFixturesV2.All )
			{
				CheckFixture( processor, fixture );
			}
			double coldMs = swEncode.ElapsedMilliseconds;

			// warm timing pass
			swEncode.Restart();
			int warmRuns = 0;
			for ( int rep = 0; rep < 5; rep++ )
			{
				foreach ( var fixture in PreprocessingFixturesV2.All )
				{
					try
					{
						var request = BuildRequest( fixture );
						processor.Encode( request, collectDiagnostics: false );
						warmRuns++;
					}
					catch
					{
						// Timing pass only; correctness failures already counted.
					}
				}
			}
			double warmPerSuiteMs = warmRuns > 0 ? swEncode.ElapsedMilliseconds / (double)warmRuns : 0;

			Log.Info( $"[GLI:PRE] parity complete passed={_passed} failed={_failed} " +
				$"fixtures={PreprocessingFixturesV2.All.Length} cold_ms={coldMs:N1} " +
				$"warm_runs={warmRuns} warm_per_request_ms={warmPerSuiteMs:N3}" );
			Log.Info( _failed == 0
				? $"[GLI:PRE] ALL PASS ({_passed}/{PreprocessingFixturesV2.All.Length})"
				: $"[GLI:PRE] FAILURES PRESENT ({_failed} of {PreprocessingFixturesV2.All.Length})" );
		}
		catch ( Exception error )
		{
			Log.Error( $"[GLI:PRE] harness failure: {error.Message}" );
		}
	}

	private static GlinerClassificationRequest BuildRequest( PreprocessingFixturesV2.Fixture f )
	{
		var candidates = new List<GlinerCandidate>( f.Labels.Length );
		for ( int i = 0; i < f.Labels.Length; i++ )
		{
			string description = null;
			for ( int d = 0; d < f.DescLabels.Length; d++ )
			{
				if ( f.DescLabels[d] == f.Labels[i] )
				{
					description = f.DescTexts[d];
					break;
				}
			}
			candidates.Add( new GlinerCandidate( f.Labels[i], description ) );
		}
		return new GlinerClassificationRequest( f.Context, f.Task, candidates );
	}

	private void CheckFixture( GlinerProcessor processor, PreprocessingFixturesV2.Fixture f )
	{
		try
		{
			var result = processor.Encode( BuildRequest( f ), collectDiagnostics: true );
			Compare( f.Name, "words", f.ContextWords, result.ContextWords );
			// Schema tokens are validated as the combined-tokens prefix
			// (combined = schema words + [SEP_TEXT] + context words).
			Compare( f.Name, "combined", f.CombinedTokens, result.CombinedTokens );
			Compare( f.Name, "pieces", FlattenPerTokenPieces( f.PerTokenPieces ), result.Pieces );
			Compare( f.Name, "ids", f.InputIds, result.InputIds );
			Compare( f.Name, "mask", f.AttentionMask, result.AttentionMask );
			Compare( f.Name, "markers", f.MarkerPositions, result.ClassificationMarkerIndices );

			var markerTokenIds = new int[f.MarkerPositions.Length];
			for ( int i = 0; i < f.MarkerPositions.Length; i++ )
			{
				markerTokenIds[i] = result.InputIds[f.MarkerPositions[i]];
			}
			Compare( f.Name, "marker_ids", f.MarkerTokenIds, markerTokenIds );

			if ( result.TokenCount != f.SequenceLength )
			{
				throw new InvalidOperationException(
					$"length expected {f.SequenceLength}, found {result.TokenCount}" );
			}

			_passed++;
			Log.Info( $"[GLI:PRE] PASS {f.Name} len={result.TokenCount} markers={result.ClassificationMarkerIndices.Length}" );
		}
		catch ( Exception error )
		{
			_failed++;
			Log.Error( $"[GLI:PRE] FAIL {f.Name}: {error.Message}" );
		}
	}

	private static string[] FlattenPerTokenPieces( string[][] perToken )
	{
		int total = 0;
		foreach ( var t in perToken )
		{
			total += t.Length;
		}
		var all = new string[total];
		int i = 0;
		foreach ( var t in perToken )
		{
			foreach ( var p in t )
			{
				all[i++] = p;
			}
		}
		return all;
	}

	private static void Compare( string name, string stage, string[] expected, string[] actual )
	{
		if ( expected.Length != actual.Length )
		{
			throw new InvalidOperationException(
				$"{stage}: count expected {expected.Length}, found {actual.Length}" );
		}
		for ( int i = 0; i < expected.Length; i++ )
		{
			if ( expected[i] != actual[i] )
			{
				throw new InvalidOperationException(
					$"{stage}[{i}]: expected '{expected[i]}', found '{actual[i]}'" );
			}
		}
	}

	private static void Compare( string name, string stage, int[] expected, int[] actual )
	{
		if ( expected.Length != actual.Length )
		{
			throw new InvalidOperationException(
				$"{stage}: count expected {expected.Length}, found {actual.Length}" );
		}
		for ( int i = 0; i < expected.Length; i++ )
		{
			if ( expected[i] != actual[i] )
			{
				throw new InvalidOperationException(
					$"{stage}[{i}]: expected {expected[i]}, found {actual[i]}" );
			}
		}
	}

	private static void Compare( string name, string stage, byte[] expected, byte[] actual )
	{
		if ( expected.Length != actual.Length )
		{
			throw new InvalidOperationException(
				$"{stage}: count expected {expected.Length}, found {actual.Length}" );
		}
		for ( int i = 0; i < expected.Length; i++ )
		{
			if ( expected[i] != actual[i] )
			{
				throw new InvalidOperationException(
					$"{stage}[{i}]: expected {expected[i]}, found {actual[i]}" );
			}
		}
	}
}