using Editor.Prism.Core;
using Editor.Prism.Text.Lexer;

namespace Editor.Prism.Text;

/// <summary>
/// Drives an <see cref="ILexer"/> over a <see cref="TextDocument"/> one line at a time and caches the
/// result. After an edit at line <i>N</i> the lexer re-runs from <i>N</i> forward until a line's
/// computed entry state equals the state it already had — typically one to three lines per keystroke,
/// which is what keeps highlighting free even in a ten-thousand-line shader.
/// </summary>
public sealed class IncrementalLexer
{
	struct LineInfo
	{
		public Token[] Tokens;
		public LexState Entry;
		public LexState Exit;
		public bool Valid;
		public bool Dirty;
	}

	static readonly Token[] s_noTokens = Array.Empty<Token>();

	readonly TextDocument _document;
	readonly List<LineInfo> _lines = new();
	readonly List<Token> _scratch = new();

	ILexer _lexer;
	int _firstDirty;
	int _lastDirty = -1;
	bool _reportedFailure;
	bool _attached;

	/// <summary>Creates a lexer driver bound to a document.</summary>
	public IncrementalLexer( TextDocument document, ILexer lexer = null )
	{
		_document = document;
		_lexer = lexer;

		if ( _document is not null )
		{
			_document.Changed += OnDocumentChanged;
			_attached = true;
		}

		Rebuild();
	}

	/// <summary>The document being lexed.</summary>
	public TextDocument Document => _document;

	/// <summary>
	/// The active lexer. Assigning a different one invalidates every cached line. Null is legal and
	/// produces one plain token per line, so the editor still renders.
	/// </summary>
	public ILexer Lexer
	{
		get => _lexer;
		set
		{
			if ( ReferenceEquals( _lexer, value ) )
				return;

			_lexer = value;
			_reportedFailure = false;
			Rebuild();
		}
	}

	/// <summary>The language id the active lexer reports, or <c>"text"</c>.</summary>
	public string Language => _lexer?.Language ?? "text";

	/// <summary>Lines longer than this are not tokenised, as a guard against pathological input.</summary>
	public int MaxLineLength { get; set; } = 20000;

	/// <summary>Whether any line still needs re-lexing.</summary>
	public bool IsDirty => _firstDirty < _lines.Count;

	/// <summary>The first line that still needs re-lexing.</summary>
	public int FirstDirtyLine => _firstDirty;

	/// <summary>
	/// Resolves a language or file extension through <see cref="Lexers.For"/> and installs it.
	/// Degrades to no lexer when the language is unknown.
	/// </summary>
	public void SetLanguage( string languageOrExtension )
	{
		ILexer resolved = null;

		PrismLog.Guard( "Prism.Text: resolve lexer", () =>
		{
			resolved = Lexers.For( languageOrExtension );
		} );

		Lexer = resolved;
	}

	/// <summary>Marks every line from <paramref name="fromLine"/> onwards as needing a re-lex.</summary>
	public void Invalidate( int fromLine = 0 )
	{
		fromLine = Math.Max( 0, fromLine );

		for ( var i = fromLine; i < _lines.Count; i++ )
		{
			var info = _lines[i];
			info.Valid = false;
			info.Dirty = true;
			_lines[i] = info;
		}

		_firstDirty = Math.Min( _firstDirty, fromLine );
		_lastDirty = _lines.Count - 1;
	}

	/// <summary>Tokens for one line, lexing on demand. Never null.</summary>
	public IReadOnlyList<Token> GetTokens( int line )
	{
		if ( line < 0 || line >= _lines.Count )
			return s_noTokens;

		EnsureLexed( line );

		return _lines[line].Tokens ?? s_noTokens;
	}

	/// <summary>The lexer state a line begins in.</summary>
	public LexState GetEntryState( int line )
	{
		if ( line <= 0 )
			return LexState.Default;

		if ( line > _lines.Count )
			return LexState.Default;

		EnsureLexed( line - 1 );
		return _lines[line - 1].Exit;
	}

