using Editor.Prism.Text.Lexer;

namespace Editor.Prism.Text;

/// <summary>What produced a fold region.</summary>
public enum FoldKind
{
	/// <summary>A <c>{ ... }</c> block.</summary>
	Braces,

	/// <summary>A top-level VFX block such as <c>PS</c> or <c>FEATURES</c>.</summary>
	Block,

	/// <summary>A <c>#if ... #endif</c> span.</summary>
	Preprocessor,

	/// <summary>A run of comment-only lines.</summary>
	Comment,

	/// <summary>An explicit <c>//region ... //endregion</c> span.</summary>
	Region
}

/// <summary>One collapsible span of lines. <see cref="StartLine"/> stays visible when collapsed.</summary>
public sealed record FoldRegion( int StartLine, int EndLine, FoldKind Kind, string Label )
{
	/// <summary>Number of lines hidden when this region is collapsed.</summary>
	public int HiddenCount => Math.Max( 0, EndLine - StartLine );

	/// <summary>Whether a line falls inside the region, header included.</summary>
	public bool Contains( int line ) => line >= StartLine && line <= EndLine;
}

/// <summary>
/// Computes and owns the fold regions of a document, and maps between document lines and the visual
/// rows the editor actually draws. With nothing collapsed the mapping is the identity and costs
/// nothing.
/// </summary>
public sealed class FoldingModel
{
	static readonly string[] s_vfxBlocks =
	{
		"HEADER", "MODES", "FEATURES", "COMMON", "VS", "PS", "GS", "CS", "PS_RENDER_STATE", "RTX"
	};

	readonly TextDocument _document;
	readonly List<FoldRegion> _regions = new();
	readonly HashSet<int> _collapsed = new();
	readonly List<int> _visible = new();

	bool _regionsDirty = true;
	bool _mappingDirty = true;
	bool _attached;

