using System.Text;

namespace Nbt;

internal sealed class QuotedWriter(TextWriter baseWriter, char quoteChar, char escapeChar, bool leaveOpen = false) : TextWriter
{
    private readonly TextWriter baseWriter = baseWriter;
    private readonly char quoteChar = quoteChar;
    private readonly char escapeChar = escapeChar;
    private readonly bool leaveOpen = leaveOpen;
    private bool hasWrittenOneChar = false;

    public QuotedWriter(StringBuilder sb, char quoteChar, char escapeChar) : this(new StringWriter(sb), quoteChar, escapeChar) { }

    public override Encoding Encoding => baseWriter.Encoding;

    public override void Write(char value)
    {
        if (!hasWrittenOneChar)
        {
            baseWriter.Write(quoteChar);
            hasWrittenOneChar = true;
        }

        foreach (var c in Escape(value))
        {
            baseWriter.Write(c);
        }
    }

    public override async Task WriteAsync(char value)
    {
        if (!hasWrittenOneChar)
        {
            await baseWriter.WriteAsync(quoteChar);
            hasWrittenOneChar = true;
        }

        foreach (var c in Escape(value))
        {
            await baseWriter.WriteAsync(c);
        }
    }

    private IEnumerable<char> Escape(char c)
    {
        var escape = false;

        switch (c)
        {
            case '\t':
                escape = true;
                c = 't';
                break;

            case '\r':
                escape = true;
                c = 'r';
                break;

            case '\n':
                escape = true;
                c = 'n';
                break;

            default:
                if (c == quoteChar || c == escapeChar)
                {
                    escape = true;
                }
                break;
        }

        if (escape)
        {
            yield return escapeChar;
        }

        yield return c;
    }

    protected override void Dispose(bool disposing)
    {
        try
        {
            if (disposing)
            {
                if (!hasWrittenOneChar)
                {
                    baseWriter.Write(quoteChar);
                }

                baseWriter.Write(quoteChar);

                if (!leaveOpen)
                {
                    baseWriter.Dispose();
                }
            }
        }
        finally
        {
            base.Dispose(disposing);
        }
    }

    public override async ValueTask DisposeAsync()
    {
        try
        {
            if (!hasWrittenOneChar)
            {
                await baseWriter.WriteAsync(quoteChar);
            }

            await baseWriter.WriteAsync(quoteChar);

            if (!leaveOpen)
            {
                await baseWriter.DisposeAsync();
            }
        }
        finally
        {
            await base.DisposeAsync();
        }
    }
}