/*
 * A faster replacement for some subset of std.stdio
 * Copyright © 2022, Siarhei Siamashka
 *
 * Boost Software License - Version 1.0 - August 17th, 2003
 *
 * Permission is hereby granted, free of charge, to any person or organization
 * obtaining a copy of the software and accompanying documentation covered by
 * this license (the "Software") to use, reproduce, display, distribute,
 * execute, and transmit the Software, and to prepare derivative works of the
 * Software, and to permit third-parties to whom the Software is furnished to
 * do so, all subject to the following:
 *
 * The copyright notices in the Software and this entire statement, including
 * the above license grant, this restriction and the following disclaimer,
 * must be included in all copies of the Software, in whole or in part, and
 * all derivative works of the Software, unless such copies or derivative
 * works are solely in the form of machine-executable object code generated by
 * a source language processor.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
 * SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
 * FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */
module speedy.stdio;
@safe:

private const BUFSIZE = 32 * 1024;

public import std.stdio;
import std.range, std.format, std.algorithm, std.traits, core.bitop;
import std.exception;

private enum SpeedySafety
{
    Safe,
    Unsafe
}

package template SpeedyWriter(SpeedySafety safety)
{
    private static bool isChar(T)()
    {
        return isSomeChar!T && T.sizeof == 1;
    }

    private static bool isString(T)()
    {
        return is(T == string) || is(T == immutable(string));
    }

    private static bool isNogcCompatibleNested(T)()
    {
        static if (isIntegral!T)
            return true;
        else static if (isInputRange!T || isArray!T)
            return isNogcCompatibleNested!(ElementType!T);
        else
            return false;
    }

    private static bool isNogcCompatible(T)()
    {
        static if (isIntegral!T || isString!T || isChar!T)
            return true;
        else static if (isInputRange!T || isArray!T)
            return isNogcCompatibleNested!(ElementType!T);
        else
            return false;
    }

    static if (safety == SpeedySafety.Safe)
    {
        static char[BUFSIZE] outbuf;
        static ptrdiff_t outbuf_size = 0;
    }
    else
    {
        shared static char[BUFSIZE] outbuf;
        shared static ptrdiff_t outbuf_size = 0;
    }

    version (unittest)
    {
        static void rawWriteStdout()(char[] buffer)
        {
        }

        static bool buffer_contains()(string expected_data) @trusted
        {
            return expected_data == cast(string) outbuf[0 .. outbuf_size];
        }
    }
    else static if (safety == SpeedySafety.Safe)
    {
        version (Posix)
        {
            import core.sys.posix.stdio;
        }
        version (Windows)
        {
            alias flockfile = _lock_file;
            alias funlockfile = _unlock_file;
        }
        shared static FILE* fp;
        shared static this() @trusted
        {
            fp = std.stdio.stdout.getFP;
        }

        static bool flocked;
        static void rawWriteStdout()(char[] buffer) @trusted
        {
            if (!flocked)
            {
                flocked = true;
                flockfile(fp);
            }
            fwrite(cast(void*) buffer.ptr, 1, buffer.length, fp);
        }
    }
    else version (Posix)
    {

        static void rawWriteStdout()(char[] buffer) @trusted
        {
            import core.sys.posix.unistd : write, STDOUT_FILENO;
            import core.sys.posix.fcntl, core.stdc.errno;

            while (buffer.length > 0)
            {
                auto result = write(STDOUT_FILENO, cast(void*) buffer.ptr, buffer.length);
                if (result > 0)
                {
                    buffer = buffer[result .. $];
                    continue;
                }
                assert(result == -1 && (errno == EAGAIN || errno == EINTR), "unexpected error type");
            }
        }
    }
    else version (Windows)
    {
        static void rawWriteStdout()(char[] buffer) @trusted
        {
            import core.sys.windows.windows;

            auto h = GetStdHandle(STD_OUTPUT_HANDLE);
            while (buffer.length > 0)
            {
                uint actual_written;
                auto result = WriteFile(h, cast(void*) buffer.ptr, cast(uint) buffer.length, &actual_written, null);
                if (result && actual_written > 0)
                {
                    buffer = buffer[actual_written .. $];
                    continue;
                }
                assert(result, "unexpected error");
            }
        }
    }

    static void flush()() @trusted
    {
        if (outbuf_size > 0)
            rawWriteStdout(cast(char[]) outbuf[0 .. outbuf_size]);
        outbuf_size = 0;
    }

    static void done()() @trusted
    {
        static if (safety == SpeedySafety.Safe)
        {
            if (flocked)
            {
                flush();
                funlockfile(fp);
                flocked = false;
            }
            else
            {
                flocked = true;
                flush();
                flocked = false;
            }
        }
    }

    shared static ~this()
    {
        flush();
    }

    static void put(T)(T ch) @trusted if (isChar!T)
    {
        if (outbuf_size < outbuf.sizeof)
        {
            outbuf[outbuf_size] = ch;
            outbuf_size = outbuf_size + 1;
        }
        else
        {
            rawWriteStdout(cast(char[]) outbuf[0 .. outbuf_size]);
            outbuf[0] = ch;
            outbuf_size = 1;
        }
    }

    static void ensure_outbuf_capacity()(ptrdiff_t size) @trusted
    {
        if (outbuf_size + size > outbuf.sizeof)
        {
            rawWriteStdout(cast(char[]) outbuf[0 .. outbuf_size]);
            outbuf_size = 0;
            assert(size <= outbuf.sizeof, "write buffer is too small");
        }
    }

    // Compute the number of digits of an integer quickly. Public domain code from Daniel Lemire:
    //  * https://github.com/lemire/Code-used-on-Daniel-Lemire-s-blog
    //  * https://lemire.me/blog/2021/06/03/computing-the-number-of-digits-of-an-integer-even-faster/
    static int fast_digit_count()(uint x)
    {
        static immutable table = [
            4294967296, 8589934582, 8589934582, 8589934582, 12884901788,
            12884901788, 12884901788, 17179868184, 17179868184, 17179868184,
            21474826480, 21474826480, 21474826480, 21474826480, 25769703776,
            25769703776, 25769703776, 30063771072, 30063771072, 30063771072,
            34349738368, 34349738368, 34349738368, 34349738368, 38554705664,
            38554705664, 38554705664, 41949672960, 41949672960, 41949672960,
            42949672960, 42949672960
        ];
        assert(x != 0);
        return (x + table[bsr(x)]) >> 32;
    }

    static int fast_digit_count()(ulong x)
    {
        static immutable table = [
            4503599627370495, 9007199254740982, 9007199254740982,
            9007199254740982, 13510798882111438, 13510798882111438,
            13510798882111438, 18014398509481484, 18014398509481734,
            18014398509481734, 22517998136849980, 22517998136849980,
            22517998136851230, 22517998136851230, 27021597764210476,
            27021597764210476, 27021597764216726, 31525197391530972,
            31525197391530972, 31525197391530972, 36028797018651468,
            36028797018651468, 36028797018651468, 36028797018651468,
            40532396644771964, 40532396644771964, 40532396644771964,
            45035996258079960, 45035996265892460, 45035996265892460,
            49539595822950456, 49539595822950456, 49539595862012956,
            49539595862012956, 54043195137820952, 54043195137820952,
            54043195333133452, 58546793202691448, 58546793202691448,
            58546793202691448, 63050385017561944, 63050385017561944,
            63050385017561944, 63050385017561944, 67553945582432440,
            67553945582432440, 67553945582432440, 72057105756677936,
            72057349897302936, 72057349897302936, 76558752259048432,
            76558752259048432, 76559972962173432, 76559972962173432,
            81052586261418928, 81052586261418928, 81058689777043928,
            85507357763789424, 85507357763789424, 85507357763789424,
            89766816766159920, 89766816766159920, 89766816766159920,
            89766816766159920
        ];
        assert(x != 0);
        int floor_log2_x = bsr(x);
        return (table[floor_log2_x] + (x >> (floor_log2_x / 4))) >> 52;
    }

    // Construct digits conversion table at compile time

    static immutable digitsbatch = 4;

    static immutable tbl = {
        char[(10 ^^ digitsbatch) * digitsbatch] result;
        foreach (i; 0 .. (10 ^^ digitsbatch))
        {
            int v = i;
            foreach_reverse (j; 0 .. digitsbatch)
            {
                result[i * digitsbatch + j] = '0' + v % 10;
                v /= 10;
            }
        }
        return result;
    }();

    static void put(T)(T val) if (isIntegral!T)
    {
        if (val == 0)
        {
            put('0');
            return;
        }
        if (val < 0)
        {
            put('-');
            val = -val;
        }
        Unsigned!T v = val;

        ptrdiff_t ndigits = fast_digit_count(v);
        ensure_outbuf_capacity(ndigits);

        ptrdiff_t i = outbuf_size + ndigits - digitsbatch;
        while (i >= outbuf_size)
        {
            auto newv = v / (10 ^^ digitsbatch);
            auto rem = v - newv * (10 ^^ digitsbatch);
            v = newv;
            ptrdiff_t j = cast(ptrdiff_t) rem * digitsbatch;
            outbuf[i .. i + digitsbatch] = tbl[j .. j + digitsbatch];
            i -= digitsbatch;
        }
        if (i < outbuf_size)
        {
            size_t lastbatch = digitsbatch - (outbuf_size - i);
            ptrdiff_t j = cast(ptrdiff_t) v * digitsbatch + digitsbatch - lastbatch;
            static if (digitsbatch <= 4)
            {
                i = outbuf_size;
                if (lastbatch & 2)
                {
                    outbuf[i .. i + 2] = tbl[j .. j + 2];
                    i += 2;
                    j += 2;
                }
                if (lastbatch & 1)
                    outbuf[i] = tbl[j];
            }
            else
            {
                outbuf[outbuf_size .. outbuf_size + lastbatch] = tbl[j .. j + lastbatch];
            }
        }
        outbuf_size = outbuf_size + ndigits;
    }

    static void put()(string str) @trusted
    {
        while (str.length > 0)
        {
            if (outbuf_size == outbuf.sizeof)
            {
                rawWriteStdout(cast(char[]) outbuf[0 .. outbuf_size]);
                outbuf_size = 0;
            }
            ptrdiff_t chunksize = min(str.length, outbuf.sizeof - outbuf_size);
            outbuf[outbuf_size .. outbuf_size + chunksize] = str[0 .. chunksize];
            str = str[chunksize .. $];
            outbuf_size = outbuf_size + chunksize;
        }
    }

    static void put(Range)(Range r) if ((isInputRange!Range || isArray!Range) && isNogcCompatible!(ElementType!Range))
    {
        static if (isChar!(ElementType!Range))
        {
            put_with_joiner(r, "");
        }
        else
        {
            put('[');
            put_with_joiner(r, ", ");
            put(']');
        }
    }

    static void put_with_joiner(Range, T)(Range r, T sep)
            if ((isInputRange!Range || isArray!Range) && (isString!T || isChar!T))
    {
        bool first = true;
        foreach (x; r)
        {
            if (!first)
                put(sep);
            static if (isNogcCompatible!(typeof(x)))
                put(x);
            else
                put_gc(x);
            first = false;
        }
    }

    // fallback to std.conv for the types, which are not supported
    static void put_gc(T)(T val)
    {
        import std.conv;

        put(val.to!string);
    }

    private struct SpeedyWriterSink
    {
        void put(T)(T arg)
        {
            SpeedyWriter!safety.put(arg);
        }
    }

    static SpeedyWriterSink s;

    static void writef_speedy(alias fmt, A...)(A args)
    {
        static immutable xfmt = parsefmt(fmt);
        put(xfmt.prefix);
        foreach (i, arg; args)
        {
            static if (xfmt.tokens[i].compound)
                put_with_joiner(arg, xfmt.tokens[i].separator);
            else
                put(arg);
            put(xfmt.tokens[i].suffix);
        }
    }

    public static void writef(alias fmt, A...)(A args)
    {
        static if (isSpeedyCompatibleFormat!(fmt, A))
            writef_speedy!fmt(args);
        else
            formattedWrite!fmt(s, args);
        done();
    }

    public static void writef(Char, A...)(in Char[] fmt, A args)
    {
        formattedWrite(s, fmt, args);
        done();
    }

    public static void writefln(alias fmt, A...)(A args)
    {
        static if (isSpeedyCompatibleFormat!(fmt, A))
            writef_speedy!fmt(args);
        else
            formattedWrite!fmt(s, args);
        put('\n');
        done();
    }

    public static void writefln(Char, A...)(in Char[] fmt, A args)
    {
        formattedWrite(s, fmt, args);
        put('\n');
        done();
    }

    public static void write(T...)(T args)
    {
        foreach (arg; args)
            static if (isNogcCompatible!(typeof(arg)))
                put(arg);
            else
                put_gc(arg);
        done();
    }

    public static void writeln(T...)(T args)
    {
        foreach (arg; args)
            static if (isNogcCompatible!(typeof(arg)))
                put(arg);
            else
                put_gc(arg);
        put('\n');
        done();
    }

    // Similar to checkFormatException from Phobos, but only accepts a small
    // subset of possible formats.
    enum isSpeedyCompatibleFormat(alias fmt, Args...) =
    {
        try
        {
            void tassert(T)(T cond)
            {
                if (!cond)
                    throw new Exception("");
            }

            static immutable xfmt = parsefmt(fmt);
            tassert(xfmt.tokens.length == Args.length);
            foreach (i, arg; Args.init)
            {
                if (xfmt.tokens[i].compound == 1 && (isInputRange!(typeof(arg)) || isArray!(typeof(arg))))
                {
                    if (isIntegral!(ElementType!(typeof(arg))))
                        tassert(!find("ds", xfmt.tokens[i].spec).empty);
                    else if (isNogcCompatibleNested!(ElementType!(typeof(arg))))
                        tassert(!find("s", xfmt.tokens[i].spec).empty);
                    else
                        tassert(false);
                }
                else if (xfmt.tokens[i].compound == 2 && (isInputRange!(typeof(arg)) || isArray!(typeof(arg))))
                {
                    if (isIntegral!(ElementType!(typeof(arg))))
                        tassert(!find("ds", xfmt.tokens[i].spec).empty);
                    else if (isNogcCompatible!(ElementType!(typeof(arg))))
                        tassert(!find("s", xfmt.tokens[i].spec).empty);
                    else
                        tassert(false);
                }
                else if (!xfmt.tokens[i].compound)
                {
                    if (isIntegral!(typeof(arg)))
                        tassert(!find("ds", xfmt.tokens[i].spec).empty);
                    else if (isNogcCompatible!(typeof(arg)))
                        tassert(!find("s", xfmt.tokens[i].spec).empty);
                    else
                        tassert(false);
                }
                else
                {
                    tassert(false);
                }

            }
        }
        catch (Exception e)
            return false;
        return true;
    }();

}