	/// <summary>Creates a folding model over a document.</summary>
	public FoldingModel( TextDocument document, IncrementalLexer lexer = null )
	{
		_document = document;
		Lexer = lexer;

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

	/// <summary>Raised when regions or the collapsed set change.</summary>
	public event Action Changed;

	/// <summary>Optional lexer, used to ignore braces inside comments and strings.</summary>
	public IncrementalLexer Lexer { get; set; }

	/// <summary>Whether folding is computed at all.</summary>
	public bool Enabled { get; set; } = true;

	/// <summary>Every fold region, ordered by start line.</summary>
	public IReadOnlyList<FoldRegion> Regions
	{
		get
		{
			EnsureRegions();
			return _regions;
		}
	}

	/// <summary>Whether at least one region is collapsed.</summary>
	public bool HasCollapsed => _collapsed.Count > 0;

	/// <summary>Number of visual rows the editor draws.</summary>
	public int VisualCount
	{
		get
		{
			EnsureMapping();

			if ( _visible.Count == 0 )
				return _document?.LineCount ?? 0;

			return _visible.Count;
		}
	}

	/// <summary>Marks the region list stale. The next query recomputes it.</summary>
	public void Invalidate()
	{
		_regionsDirty = true;
		_mappingDirty = true;
	}

	/// <summary>The region that starts on a line, if any.</summary>
	public bool TryGetRegionAt( int line, out FoldRegion region )
	{
		EnsureRegions();

		region = null;

		for ( var i = 0; i < _regions.Count; i++ )
		{
			if ( _regions[i].StartLine != line )
				continue;

			// Prefer the largest region starting here.
			if ( region is null || _regions[i].EndLine > region.EndLine )
				region = _regions[i];
		}

		return region is not null;
	}

	/// <summary>Whether a line carries a fold arrow.</summary>
	public bool IsFoldStart( int line ) => TryGetRegionAt( line, out _ );

	/// <summary>Whether the region starting on a line is collapsed.</summary>
	public bool IsCollapsed( int line ) => _collapsed.Contains( line );

	/// <summary>Whether a line is hidden inside a collapsed region.</summary>
	public bool IsHidden( int line )
	{
		if ( _collapsed.Count == 0 )
			return false;

		EnsureMapping();
		return ToVisualInternal( line, out var exact ) >= 0 && !exact;
	}

	/// <summary>Collapses the region starting on a line.</summary>
	public bool Collapse( int line )
	{
		if ( !TryGetRegionAt( line, out _ ) )
			return false;

		if ( !_collapsed.Add( line ) )
			return false;

		_mappingDirty = true;
		Changed?.Invoke();
		return true;
	}

	/// <summary>Expands the region starting on a line.</summary>
	public bool Expand( int line )
	{
		if ( !_collapsed.Remove( line ) )
			return false;

		_mappingDirty = true;
		Changed?.Invoke();
		return true;
	}

	/// <summary>Toggles the region starting on a line.</summary>
	public bool Toggle( int line ) => IsCollapsed( line ) ? Expand( line ) : Collapse( line );

	/// <summary>Collapses every region.</summary>
	public void CollapseAll()
	{
		EnsureRegions();

		_collapsed.Clear();

		for ( var i = 0; i < _regions.Count; i++ )
			_collapsed.Add( _regions[i].StartLine );

		_mappingDirty = true;
		Changed?.Invoke();
	}

	/// <summary>Expands everything.</summary>
	public void ExpandAll()
	{
		if ( _collapsed.Count == 0 )
			return;

		_collapsed.Clear();
		_mappingDirty = true;
		Changed?.Invoke();
	}

	/// <summary>Expands every collapsed region that hides a line, so the line can be revealed.</summary>
	public void EnsureVisible( int line )
	{
		if ( _collapsed.Count == 0 )
			return;

		EnsureRegions();

		var changed = false;

		for ( var i = 0; i < _regions.Count; i++ )
		{
			var region = _regions[i];

			if ( line <= region.StartLine || line > region.EndLine )
				continue;

			if ( _collapsed.Remove( region.StartLine ) )
				changed = true;
		}

		if ( !changed )
			return;

		_mappingDirty = true;
		Changed?.Invoke();
	}

	/// <summary>Maps a document line to the visual row that draws it, clamping hidden lines to their header.</summary>
	public int ToVisual( int documentLine )
	{
		if ( _collapsed.Count == 0 )
			return Math.Clamp( documentLine, 0, Math.Max( 0, (_document?.LineCount ?? 1) - 1 ) );

		EnsureMapping();
		var row = ToVisualInternal( documentLine, out _ );
		return row < 0 ? 0 : row;
	}

	/// <summary>Maps a visual row back to a document line.</summary>
	public int ToDocument( int visualRow )
	{
		var lineCount = _document?.LineCount ?? 0;

		if ( lineCount == 0 )
			return 0;

		if ( _collapsed.Count == 0 )
			return Math.Clamp( visualRow, 0, lineCount - 1 );

		EnsureMapping();

		if ( _visible.Count == 0 )
			return Math.Clamp( visualRow, 0, lineCount - 1 );

		return _visible[Math.Clamp( visualRow, 0, _visible.Count - 1 )];
	}

	/// <summary>The next visible document line in a direction, or the input when there is none.</summary>
	public int StepVisible( int documentLine, int direction )
	{
		var lineCount = _document?.LineCount ?? 0;

		if ( lineCount == 0 )
			return 0;

		if ( _collapsed.Count == 0 )
			return Math.Clamp( documentLine + direction, 0, lineCount - 1 );

		EnsureMapping();

		var row = ToVisual( documentLine ) + direction;
		row = Math.Clamp( row, 0, Math.Max( 0, _visible.Count - 1 ) );
		return ToDocument( row );
	}

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

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

	void OnDocumentChanged( TextDocument document, TextChange change )
	{
		_regionsDirty = true;
		_mappingDirty = true;

		if ( _collapsed.Count == 0 )
			return;

		var remapped = new List<int>( _collapsed.Count );

		foreach ( var line in _collapsed )
			remapped.Add( change.Map( new TextPosition( line, 0 ) ).Line );

		_collapsed.Clear();

		for ( var i = 0; i < remapped.Count; i++ )
			_collapsed.Add( remapped[i] );
	}

	int ToVisualInternal( int documentLine, out bool exact )
	{
		exact = true;

		if ( _visible.Count == 0 )
			return documentLine;

		var low = 0;
		var high = _visible.Count - 1;
		var best = -1;

		while ( low <= high )
		{
			var mid = (low + high) / 2;
			var value = _visible[mid];

			if ( value == documentLine )
				return mid;

			if ( value < documentLine )
			{
				best = mid;
				low = mid + 1;
			}
			else
			{
				high = mid - 1;
			}
		}

		exact = false;
		return best < 0 ? 0 : best;
	}

	void EnsureMapping()
	{
		if ( !_mappingDirty )
			return;

		_mappingDirty = false;
		_visible.Clear();

		var lineCount = _document?.LineCount ?? 0;

		if ( lineCount == 0 || _collapsed.Count == 0 )
			return;

		EnsureRegions();

		var hidden = new bool[lineCount];

		for ( var i = 0; i < _regions.Count; i++ )
		{
			var region = _regions[i];

			if ( !_collapsed.Contains( region.StartLine ) )
				continue;

			if ( region.StartLine < 0 || region.StartLine >= lineCount )
				continue;

			if ( hidden[region.StartLine] )
				continue;

			var last = Math.Min( region.EndLine, lineCount - 1 );

			for ( var line = region.StartLine + 1; line <= last; line++ )
				hidden[line] = true;
		}

		for ( var line = 0; line < lineCount; line++ )
		{
			if ( !hidden[line] )
				_visible.Add( line );
		}

		if ( _visible.Count == 0 )
			_visible.Add( 0 );
	}

	void EnsureRegions()
	{
		if ( !_regionsDirty )
			return;

		_regionsDirty = false;
		_mappingDirty = true;
		_regions.Clear();

		if ( !Enabled || _document is null )
			return;

		BuildBraceRegions();
		BuildPreprocessorRegions();
		BuildCommentRegions();

		_regions.Sort( static ( a, b ) =>
		{
			var c = a.StartLine.CompareTo( b.StartLine );
			return c != 0 ? c : b.EndLine.CompareTo( a.EndLine );
		} );

		// Drop collapsed markers that no longer correspond to a region.
		if ( _collapsed.Count == 0 )
			return;

		var live = new HashSet<int>();

		for ( var i = 0; i < _regions.Count; i++ )
			live.Add( _regions[i].StartLine );

		_collapsed.RemoveWhere( line => !live.Contains( line ) );
	}

	void BuildBraceRegions()
	{
		var stack = new List<int>();

		for ( var line = 0; line < _document.LineCount; line++ )
		{
			var text = _document.GetLine( line );

			for ( var column = 0; column < text.Length; column++ )
			{
				var c = text[column];

				if ( c != '{' && c != '}' )
					continue;

				if ( IsInert( line, column ) )
					continue;

				if ( c == '{' )
				{
					stack.Add( line );
					continue;
				}

				if ( stack.Count == 0 )
					continue;

				var start = stack[^1];
				stack.RemoveAt( stack.Count - 1 );

				if ( line <= start )
					continue;

				var kind = IsVfxBlockHeader( start ) ? FoldKind.Block : FoldKind.Braces;
				_regions.Add( new FoldRegion( start, line, kind, MakeLabel( start ) ) );
			}
		}
	}

	void BuildPreprocessorRegions()
	{
		var stack = new List<int>();

		for ( var line = 0; line < _document.LineCount; line++ )
		{
			var text = _document.GetLine( line ).TrimStart();

			if ( text.Length == 0 || text[0] != '#' )
				continue;

			var directive = ReadDirective( text );

			switch ( directive )
			{
				case "if":
				case "ifdef":
				case "ifndef":
					stack.Add( line );
					break;

				case "endif":
					if ( stack.Count == 0 )
						break;

					var start = stack[^1];
					stack.RemoveAt( stack.Count - 1 );

					if ( line > start + 1 )
						_regions.Add( new FoldRegion( start, line, FoldKind.Preprocessor, MakeLabel( start ) ) );

					break;
			}
		}
	}

	void BuildCommentRegions()
	{
		var lexer = Lexer;

		if ( lexer is null )
			return;

		var runStart = -1;

		for ( var line = 0; line <= _document.LineCount; line++ )
		{
			var isComment = line < _document.LineCount && IsCommentOnlyLine( lexer, line );

			if ( isComment )
			{
				if ( runStart < 0 )
					runStart = line;

				continue;
			}

			if ( runStart >= 0 && line - 1 > runStart )
				_regions.Add( new FoldRegion( runStart, line - 1, FoldKind.Comment, MakeLabel( runStart ) ) );

			runStart = -1;
		}
	}

	bool IsCommentOnlyLine( IncrementalLexer lexer, int line )
	{
		var text = _document.GetLine( line );

		if ( text.Trim().Length == 0 )
			return false;

		var tokens = lexer.GetCachedTokens( line );

		if ( tokens.Count == 0 )
			return false;

		for ( var i = 0; i < tokens.Count; i++ )
		{
			switch ( tokens[i].Kind )
			{
				case TokenKind.Comment:
				case TokenKind.DocComment:
				case TokenKind.Whitespace:
					continue;
				default:
					return false;
			}
		}

		return true;
	}

	bool IsInert( int line, int column )
	{
		var lexer = Lexer;

		if ( lexer is null )
			return false;

		// Cached only: rebuilding regions must never force a whole-document lex.
		if ( !lexer.TryGetCachedToken( new TextPosition( line, column ), out var token ) )
			return false;

		switch ( token.Kind )
		{
			case TokenKind.Comment:
			case TokenKind.DocComment:
			case TokenKind.String:
			case TokenKind.IncludePath:
				return true;
			default:
				return false;
		}
	}

	bool IsVfxBlockHeader( int line )
	{
		var text = _document.GetLine( line ).Trim();

		if ( text.EndsWith( "{", StringComparison.Ordinal ) )
			text = text[..^1].Trim();

		if ( text.Length == 0 && line > 0 )
			text = _document.GetLine( line - 1 ).Trim();

		for ( var i = 0; i < s_vfxBlocks.Length; i++ )
		{
			if ( string.Equals( text, s_vfxBlocks[i], StringComparison.Ordinal ) )
				return true;
		}

		return false;
	}

	string MakeLabel( int line )
	{
		var text = _document.GetLine( line ).Trim();

		if ( text.Length == 0 && line > 0 )
			text = _document.GetLine( line - 1 ).Trim();

		if ( text.Length > 60 )
			text = text[..60] + "…";

		return text;
	}

	static string ReadDirective( string trimmedLine )
	{
		var index = 1;

		while ( index < trimmedLine.Length && char.IsWhiteSpace( trimmedLine[index] ) )
			index++;

		var start = index;

		while ( index < trimmedLine.Length && char.IsLetter( trimmedLine[index] ) )
			index++;

		return trimmedLine[start..index];
	}
}