using System;
using System.Collections.Generic;
using System.Text;

namespace GlinerPoc.Preprocessing;

/// <summary>
/// Exact native implementation of the checkpoint's tokenizer pipeline for one
/// processor token (Phase 3.4–3.9):
///
///   raw token string
///     → split on added/special tokens (raw-text match, longest match)
///     → per non-special segment: normalize (regex \s{2,}|[\n\r\t] → " ", NFC,
///       right-strip whitespace) → Metaspace (split on whitespace, prefix ▁)
///     → per ▁-piece: Unigram Viterbi over the 128k vocabulary (double scores)
///     → unmatched characters become one [UNK] each (byte_fallback = false)
///
/// Built from the packaged tokenizer.json bytes; validated against the pinned
/// checkpoint identity. Tokenizer scores are f64 in the official Rust
/// implementation; all path selection here is double-precision.
/// </summary>
public sealed class GlinerTokenizer
{
	public const int ExpectedVocabSize = 128000;
	public const int UnkId = 3;
	private const double UnkPenalty = 10.0;

	private readonly Dictionary<string, int> _idByPiece = new( 128100, StringComparer.Ordinal );
	private double[] _scoreById = new double[128100];
	private readonly List<AddedToken> _addedTokens = new();
	private double _unkScore;
	private int _maxPieceChars;

	private sealed class AddedToken
	{
		public string Content;
		public int Id;
	}

	private GlinerTokenizer()
	{
	}

	public int VocabularySize => _idByPiece.Count;

	/// <summary>Map piece → id lookup used by tests/diagnostics.</summary>
	public int? TryGetId( string piece ) =>
		_idByPiece.TryGetValue( piece, out int id ) ? id : (int?)null;

	/// <summary>
	/// Parse + validate the packaged tokenizer JSON (HF tokenizers format,
	/// Unigram model). Throws with actionable messages on any identity or
	/// structure mismatch. Reports measured load milliseconds.
	/// </summary>
	public static GlinerTokenizer Load( byte[] tokenizerJsonBytes, out double loadMilliseconds )
	{
		var sw = System.Diagnostics.Stopwatch.StartNew();
		if ( tokenizerJsonBytes is null || tokenizerJsonBytes.Length == 0 )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Tokenizer data is empty." );
		}

		string json = Encoding.UTF8.GetString( tokenizerJsonBytes );
		var dto = Json.Deserialize<TokenizerJsonDto>( json )
			?? throw new InvalidOperationException( "[GLI:ERROR] Tokenizer JSON deserialized to null." );

