using System.Text;
using System.Text.RegularExpressions;
using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public sealed class Gpt2ByteBpeTokenizer
{
	private const string PairSeparator = "\0";
	private const string Gpt2Pattern = @"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+";

	private static readonly Regex TokenPattern = new(
		Gpt2Pattern,
		RegexOptions.CultureInvariant );

	private readonly Dictionary<string, int> _vocabulary;
	private readonly string[] _tokensById;
	private readonly Dictionary<string, int> _mergeRanks;
	private readonly Dictionary<string, string[]> _bpeCache = new( StringComparer.Ordinal );
	private readonly char[] _byteEncoder = new char[256];
	private readonly Dictionary<char, byte> _byteDecoder = new();
	private readonly List<Gpt2AddedTokenDocument> _addedTokens;
	private readonly Dictionary<int, Gpt2AddedTokenDocument> _addedTokensById;

	public int VocabularySize => _tokensById.Length;
	public int MergeCount => _mergeRanks.Count;

	private Gpt2ByteBpeTokenizer( Gpt2TokenizerDocument document )
	{
		ValidateDocument( document );
		_vocabulary = new Dictionary<string, int>( document.Model.Vocabulary, StringComparer.Ordinal );
		_tokensById = new string[_vocabulary.Count];

		foreach ( KeyValuePair<string, int> item in _vocabulary )
		{
			if ( item.Value < 0 || item.Value >= _tokensById.Length )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Token '{item.Key}' has ID {item.Value}, outside " +
					$"[0,{_tokensById.Length})." );
			}
			if ( _tokensById[item.Value] is not null )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Token ID {item.Value} is assigned to both " +
					$"'{_tokensById[item.Value]}' and '{item.Key}'." );
			}
			_tokensById[item.Value] = item.Key;
		}

		for ( int id = 0; id < _tokensById.Length; id++ )
		{
			if ( _tokensById[id] is null )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Tokenizer vocabulary has no token for ID {id}." );
			}
		}

		_mergeRanks = new Dictionary<string, int>( document.Model.Merges.Count, StringComparer.Ordinal );
		for ( int rank = 0; rank < document.Model.Merges.Count; rank++ )
		{
			string[] pair = document.Model.Merges[rank];
			if ( pair is null || pair.Length != 2 || pair[0] is null || pair[1] is null )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Tokenizer merge rank {rank} must contain exactly two strings." );
			}

			string key = PairKey( pair[0], pair[1] );
			if ( !_mergeRanks.TryAdd( key, rank ) )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Duplicate tokenizer merge pair at rank {rank}: " +
					$"'{pair[0]}' + '{pair[1]}'." );
			}
		}

		_addedTokens = document.AddedTokens?
			.Where( token => token is not null && !string.IsNullOrEmpty( token.Content ) )
			.OrderByDescending( token => token.Content.Length )
			.ToList() ?? new List<Gpt2AddedTokenDocument>();
		_addedTokensById = _addedTokens.ToDictionary( token => token.Id );

		BuildByteMaps();
	}

	public static Gpt2ByteBpeTokenizer LoadFromMounted( string path )
	{
		if ( !FileSystem.Mounted.FileExists( path ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer file '{path}' is not present in FileSystem.Mounted." );
		}

		FastTimer timer = FastTimer.StartNew();
		Gpt2TokenizerDocument document = FileSystem.Mounted.ReadJson<Gpt2TokenizerDocument>( path );
		if ( document is null )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer file '{path}' deserialized to null." );
		}

		Gpt2ByteBpeTokenizer tokenizer = new( document );
		LlmLog.Info(
			"TOKEN",
			$"Loaded GPT-2 byte-level BPE tokenizer with {tokenizer.VocabularySize:N0} tokens " +
			$"and {tokenizer.MergeCount:N0} merges in {timer.ElapsedMilliSeconds:N2} ms." );
		return tokenizer;
	}

	public static Gpt2ByteBpeTokenizer LoadFromJson(
		string json,
		string sourceLabel = "in-memory tokenizer JSON" )
	{
		if ( string.IsNullOrWhiteSpace( json ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer JSON from '{sourceLabel}' is empty." );
		}

		FastTimer timer = FastTimer.StartNew();
		Gpt2TokenizerDocument document = Json.Deserialize<Gpt2TokenizerDocument>( json );
		if ( document is null )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer JSON from '{sourceLabel}' deserialized to null." );
		}

		Gpt2ByteBpeTokenizer tokenizer = new( document );
		LlmLog.Info(
			"TOKEN",
			$"Loaded packaged GPT-2 byte-level BPE tokenizer with " +
			$"{tokenizer.VocabularySize:N0} tokens and {tokenizer.MergeCount:N0} merges " +
			$"in {timer.ElapsedMilliSeconds:N2} ms." );
		return tokenizer;
	}

	public int[] Encode( string text )
	{
		if ( text is null )
		{
			throw new ArgumentNullException( nameof( text ) );
		}

		FastTimer timer = FastTimer.StartNew();
		List<int> result = new();
		int position = 0;

		while ( position < text.Length )
		{
			Gpt2AddedTokenDocument nextSpecial = null;
			int nextSpecialIndex = int.MaxValue;
			foreach ( Gpt2AddedTokenDocument token in _addedTokens )
			{
				int index = text.IndexOf( token.Content, position, StringComparison.Ordinal );
				if ( index >= 0 && index < nextSpecialIndex )
				{
					nextSpecial = token;
					nextSpecialIndex = index;
				}
			}

			if ( nextSpecial is null )
			{
				EncodeOrdinary( text[position..], result );
				break;
			}

			if ( nextSpecialIndex > position )
			{
				EncodeOrdinary( text[position..nextSpecialIndex], result );
			}
			result.Add( nextSpecial.Id );
			position = nextSpecialIndex + nextSpecial.Content.Length;
		}

		LlmLog.Trace(
			"TOKEN",
			$"Encoded chars={text.Length:N0} tokens={result.Count:N0} " +
			$"ids=[{string.Join( ",", result )}] in {timer.ElapsedMilliSeconds:N3} ms." );
		return result.ToArray();
	}

	public string Decode( IReadOnlyList<int> tokenIds )
	{
		if ( tokenIds is null )
		{
			throw new ArgumentNullException( nameof( tokenIds ) );
		}

		StringBuilder output = new();
		List<byte> bytes = new();

		for ( int index = 0; index < tokenIds.Count; index++ )
		{
			int tokenId = tokenIds[index];
			ValidateTokenId( tokenId, index );

			if ( _addedTokensById.TryGetValue( tokenId, out Gpt2AddedTokenDocument addedToken ) )
			{
				FlushBytes( bytes, output );
				output.Append( addedToken.Content );
				continue;
			}

			string token = _tokensById[tokenId];
			for ( int characterIndex = 0; characterIndex < token.Length; characterIndex++ )
			{
				char character = token[characterIndex];
				if ( !_byteDecoder.TryGetValue( character, out byte value ) )
				{
					throw new InvalidOperationException(
						$"[LLM:ERROR] Token ID {tokenId} contains character U+{(int)character:X4} " +
						"which is absent from the GPT-2 byte decoder." );
				}
				bytes.Add( value );
			}
		}

		FlushBytes( bytes, output );
		string decoded = output.ToString();
		LlmLog.Trace(
			"TOKEN",
			$"Decoded tokens={tokenIds.Count:N0} chars={decoded.Length:N0}." );
		return decoded;
	}

	private void EncodeOrdinary( string text, List<int> output )
	{
		MatchCollection matches = TokenPattern.Matches( text );
		foreach ( Match match in matches )
		{
			byte[] utf8 = Encoding.UTF8.GetBytes( match.Value );
			StringBuilder encoded = new( utf8.Length );
			for ( int index = 0; index < utf8.Length; index++ )
			{
				encoded.Append( _byteEncoder[utf8[index]] );
			}

			string[] pieces = ApplyBpe( encoded.ToString() );
			for ( int pieceIndex = 0; pieceIndex < pieces.Length; pieceIndex++ )
			{
				string piece = pieces[pieceIndex];
				if ( !_vocabulary.TryGetValue( piece, out int tokenId ) )
				{
					throw new InvalidOperationException(
						$"[LLM:ERROR] GPT-2 BPE produced token '{Escape( piece )}' " +
						$"which is absent from the {VocabularySize:N0}-entry vocabulary." );
				}
				output.Add( tokenId );
			}
		}
	}

	private string[] ApplyBpe( string token )
	{
		if ( _bpeCache.TryGetValue( token, out string[] cached ) )
		{
			return cached;
		}

		List<string> word = new( token.Length );
		for ( int index = 0; index < token.Length; index++ )
		{
			word.Add( token[index].ToString() );
		}

		while ( word.Count > 1 )
		{
			int bestRank = int.MaxValue;
			string bestLeft = null;
			string bestRight = null;

			for ( int index = 0; index < word.Count - 1; index++ )
			{
				if ( _mergeRanks.TryGetValue( PairKey( word[index], word[index + 1] ), out int rank ) &&
					rank < bestRank )
				{
					bestRank = rank;
					bestLeft = word[index];
					bestRight = word[index + 1];
				}
			}

			if ( bestLeft is null )
			{
				break;
			}

			List<string> merged = new( word.Count );
			int position = 0;
			while ( position < word.Count )
			{
				if ( position < word.Count - 1 &&
					word[position] == bestLeft && word[position + 1] == bestRight )
				{
					merged.Add( bestLeft + bestRight );
					position += 2;
				}
				else
				{
					merged.Add( word[position] );
					position++;
				}
			}
			word = merged;
		}

		string[] result = word.ToArray();
		_bpeCache[token] = result;
		return result;
	}

	private void ValidateTokenId( int tokenId, int position )
	{
		if ( tokenId < 0 || tokenId >= _tokensById.Length )
		{
			throw new IndexOutOfRangeException(
				$"[LLM:ERROR] Token ID {tokenId} at sequence index {position} is outside " +
				$"[0,{_tokensById.Length})." );
		}
	}

	private static void ValidateDocument( Gpt2TokenizerDocument document )
	{
		if ( document is null )
		{
			throw new ArgumentNullException( nameof( document ) );
		}
		if ( document.Version != "1.0" )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer JSON version expected 1.0, found '{document.Version}'." );
		}
		if ( document.Normalizer is not null )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Exported tokenizer unexpectedly defines a normalizer." );
		}
		if ( document.PreTokenizer is null || document.PreTokenizer.Type != "ByteLevel" ||
			document.PreTokenizer.AddPrefixSpace || !document.PreTokenizer.UseRegex )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Tokenizer requires ByteLevel pre-tokenization with " +
				"add_prefix_space=false and use_regex=true." );
		}
		if ( document.Decoder is null || document.Decoder.Type != "ByteLevel" )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Tokenizer requires a ByteLevel decoder." );
		}
		if ( document.Model is null || document.Model.Type != "BPE" )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Tokenizer model expected BPE, found '{document.Model?.Type}'." );
		}
		if ( document.Model.Dropout is not null || document.Model.ByteFallback ||
			document.Model.IgnoreMerges || document.Model.FuseUnknown ||
			!string.IsNullOrEmpty( document.Model.UnknownToken ) ||
			!string.IsNullOrEmpty( document.Model.ContinuingSubwordPrefix ) ||
			!string.IsNullOrEmpty( document.Model.EndOfWordSuffix ) )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Exported tokenizer uses unsupported BPE options." );
		}
		if ( document.Model.Vocabulary is null || document.Model.Vocabulary.Count == 0 )
		{
			throw new InvalidOperationException( "[LLM:ERROR] Tokenizer vocabulary is empty." );
		}
		if ( document.Model.Merges is null )
		{
			throw new InvalidOperationException( "[LLM:ERROR] Tokenizer merges array is missing." );
		}
	}

	private void BuildByteMaps()
	{
		List<int> mappedBytes = new();
		for ( int value = 0; value < 256; value++ )
		{
			if ( IsDirectByte( value ) )
			{
				mappedBytes.Add( value );
			}
		}

		List<int> mappedCodePoints = new( mappedBytes );
		int extra = 0;
		for ( int value = 0; value < 256; value++ )
		{
			if ( !IsDirectByte( value ) )
			{
				mappedBytes.Add( value );
				mappedCodePoints.Add( 256 + extra );
				extra++;
			}
		}

		for ( int index = 0; index < mappedBytes.Count; index++ )
		{
			byte source = (byte)mappedBytes[index];
			char encoded = (char)mappedCodePoints[index];
			_byteEncoder[source] = encoded;
			_byteDecoder.Add( encoded, source );
		}
	}

	private static bool IsDirectByte( int value )
	{
		return (value >= 33 && value <= 126) ||
			(value >= 161 && value <= 172) ||
			(value >= 174 && value <= 255);
	}

	private static string PairKey( string left, string right )
	{
		return left + PairSeparator + right;
	}

	private static void FlushBytes( List<byte> bytes, StringBuilder output )
	{
		if ( bytes.Count == 0 )
		{
			return;
		}
		output.Append( Encoding.UTF8.GetString( bytes.ToArray() ) );
		bytes.Clear();
	}

	private static string Escape( string value )
	{
		StringBuilder result = new();
		for ( int index = 0; index < value.Length; index++ )
		{
			char character = value[index];
			if ( char.IsControl( character ) )
			{
				result.Append( $"\\u{(int)character:X4}" );
			}
			else
			{
				result.Append( character );
			}
		}
		return result.ToString();
	}
}