	/// <summary>The lexer state a line ends in.</summary>
	public LexState GetExitState( int line )
	{
		if ( line < 0 || line >= _lines.Count )
			return LexState.Default;

		EnsureLexed( line );
		return _lines[line].Exit;
	}

	/// <summary>
	/// Tokens for a line, but only when they are already cached. Never forces a re-lex, so a caller
	/// that scans thousands of lines — bracket matching, folding — cannot turn one keystroke into a
	/// whole-document lex. Returns an empty list for lines that have not converged yet.
	/// </summary>
	public IReadOnlyList<Token> GetCachedTokens( int line )
	{
		if ( line < 0 || line >= _lines.Count || line >= _firstDirty )
			return s_noTokens;

		return _lines[line].Tokens ?? s_noTokens;
	}

	/// <summary>The cached token covering a position, if the line has already been lexed.</summary>
	public bool TryGetCachedToken( TextPosition position, out Token token )
	{
		token = default;

		var tokens = GetCachedTokens( position.Line );

		for ( var i = 0; i < tokens.Count; i++ )
		{
			var candidate = tokens[i];

			if ( position.Column >= candidate.Start && position.Column < candidate.Start + candidate.Length )
			{
				token = candidate;
				return true;
			}
		}

		return false;
	}

	/// <summary>The token covering a position, if any.</summary>
	public bool TryGetToken( TextPosition position, out Token token )
	{
		token = default;

		var tokens = GetTokens( position.Line );

		for ( var i = 0; i < tokens.Count; i++ )
		{
			var candidate = tokens[i];

			if ( position.Column >= candidate.Start && position.Column < candidate.Start + candidate.Length )
			{
				token = candidate;
				return true;
			}
		}

		return false;
	}

	/// <summary>The token class at a position, or <see cref="TokenKind.None"/>.</summary>
	public TokenKind KindAt( TextPosition position ) => TryGetToken( position, out var token ) ? token.Kind : TokenKind.None;

	/// <summary>
	/// Whether a position sits inside a comment or a string, in which case completion, bracket
	/// matching and occurrence highlighting all stand down.
	/// </summary>
	public bool IsInert( TextPosition position )
	{
		if ( TryGetToken( position, out var token ) )
		{
			switch ( token.Kind )
			{
				case TokenKind.Comment:
				case TokenKind.DocComment:
				case TokenKind.String:
				case TokenKind.IncludePath:
					return true;
			}
		}

		// A caret sitting exactly at the end of a comment token still counts as inert.
		if ( position.Column > 0 && TryGetToken( new TextPosition( position.Line, position.Column - 1 ), out var before ) )
		{
			switch ( before.Kind )
			{
				case TokenKind.Comment:
				case TokenKind.DocComment:
					return true;
			}
		}

		if ( _lexer is null )
			return false;

		return PrismLog.Guard( "Prism.Text: IsInert", () => _lexer.IsInert( GetEntryState( position.Line ) ), false );
	}

	/// <summary>
	/// Processes pending lines, at most <paramref name="budget"/> of them. Returns how many were
	/// lexed. Call from an idle tick so a large paste never blocks a frame.
	/// </summary>
	public int Sync( int budget = 2000 )
	{
		if ( !IsDirty )
			return 0;

		return LexFrom( _firstDirty, -1, Math.Max( 1, budget ) );
	}

	/// <summary>Unsubscribes from the document. Call when the editor is destroyed.</summary>
	public void Detach()
	{
		if ( !_attached || _document is null )
			return;

		_document.Changed -= OnDocumentChanged;
		_attached = false;
	}

	void EnsureLexed( int line )
	{
		if ( line < 0 || line >= _lines.Count )
			return;

		if ( _firstDirty > line )
			return;

		LexFrom( _firstDirty, line, int.MaxValue );
	}

