nub-lang/src/compiler/Nub.Lang/Frontend/Diagnostics/Diagnostic.cs

using System.Text;
using Nub.Lang.Frontend.Lexing;
using Nub.Lang.Frontend.Parsing;

namespace Nub.Lang.Frontend.Diagnostics;

public enum DiagnosticSeverity
{
    Info,
    Warning,
    Error
}

public class Diagnostic
{
    public class DiagnosticBuilder
    {
        private readonly DiagnosticSeverity _severity;
        private readonly string _message;
        private SourceText? _sourceFile;
        private SourceSpan? _span;
        private string? _help;

        public DiagnosticBuilder(DiagnosticSeverity severity, string message)
        {
            _severity = severity;
            _message = message;
        }

        public DiagnosticBuilder At(Token token)
        {
            _sourceFile = token.SourceText;
            _span = SourceLocationCalculator.GetSpan(token);
            return this;
        }

        public DiagnosticBuilder At(Node node)
        {
            if (!node.Tokens.Any())
            {
                throw new ArgumentException("Node has no tokens", nameof(node));
            }

            _sourceFile = node.Tokens[0].SourceText;
            _span = SourceLocationCalculator.GetSpan(node);
            return this;
        }

        public DiagnosticBuilder At(SourceText sourceText, SourceSpan span)
        {
            _sourceFile = sourceText;
            _span = span;
            return this;
        }

        public DiagnosticBuilder WithHelp(string help)
        {
            _help = help;
            return this;
        }

        public Diagnostic Build() => new(_severity, _message, _sourceFile, _span, _help);
    }

    public static DiagnosticBuilder Error(string message) => new(DiagnosticSeverity.Error, message);
    public static DiagnosticBuilder Warning(string message) => new(DiagnosticSeverity.Warning, message);
    public static DiagnosticBuilder Info(string message) => new(DiagnosticSeverity.Info, message);

    public DiagnosticSeverity Severity { get; }
    public string Message { get; }
    public SourceText? SourceFile { get; }
    public SourceSpan? Span { get; }
    public string? Help { get; }

    private Diagnostic(DiagnosticSeverity severity, string message, SourceText? sourceFile, SourceSpan? span, string? help)
    {
        Severity = severity;
        Message = message;
        SourceFile = sourceFile;
        Span = span;
        Help = help;
    }

    public string Format()
    {
        var sb = new StringBuilder();

        var severityText = GetSeverityText(Severity);
        sb.Append(severityText);

        if (SourceFile.HasValue)
        {
            var locationText = $" at {SourceFile.Value.Path}:{Span}";
            sb.Append(ConsoleColors.Colorize(locationText, ConsoleColors.Gray));
        }

        sb.Append(": ");
        sb.AppendLine(ConsoleColors.Colorize(Message, ConsoleColors.BrightWhite));

        if (SourceFile.HasValue && Span.HasValue)
        {
            AppendSourceContext(sb, SourceFile.Value, Span.Value);
        }

        if (!string.IsNullOrEmpty(Help))
        {
            sb.AppendLine();
            var helpText = $"help: {Help}";
            sb.AppendLine(ConsoleColors.Colorize(helpText, ConsoleColors.Cyan));
        }

        return sb.ToString();
    }

    private static string GetSeverityText(DiagnosticSeverity severity)
    {
        return severity switch
        {
            DiagnosticSeverity.Error => ConsoleColors.Colorize("error", ConsoleColors.Bold + ConsoleColors.Red),
            DiagnosticSeverity.Warning => ConsoleColors.Colorize("warning", ConsoleColors.Bold + ConsoleColors.Yellow),
            DiagnosticSeverity.Info => ConsoleColors.Colorize("info", ConsoleColors.Bold + ConsoleColors.Blue),
            _ => throw new ArgumentOutOfRangeException(nameof(severity), severity, "Unknown diagnostic severity")
        };
    }

    private static void AppendSourceContext(StringBuilder sb, SourceText sourceText, SourceSpan span)
    {
        var lines = sourceText.Content.Split('\n');
        var startLine = span.Start.Line;
        var endLine = span.End.Line;

        const int contextLines = 3;

        if (startLine < 1 || startLine > lines.Length || endLine < 1 || endLine > lines.Length)
        {
            return;
        }

        var maxLineNum = Math.Min(endLine + contextLines, lines.Length);
        var lineNumWidth = maxLineNum.ToString().Length;

        var contextStart = Math.Max(1, startLine - contextLines);

        for (var lineNum = contextStart; lineNum < startLine; lineNum++)
        {
            AppendContextLine(sb, lineNum, lines[lineNum - 1], lineNumWidth);
        }

        for (var lineNum = startLine; lineNum <= endLine; lineNum++)
        {
            var line = lines[lineNum - 1];
            AppendContextLine(sb, lineNum, line, lineNumWidth);
            AppendErrorIndicators(sb, span, lineNum, line, lineNumWidth);
        }

        var contextEnd = Math.Min(lines.Length, endLine + contextLines);
        for (var lineNum = endLine + 1; lineNum <= contextEnd; lineNum++)
        {
            AppendContextLine(sb, lineNum, lines[lineNum - 1], lineNumWidth);
        }
    }

    private static void AppendContextLine(StringBuilder sb, int lineNum, string line, int lineNumWidth)
    {
        var lineNumStr = lineNum.ToString().PadLeft(lineNumWidth);
        sb.Append(ConsoleColors.Colorize(lineNumStr, ConsoleColors.Gray));
        sb.Append(ConsoleColors.Colorize(" | ", ConsoleColors.Gray));
        sb.AppendLine(line);
    }

    private static void AppendErrorIndicators(StringBuilder sb, SourceSpan span, int lineNum, string line, int lineNumWidth)
    {
        sb.Append(new string(' ', lineNumWidth + 3));
        var indicators = GetIndicatorsForLine(span, lineNum, line);
        sb.AppendLine(ConsoleColors.Colorize(indicators, ConsoleColors.Red));
    }

    private static string GetIndicatorsForLine(SourceSpan span, int lineNum, string line)
    {
        const char indicator = '^';

        if (lineNum == span.Start.Line && lineNum == span.End.Line)
        {
            var startCol = Math.Max(0, span.Start.Column - 1);
            var endCol = Math.Min(line.Length, span.End.Column - 1);
            var length = Math.Max(1, endCol - startCol);

            return new string(' ', startCol) + new string(indicator, length);
        }

        if (lineNum == span.Start.Line)
        {
            var startCol = Math.Max(0, span.Start.Column - 1);
            return new string(' ', startCol) + new string(indicator, Math.Max(0, line.Length - startCol));
        }

        if (lineNum == span.End.Line)
        {
            var endCol = Math.Min(line.Length, span.End.Column - 1);
            return new string(indicator, Math.Max(0, endCol));
        }

        return new string(indicator, line.Length);
    }
}