		var tokenizer = new GlinerTokenizer();
		tokenizer.Initialize( dto );
		loadMilliseconds = sw.ElapsedMilliseconds;
		return tokenizer;
	}

	private void Initialize( TokenizerJsonDto dto )
	{
		if ( dto.model is null )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Tokenizer JSON has no model section." );
		}
		if ( dto.model.type != "Unigram" )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Tokenizer model type expected 'Unigram', found '{dto.model.type}'." );
		}
		if ( dto.model.byte_fallback )
		{
			throw new InvalidOperationException(
				"[GLI:ERROR] Unexpected byte_fallback=true; pinned tokenizer has it disabled." );
		}
		if ( dto.model.unk_id != UnkId )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Tokenizer unk_id expected {UnkId}, found {dto.model.unk_id}." );
		}
		if ( dto.model.vocab is null || dto.model.vocab.Count != ExpectedVocabSize )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Tokenizer vocab expected {ExpectedVocabSize} entries, found " +
				$"{dto.model.vocab?.Count ?? 0}." );
		}

		double minScore = double.PositiveInfinity;
		int maxPieceChars = 1;
		_idByPiece.Clear();
		Array.Clear( _scoreById, 0, _scoreById.Length );

		for ( int i = 0; i < dto.model.vocab.Count; i++ )
		{
			var pair = dto.model.vocab[i];
			if ( pair is null || pair.Count != 2 )
			{
				throw new InvalidOperationException(
					$"[GLI:ERROR] Tokenizer vocab entry {i} is not a [piece, score] pair." );
			}
			string piece = AsString( pair[0], $"vocab[{i}].piece" );
			double score = AsDouble( pair[1], $"vocab[{i}].score" );
			_idByPiece[piece] = i;
			_scoreById[i] = score;
			if ( score < minScore )
			{
				minScore = score;
			}
			if ( piece.Length > maxPieceChars )
			{
				maxPieceChars = piece.Length;
			}
		}

		// Added special tokens: exact pinned set, matched on raw text.
		var expectedSpecials = new (int Id, string Content)[]
		{
			(0, "[PAD]"), (1, "[CLS]"), (2, "[SEP]"), (3, "[UNK]"), (128000, "[MASK]"),
			(128001, "[SEP_STRUCT]"), (128002, "[SEP_TEXT]"), (128003, "[P]"),
			(128004, "[C]"), (128005, "[E]"), (128006, "[R]"), (128007, "[L]"),
			(128008, "[EXAMPLE]"), (128009, "[OUTPUT]"), (128010, "[DESCRIPTION]"),
		};
		if ( dto.added_tokens is null || dto.added_tokens.Length != expectedSpecials.Length )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Tokenizer added_tokens expected {expectedSpecials.Length}, found " +
				$"{dto.added_tokens?.Length ?? 0}." );
		}
		foreach ( var (id, content) in expectedSpecials )
		{
			bool found = false;
			foreach ( var at in dto.added_tokens )
			{
				if ( at.special && at.id == id && at.content == content )
				{
					found = true;
					break;
				}
			}
			if ( !found )
			{
				throw new InvalidOperationException(
					$"[GLI:ERROR] Tokenizer missing added special token id={id} '{content}'." );
			}
			_addedTokens.Add( new AddedToken { Id = id, Content = content } );
			_idByPiece[content] = id;
			if ( id >= _scoreById.Length )
			{
				Array.Resize( ref _scoreById, id + 1 );
			}
			_scoreById[id] = 0.0;
		}

		_unkScore = minScore - UnkPenalty;
		_maxPieceChars = maxPieceChars;
	}

	private static string AsString( object boxed, string what )
	{
		if ( boxed is System.Text.Json.JsonElement e && e.ValueKind == System.Text.Json.JsonValueKind.String )
		{
			return e.GetString();
		}
		throw new InvalidOperationException( $"[GLI:ERROR] Tokenizer {what} is not a string." );
	}

	private static double AsDouble( object boxed, string what )
	{
		if ( boxed is System.Text.Json.JsonElement e &&
			e.ValueKind == System.Text.Json.JsonValueKind.Number )
		{
			return e.GetDouble();
		}
		throw new InvalidOperationException( $"[GLI:ERROR] Tokenizer {what} is not a number." );
	}

	/// <summary>
	/// Tokenize one processor token (schema word or context word) exactly like
	/// the official tokenizer.tokenize(). Returns token IDs; when pieces is
	/// non-null, also appends the piece strings (debug/reference mode).
	/// </summary>
	public int[] TokenizeToken( string token, List<string> pieces )
	{
		var ids = new List<int>( 8 );
		int pos = 0;
		while ( pos < token.Length )
		{
			// Longest added-token match at the current raw position.
			int matchLen = 0;
			int matchId = -1;
			foreach ( var at in _addedTokens )
			{
				string content = at.Content;
				if ( content.Length > matchLen && pos + content.Length <= token.Length &&
					string.CompareOrdinal( token, pos, content, 0, content.Length ) == 0 )
				{
					matchLen = content.Length;
					matchId = at.Id;
				}
			}
			if ( matchLen > 0 )
			{
				ids.Add( matchId );
				pieces?.Add( token.Substring( pos, matchLen ) );
				pos += matchLen;
				continue;
			}

			// Advance to the end of the current non-special segment.
			int segEnd = pos + 1;
			while ( segEnd < token.Length )
			{
				bool isSpecial = false;
				foreach ( var at in _addedTokens )
				{
					if ( segEnd + at.Content.Length <= token.Length &&
						string.CompareOrdinal( token, segEnd, at.Content, 0, at.Content.Length ) == 0 )
					{
						isSpecial = true;
						break;
					}
				}
				if ( isSpecial )
				{
					break;
				}
				segEnd++;
			}

			EncodeSegment( token.Substring( pos, segEnd - pos ), ids, pieces );
			pos = segEnd;
		}
		return ids.ToArray();
	}

	/// <summary>
	/// Normalize + Metaspace + Unigram for one non-special segment.
	/// </summary>
	private void EncodeSegment( string segment, List<int> ids, List<string> pieces )
	{
		string normalized = Normalize( segment );
		int start = 0;
		int len = normalized.Length;
		while ( start < len )
		{
			while ( start < len && char.IsWhiteSpace( normalized[start] ) )
			{
				start++;
			}
			if ( start >= len )
			{
				break;
			}
			int end = start + 1;
			while ( end < len && !char.IsWhiteSpace( normalized[end] ) )
			{
				end++;
			}
			string piece = "\u2581" + normalized.Substring( start, end - start );
			ViterbiPiece( piece, ids, pieces );
			start = end;
		}
	}

	private static string Normalize( string input )
	{
		// 1) Replace (\s{2,} | [\n\r\t]) → " "
		string collapsed = CollapseWhitespace( input );
		// 2) NFC. NOTE: .NET Normalize throws on characters it considers invalid
		// (U+FFF9-FFFB, unpaired surrogates) where Python NFC is identity;
		// those characters have no decompositions, so identity is exact.
		string formC;
		try
		{
			formC = collapsed.Normalize( NormalizationForm.FormC );
		}
		catch ( ArgumentException )
		{
			formC = collapsed;
		}
		// 3) Strip right whitespace
		return formC.TrimEnd();
	}

	private static string CollapseWhitespace( string input )
	{
		// Official normalizer "Replace": regex (\s{2,} | [\n\r\t]) -> " ".
		// Single non-target whitespace passes through unchanged.
		var sb = new StringBuilder( input.Length );
		int i = 0;
		int n = input.Length;
		while ( i < n )
		{
			char c = input[i];
			if ( !char.IsWhiteSpace( c ) )
			{
				sb.Append( c );
				i++;
				continue;
			}
			int start = i;
			while ( i < n && char.IsWhiteSpace( input[i] ) )
			{
				i++;
			}
			int runLength = i - start;
			if ( runLength >= 2 )
			{
				sb.Append( ' ' );
			}
			else
			{
				char single = input[start];
				if ( single == '\n' || single == '\r' || single == '\t' )
				{
					sb.Append( ' ' );
				}
				else
				{
					sb.Append( single );
				}
			}
		}
		return sb.ToString();
	}

	/// <summary>
	/// Unigram Viterbi over one ▁-prefixed piece. A character with no
	/// single-char vocab piece becomes one [UNK] with score minScore − 10.0
	/// (fuse_unk = false). Strict-greater relaxation with left-to-right,
	/// shortest-first edge order reproduces the official tie-breaking.
	/// </summary>
	private void ViterbiPiece( string s, List<int> ids, List<string> pieces )
	{
		// Viterbi operates on CODE POINTS (the official tokenizer is
		// codepoint-based; UTF-16 units would split astral characters such as
		// emoji into two [UNK]s).
		int n = s.Length;
		var cpStart = new int[n + 1];
		int m = 0;
		for ( int i = 0; i < n; i++ )
		{
			cpStart[m++] = i;
			if ( char.IsHighSurrogate( s[i] ) && i + 1 < n && char.IsLowSurrogate( s[i + 1] ) )
			{
				i++;
			}
		}
		cpStart[m] = n;

		var best = new double[m + 1];
		var back = new int[m + 1];
		var backId = new int[m + 1];
		for ( int i = 1; i <= m; i++ )
		{
			best[i] = double.NegativeInfinity;
		}
		best[0] = 0.0;

		for ( int b = 0; b < m; b++ )
		{
			if ( double.IsNegativeInfinity( best[b] ) )
			{
				continue;
			}
			int maxEnd = b + _maxPieceChars;
			if ( maxEnd > m )
			{
				maxEnd = m;
			}
			bool hasSingleCodepointPiece = false;
			for ( int e = b + 1; e <= maxEnd; e++ )
			{
				string candidate = s.Substring( cpStart[b], cpStart[e] - cpStart[b] );
				if ( _idByPiece.TryGetValue( candidate, out int id ) )
				{
					if ( e == b + 1 )
					{
						hasSingleCodepointPiece = true;
					}
					double cand = best[b] + _scoreById[id];
					if ( cand > best[e] )
					{
						best[e] = cand;
						back[e] = b;
						backId[e] = id;
					}
				}
			}
			if ( !hasSingleCodepointPiece )
			{
				double cand2 = best[b] + _unkScore;
				if ( cand2 > best[b + 1] )
				{
					best[b + 1] = cand2;
					back[b + 1] = b;
					backId[b + 1] = UnkId;
				}
			}
		}

		// Backtrack (debug piece strings are recovered from the source span).
		var reversedIds = new List<int>( 8 );
		var reversedSpans = new List<(int From, int To)>( 8 );
		int pos = m;
		while ( pos > 0 )
		{
			reversedIds.Add( backId[pos] );
			reversedSpans.Add( (cpStart[back[pos]], cpStart[pos]) );
			pos = back[pos];
		}
		for ( int i = reversedIds.Count - 1; i >= 0; i-- )
		{
			ids.Add( reversedIds[i] );
			pieces?.Add( s.Substring( reversedSpans[i].From, reversedSpans[i].To - reversedSpans[i].From ) );
		}
	}

	private sealed class TokenizerJsonDto
	{
		public string version { get; set; }
		public AddedTokenDto[] added_tokens { get; set; }
		public ModelDto model { get; set; }
	}

	private sealed class AddedTokenDto
	{
		public int id { get; set; }
		public string content { get; set; }
		public bool special { get; set; }
	}

	private sealed class ModelDto
	{
		public string type { get; set; }
		public int unk_id { get; set; }
		public List<List<object>> vocab { get; set; }
		public bool byte_fallback { get; set; }
	}
}