module jaster.stream.range;

private
{
    import std.range                  : ElementType, isInputRange;
    import std.typecons               : RefCounted, refCounted, RefCountedAutoInitialize;
    import std.experimental.allocator : dispose, makeArray;
    import jaster.stream.core, jaster.stream.util;
}

/// See `byChunkAlloc`.
struct ByChunkAllocRange(Alloc)
{
    private static struct Payload
    {
        alias _buffer this;
        private ubyte[]           _buffer;
        private AllocHelper!Alloc _alloc;
        ~this()
        {
            if(this._buffer !is null)
                this._alloc.alloc.dispose(this._buffer);
        }
    }

    alias RC = RefCounted!(Payload, RefCountedAutoInitialize.no);
    private RC      _buffer;
    private Stream  _stream;
    private ubyte[] _used;
    private bool    _empty;

    static assert(isInputRange!(typeof(this)));
    static assert(is(ElementType!(typeof(this)) == ubyte[]));

    this(ubyte[] buffer, Stream stream, AllocHelper!Alloc alloc)
    {
        this._buffer = refCounted(Payload(buffer, alloc));
        this._stream = stream;
    }

    void popFront()
    {
        this._used = this._stream.readToBuffer(this._buffer);
        if(this._used.length == 0)
            this._empty = true;
    }

    @nogc
    ubyte[] front() nothrow
    {
        return this._used;
    }

    @safe @nogc
    bool empty() nothrow pure const
    {
        return this._empty;
    }
}

/// See `byChunkGC`
struct ByChunkGCRange
{
    private Stream  _stream;
    private ubyte[] _buffer;
    private ubyte[] _used;
    private bool    _empty;

    static assert(isInputRange!(typeof(this)));
    static assert(is(ElementType!(typeof(this)) == ubyte[]));

    this(ubyte[] buffer, Stream stream)
    {
        this._buffer = buffer;
        this._stream = stream;
    }

    void popFront()
    {
        this._used = this._stream.readToBuffer(this._buffer);
        if(this._used.length == 0)
            this._empty = true;
    }

    @nogc
    ubyte[] front() nothrow
    {
        return this._used;
    }

    @safe @nogc
    bool empty() nothrow pure const
    {
        return this._empty;
    }
}

/++
 + Allows access to a stream's data as an input range, chunk-by-chunk.
 +
 + The resulting range is only empty once the stream has no more data to read.
 +
 + Notes:
 +  The stream's $(B current) position will be used.
 +
 +  No part of this function or the resulting range uses GC memory. GC memory is only
 +  used if `Alloc` uses the GC, and/or if the given `stream` uses the GC.
 +
 +  The resulting range uses `std.typecons.RefCounted` for it's internal buffer, so
 +  copying the range has a small overhead.
 +
 + Params:
 +  stream      = The `Stream` to use.
 +  bufferSize  = The size of the buffer, this also determines the max size of each chunk.
 +  alloc       = [Non-static allocators only] The allocator instance to use.
 + ++/
ByChunkAllocRange!Alloc byChunkAlloc(Alloc)(Stream stream, size_t bufferSize)
if(AllocHelper!Alloc.IsStatic)
{
    auto r = ByChunkAllocRange!Alloc(Alloc.instance.makeArray!ubyte(bufferSize), stream, AllocHelper!Alloc());
    r.popFront();
    return r;
}
///
unittest
{
    import std.experimental.allocator.mallocator : Mallocator;
    import std.algorithm                         : all;
    import jaster.stream.memory                  : MemoryStreamGC;

    auto stream = new MemoryStreamGC();
    foreach(i; 0..129)
        stream.write([cast(ubyte)i]);

    stream.position = 0;
    assert(stream.byChunkAlloc!Mallocator(4).all!(arr => arr.length > 0));
}

/// ditto.
ByChunkAllocRange!Alloc byChunkAlloc(Alloc)(Stream stream, size_t bufferSize, Alloc alloc)
{
    auto r = ByChunkAllocRange!Alloc(alloc.makeArray!ubyte(bufferSize), stream, AllocHelper!Alloc(alloc));
    r.popFront();
    return r;
}
///
unittest
{
    import std.experimental.allocator.building_blocks : StatsCollector, Options;
    import std.experimental.allocator.mallocator      : Mallocator;
    import std.algorithm                              : all;
    import jaster.stream.memory                       : MemoryStreamGC;

    alias Alloc = StatsCollector!(Mallocator, Options.all, Options.all);
    auto alloc  = new Alloc();

    auto stream = new MemoryStreamGC();
    foreach(i; 0..129)
        stream.write([cast(ubyte)i]);

    stream.position = 0;
    assert(stream.byChunkAlloc!(Alloc*)(4, alloc).all!(arr => arr.length > 0));
    assert(alloc.bytesUsed == 0);

    // TODO: Test with more complex range chains, including holding the resulting range in a variable and doing stuff with it manually.
    //       If that all passes with 0 bytes used in the end, and 0 crashes, then it should all be good :)
}

/// Variation of `byChunkAlloc` that uses a GC allocated array, and doesn't use ref counting.
ByChunkGCRange byChunkGC(Stream stream, size_t bufferSize)
{
    auto r = ByChunkGCRange(new ubyte[bufferSize], stream);
    r.popFront();
    return r;
}
///
unittest
{
    import jaster.stream.memory : MemoryStreamGC;

    auto stream = new MemoryStreamGC();
    stream.write([0, 1, 2, 3, 4, 5]);
    stream.position = 0;

    auto range = stream.byChunkGC(4);
    assert(!range.empty);
    assert(range.front == [0, 1, 2, 3]);
    
    range.popFront();
    assert(!range.empty);
    assert(range.front == [4, 5]);

    range.popFront();
    assert(range.empty);
}