	int LexFrom( int start, int requiredLine, int budget )
	{
		var i = Math.Max( 0, start );
		var processed = 0;

		while ( i < _lines.Count )
		{
			var entry = i == 0 ? LexState.Default : _lines[i - 1].Exit;
			var info = _lines[i];

			if ( info.Valid && !info.Dirty && info.Entry == entry )
			{
				// This line's cached tokens are still right, so every line after it that nothing
				// touched is right too — but a second edit further down leaves its own dirty island,
				// and stopping here would leave that island showing stale colours until the next edit
				// happened to reach it. Multi-cursor typing and replace-all both produce exactly that.
				var next = NextDirty( i + 1 );

				if ( next < 0 )
					break;

				i = next;
				continue;
			}

			LexLine( i, entry );
			processed++;
			i++;

			if ( processed >= budget && i > requiredLine )
			{
				_firstDirty = i;
				return processed;
			}
		}

		_firstDirty = _lines.Count;
		_lastDirty = -1;

		return processed;
	}

	/// <summary>The next line at or after <paramref name="from"/> that still needs a re-lex, or -1.</summary>
	int NextDirty( int from )
	{
		var last = Math.Min( _lastDirty, _lines.Count - 1 );

		for ( var i = Math.Max( 0, from ); i <= last; i++ )
		{
			if ( !_lines[i].Valid || _lines[i].Dirty )
				return i;
		}

		return -1;
	}

	void LexLine( int index, LexState entry )
	{
		var text = _document is null ? string.Empty : _document.GetLine( index );

		var info = new LineInfo
		{
			Entry = entry,
			Exit = entry,
			Valid = true,
			Dirty = false,
			Tokens = s_noTokens
		};

		if ( _lexer is null )
		{
			if ( text.Length > 0 )
				info.Tokens = new[] { new Token( 0, text.Length, TokenKind.None ) };

			info.Exit = LexState.Default;
			_lines[index] = info;
			return;
		}

		if ( text.Length > MaxLineLength )
		{
			info.Tokens = new[] { new Token( 0, text.Length, TokenKind.None ) };
			_lines[index] = info;
			return;
		}

		_scratch.Clear();

		try
		{
			info.Exit = _lexer.Lex( text, entry, _scratch );
		}
		catch ( Exception e )
		{
			_scratch.Clear();

			if ( text.Length > 0 )
				_scratch.Add( new Token( 0, text.Length, TokenKind.Error ) );

			info.Exit = entry;

			if ( !_reportedFailure )
			{
				_reportedFailure = true;
				PrismLog.Error( e, $"Prism.Text: {Language} lexer threw on line {index + 1}; that line is shown unhighlighted" );
			}
		}

		info.Tokens = _scratch.Count > 0 ? _scratch.ToArray() : s_noTokens;
		_lines[index] = info;
	}

	void OnDocumentChanged( TextDocument document, TextChange change )
	{
		var start = change.Removed.Start.Line;
		var oldEnd = change.Removed.End.Line;
		var newEnd = change.InsertedEnd.Line;

		if ( start >= 0 && start < _lines.Count && oldEnd >= start )
		{
			var oldCount = oldEnd - start + 1;
			var newCount = newEnd - start + 1;

			if ( oldEnd < _lines.Count )
			{
				if ( newCount < oldCount )
				{
					_lines.RemoveRange( start + newCount, oldCount - newCount );
				}
				else if ( newCount > oldCount )
				{
					for ( var i = 0; i < newCount - oldCount; i++ )
						_lines.Insert( start + oldCount, default );
				}
			}
		}

		if ( _document is null || _lines.Count != _document.LineCount )
		{
			Rebuild();
			return;
		}

		var last = Math.Min( newEnd, _lines.Count - 1 );

		for ( var i = Math.Max( 0, start ); i <= last; i++ )
		{
			var info = _lines[i];
			info.Valid = false;
			info.Dirty = true;
			_lines[i] = info;
		}

		_firstDirty = Math.Min( _firstDirty, Math.Max( 0, start ) );
		_lastDirty = Math.Max( _lastDirty, last );
	}

	void Rebuild()
	{
		_lines.Clear();

		var count = _document?.LineCount ?? 0;

		for ( var i = 0; i < count; i++ )
			_lines.Add( default );

		_firstDirty = 0;
		_lastDirty = count - 1;
	}
}