using System;
using System.Collections.Generic;

namespace GlinerPoc.Preprocessing;

/// <summary>
/// Native GLiNER classification preprocessing (Phase 3.10–3.15). Reproduces
/// the pinned official path exactly:
///
///   context + task + candidates (+ descriptions)
///     → schema word structure: ( [P] &lt;task+descriptions&gt; ( [L] label ... ) )
///     → combined tokens: schema words + [SEP_TEXT] + context words
///     → per combined token: GlinerTokenizer.TokenizeToken
///     → input IDs (no [CLS], no [SEP]); markers = first-subword positions of
///       the [P] slot (index 1) and every [L] slot (word indices 4, 6, ...,
///       len-3), matching the official routing; classification scoring drops
///       the first ([P]) marker.
///
/// Pure managed data in/out; safe to call from any thread (see architecture
/// doc §15). V1 rejects requests whose FINAL assembled encoded sequence
/// exceeds the plan's 256-token budget with a clear error (no truncation).
/// </summary>
public sealed class GlinerProcessor
{
	public const int MaxEncodedTokens = 256;

	public const string PToken = "[P]";
	public const string LToken = "[L]";
	public const string SepTextToken = "[SEP_TEXT]";
	public const string SepStructToken = "[SEP_STRUCT]";
	public const string DescriptionToken = "[DESCRIPTION]";

	private readonly GlinerTokenizer _tokenizer;

	public GlinerProcessor( GlinerTokenizer tokenizer )
	{
		_tokenizer = tokenizer ?? throw new ArgumentNullException( nameof( tokenizer ) );
	}

	public GlinerTokenizer Tokenizer => _tokenizer;

	public GlinerEncodedRequest Encode( GlinerClassificationRequest request )
	{
		return Encode( request, collectDiagnostics: true );
	}

	public GlinerEncodedRequest Encode( GlinerClassificationRequest request, bool collectDiagnostics )
	{
		ValidateRequest( request );

		// ---- schema word structure ------------------------------------------
		var hasDescriptions = false;
		foreach ( var c in request.Candidates )
		{
			// Official semantics: description MODE activates when any label has a
			// description entry, even an empty string; only present entries are
			// appended.
			if ( c.Description is not null )
			{
				hasDescriptions = true;
				break;
			}
		}

		string prompt = request.Task;
		if ( hasDescriptions )
		{
			var sb = new System.Text.StringBuilder( prompt );
			foreach ( var c in request.Candidates )
			{
				if ( c.Description is null )
				{
					continue;
				}
				sb.Append( ' ' ).Append( DescriptionToken ).Append( ' ' )
					.Append( c.Label ).Append( ": " ).Append( c.Description );
			}
			prompt = sb.ToString();
		}

		var schema = new List<string>( 8 + request.Candidates.Count * 2 )
		{
			"(", PToken, prompt, "(",
		};
		foreach ( var c in request.Candidates )
		{
			schema.Add( LToken );
			schema.Add( c.Label );
		}
		schema.Add( ")" );
		schema.Add( ")" );

		// ---- context words -----------------------------------------------------
		string context = GlinerTextPreprocessor.ApplyContextRule( request.Context );
		var words = GlinerTextPreprocessor.SplitWords( context );

		// ---- combined tokens ----------------------------------------------------
		var combined = new List<string>( schema.Count + 1 + words.Count );
		combined.AddRange( schema );
		combined.Add( SepTextToken );
		combined.AddRange( words );

		// ---- marker routing (official rule) --------------------------------------
		// cls_marker_indices contains ONLY the [L] slots (word indices 4, 6, ...,
		// len-3). The [P] marker lives in schema_special_positions[0] on the
		// Python side and is stripped when the batch routes classification
		// markers; legacy decoding re-adds a zero row via embs[1:]. Native
		// Phase 4 gathers these [L] positions directly.
		var markerWordIndices = new List<int>( request.Candidates.Count );
		for ( int i = 4; i < schema.Count - 2; i += 2 )
		{
			markerWordIndices.Add( i );
		}

		// ---- tokenize every combined token ---------------------------------------
		var ids = new List<int>( 64 );
		var allPieces = collectDiagnostics ? new List<string>( 64 ) : null;
		var tokenPieces = collectDiagnostics ? new List<List<string>>( combined.Count ) : null;
		var markerPositions = new List<int>( markerWordIndices.Count );
		var tokenStarts = new int[combined.Count];

		for ( int i = 0; i < combined.Count; i++ )
		{
			tokenStarts[i] = ids.Count;
			var pieceList = collectDiagnostics ? new List<string>() : null;
			int[] tokenIds = _tokenizer.TokenizeToken( combined[i], pieceList );
			ids.AddRange( tokenIds );
			if ( collectDiagnostics )
			{
				tokenPieces.Add( pieceList );
				if ( pieceList is not null )
				{
					allPieces.AddRange( pieceList );
				}
			}
		}
		for ( int i = 0; i < combined.Count; i++ )
		{
			foreach ( int wordIndex in markerWordIndices )
			{
				if ( wordIndex == i )
				{
					markerPositions.Add( tokenStarts[i] );
				}
			}
		}

		if ( ids.Count > MaxEncodedTokens )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Encoded request exceeds the supported native limit: " +
				$"{ids.Count} encoded tokens > {MaxEncodedTokens}. Truncation is disabled in V1." );
		}

		var markerArray = markerPositions.ToArray();
		foreach ( int m in markerArray )
		{
			if ( m < 0 || m >= ids.Count )
			{
				throw new InvalidOperationException(
					$"[GLI:ERROR] Classification marker position {m} outside sequence length {ids.Count}." );
			}
		}
		if ( markerArray.Length != request.Candidates.Count )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Marker count {markerArray.Length} != candidate count " +
				$"{request.Candidates.Count}." );
		}

		var order = new List<string>( request.Candidates.Count );
		foreach ( var c in request.Candidates )
		{
			order.Add( c.Label );
		}

		var mask = new byte[ids.Count];
		Array.Fill( mask, (byte)1 );

		return new GlinerEncodedRequest(
			ids.ToArray(),
			mask,
			markerArray,
			order,
			words.Count,
			collectDiagnostics ? schema.ToArray() : null,
			collectDiagnostics ? words.ToArray() : null,
			collectDiagnostics ? combined.ToArray() : null,
			collectDiagnostics ? allPieces.ToArray() : null );
	}

	private static void ValidateRequest( GlinerClassificationRequest request )
	{
		if ( request is null )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Request is null." );
		}
		if ( string.IsNullOrEmpty( request.Task ) )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Task is null or empty." );
		}
		if ( request.Candidates is null || request.Candidates.Count == 0 )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Request has no candidates." );
		}
		foreach ( var c in request.Candidates )
		{
			if ( c is null || c.Label is null )
			{
				throw new InvalidOperationException( "[GLI:ERROR] Candidate or label is null." );
			}
			if ( c.Label.Length == 0 )
			{
				// Official behaviour: an empty label tokenizes to nothing but
				// keeps its [L] marker. V1 parity keeps that behaviour.
				continue;
			}
		}
	}
}