///////////////////////////////////////////////////////////////////////////////

private struct FmtToken
{
    char spec;
    byte compound; /* 0 - not a compound, 1 - escaped, 2 - non-escaped */
    string separator;
    string suffix;
}

private struct ParsedFmt
{
    string prefix;
    FmtToken[] tokens;
}

/*
 * This is a minimalistic parser, which recognizes the following patterns:
 *  /%%/                 - escaped '%' character
 *  /%d/                 - decimal number
 *  /%s/                 - string
 *  /%\(%[ds][^%]*%\)/   - escaped compound
 *  /%\-\(%[ds][^%]*%\)/ - non-escaped compound
 */
private ParsedFmt parsefmt(string fmt)
{
    string scan_until_spec(ref string fmt)
    {
        string prefix = "";
        while (!fmt.empty)
        {
            auto tmp = fmt.find('%');
            if (tmp.empty)
            {
                prefix ~= fmt;
                fmt = fmt[$ .. $];
                return prefix;
            }
            else
            {
                enforce(tmp.length >= 2);
                if (tmp[1] == '%')
                {
                    prefix ~= fmt[0 .. fmt.length - tmp.length + 1];
                    fmt = fmt[fmt.length - tmp.length + 2 .. $];
                }
                else
                {
                    prefix ~= fmt[0 .. fmt.length - tmp.length];
                    fmt = fmt[fmt.length - tmp.length .. $];
                    return prefix;
                }
            }
        }
        return prefix;
    }

    ParsedFmt result;
    result.prefix = scan_until_spec(fmt);
    while (!fmt.empty)
    {
        enforce(fmt.length >= 2 && fmt[0] == '%');
        if (fmt[1] == '(' || (fmt.length >= 3 && fmt[1] == '-' && fmt[2] == '('))
        {
            byte compound = 1;
            if (fmt[1] == '-')
            {
                compound = 2;
                fmt = fmt[1 .. $];
            }
            enforce(fmt.length >= 6 && fmt[2] == '%'); // at least %(%s%)
            char spec = fmt[3];
            fmt = fmt[4 .. $];
            string separator = scan_until_spec(fmt);
            enforce(fmt.length >= 2 && fmt[1] == ')');
            fmt = fmt[2 .. $];
            string suffix = scan_until_spec(fmt);
            result.tokens ~= FmtToken(spec, compound, separator, suffix);
        }
        else
        {
            byte spec = fmt[1];
            fmt = fmt[2 .. $];
            string suffix = scan_until_spec(fmt);
            result.tokens ~= FmtToken(spec, 0, "", suffix);
        }
    }
    return result;
}

alias write = SpeedyWriter!(SpeedySafety.Safe).write;
alias writeln = SpeedyWriter!(SpeedySafety.Safe).writeln;
alias writef = SpeedyWriter!(SpeedySafety.Safe).writef;
alias writefln = SpeedyWriter!(SpeedySafety.Safe).writefln;

alias unsafe_write = SpeedyWriter!(SpeedySafety.Unsafe).write;
alias unsafe_writeln = SpeedyWriter!(SpeedySafety.Unsafe).writeln;
alias unsafe_writef = SpeedyWriter!(SpeedySafety.Unsafe).writef;
alias unsafe_writefln = SpeedyWriter!(SpeedySafety.Unsafe).writefln;
alias unsafe_stdout_flush = SpeedyWriter!(SpeedySafety.Unsafe).flush;

version (unittest)
{
    alias unsafe_stdout_buffer_contains = SpeedyWriter!(SpeedySafety.Unsafe).buffer_contains;
}