module jaster.stream.memory;

private
{
    import std.experimental.allocator, std.experimental.allocator.gc_allocator, std.experimental.allocator.mallocator;
    import jaster.stream.core, jaster.stream.util;
}

/// A `MemoryStream` that uses the GC.
alias MemoryStreamGC = MemoryStream!GCAllocator;

/// A `MemoryStream` that uses malloc
alias MemoryStreamMalloc = MemoryStream!Mallocator;

/++
 + A `Stream` that writes into a memory buffer.
 +
 + This stream supports `std.experimental.allocator` for all of it's allocations and freeing.
 +
 + This stream will automatically call `dispose` during it's deconstructor.
 +
 + Outside of exceptions, this stream doesn't allocate GC memory unless the given allocator does.
 + ++/
class MemoryStream(Alloc) : Stream
{
    private
    {
        ubyte[]           _buffer;
        bool              _isDisposed;
        size_t            _position;
        AllocHelper!Alloc _alloc;
    }

    static if(_alloc.IsStatic)
    this()
    {
        super(Stream.Can.Write | Stream.Can.Read | Stream.Can.Seek);
    }

    static if(!_alloc.IsStatic)
    this(Alloc alloc)
    {
        this._alloc.alloc = alloc;
        super(Stream.Can.Write | Stream.Can.Read | Stream.Can.Seek);
    }

    ~this()
    {
        this.dispose();
    }

    public override
    {
        const(ubyte[]) write(scope const ubyte[] data)
        {
            enforceNotDisposed(this);
            enforceCanWrite(this);

            auto end = (this.position + data.length);
            if(end > this.length)
                this.length = end;

            this._buffer[this.position..end] = data[];
            this.position = end;

            return data;
        }

        ubyte[] readToBuffer(scope ref ubyte[] buffer)
        {
            enforceNotDisposed(this);
            enforceCanRead(this);

            auto end = (this.position + buffer.length);
            if(end > this.length)
                end = this.length;

            auto amount = (end - this.position);

            buffer[0..amount] = this._buffer[this.position..end];
            this.position = end;

            return buffer[0..amount];
        }

        void dispose()
        {
            if(!this.isDisposed)
            {
                this._isDisposed = true;
                this._alloc.dispose(this._buffer);
            }
        }

        void flush(){}

        void seek(SeekFrom from, ptrdiff_t amount)
        {
            auto start = (from == SeekFrom.Start)
                         ? 0
                         : (from == SeekFrom.Current)
                            ? this.position
                            : this.length;
            this.position = start + amount;
        }

        @property
        bool isDisposed() const
        {
            return this._isDisposed;
        }

        @property
        void writeTimeout(Duration dur)
        {
            enforceNotDisposed(this);
            throw new StreamCannotTimeoutException();
        }

        @property
        Duration writeTimeout()
        {
            enforceNotDisposed(this);
            throw new StreamCannotTimeoutException();
        }

        @property
        void readTimeout(Duration dur)
        {
            enforceNotDisposed(this);
            throw new StreamCannotTimeoutException();
        }

        @property
        Duration readTimeout()
        {
            enforceNotDisposed(this);
            throw new StreamCannotTimeoutException();
        }

        @property
        void length(size_t size)
        {
            enforceNotDisposed(this);

            if(this._buffer.ptr is null)
            {
                this._buffer = this._alloc.makeArray!ubyte(size);
            }
            else if(size > this.length)
            {
                auto success = this._alloc.expandArray(this._buffer, (size - this.length));
                assert(success);
            }
            else if(size < this.length)
            {
                auto success = this._alloc.shrinkArray(this._buffer, (this.length - size));
                assert(success);
            }
        }

        @property
        size_t length()
        {
            enforceNotDisposed(this);
            return this._buffer.length;
        }

        @property
        void position(size_t pos)
        {
            import std.exception : enforce;

            enforceNotDisposed(this);
            enforce!StreamException(pos <= this.length, "Attempted to seek past the end of the stream.");
            this._position = pos;
        }

        @property
        size_t position()
        {
            enforceNotDisposed(this);
            return this._position;
        }
    }

    public
    {
        /++
         + Gets the internal buffer of the stream.
         +
         + Notes:
         +  While this function by itself is `@safe`, usage of this function is unsafe.
         +
         +  This is because the buffer can be reallocated and deallocated during certain usage of the stream, so
         +  the slice returned by this function can point to freed/invalid memory, causing a crash (hopefully).
         + ++/
        @property @nogc
        inout(ubyte[]) data() nothrow inout
        {
            return this._buffer;
        }
    }
}
///
unittest
{
    import std.experimental.allocator.mallocator      : Mallocator;
    import std.experimental.allocator.building_blocks : StatsCollector, Options;

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

    auto stream = new MStream(alloc);
    assert(stream.position == 0);
    assert(stream.length   == 0);
    assert(stream.canWrite);
    assert(stream.canRead);
    assert(stream.canSeek);
    assert(!stream.canTimeout);
    assert(!stream.isDisposed);

    stream.write([0, 0, 0, 0]); // Testing to see if it cleans memory after use.
    stream.dispose();
    assert(stream.isDisposed);

    stream = new MStream(alloc);
    stream.write([0, 1, 2, 3, 4]);
    assert(stream.position == 5);
    assert(stream.length == 5);
    assert(stream.data == [0, 1, 2, 3, 4]);

    stream.length = 20;
    assert(stream.length == 20);

    stream.length = 5;
    assert(stream.length == 5);

    stream.position = stream.position - 2;
    assert(stream.read(2) == [3, 4]);

    assert(alloc.bytesUsed > 0);
    destroy(stream); // To simulate the GC collecting it
    // import std.stdio;
    // alloc.reportStatistics(stdout);
    
    assert(alloc.bytesUsed == 0);
}

// CopyTo test
unittest
{
    import std.experimental.allocator.mallocator      : Mallocator;
    import std.experimental.allocator.building_blocks : StatsCollector, Options;

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

    auto master = new MStream(alloc);
    auto slave  = new MStream(alloc);

    void doAssert(ubyte[] buffer1 = null, ubyte[] buffer2 = null)
    {
        import std.format : format;

        if(buffer1 is null)
            buffer1 = slave.data;
        
        if(buffer2 is null)
            buffer2 = master.data;

        assert(buffer1 == buffer2, format("%s\n\n\n%s", buffer1, buffer2));
    }

    // Optimisation example:
    //  Without the line below, `alloc` reports 1000 reallocations.
    //  With it, it reports a single one.
    master.length = 1_000;
    foreach(i; 0..master.length)
    {
        auto data = cast(ubyte)i;
        master.write((&data)[0..1]);
    }
    assert(master.position == master.length);

    // Test copyToAlloc, with a small buffer, and a static allocator.
    master.position = 0;
    master.copyToAlloc!Mallocator(slave, 32);
    assert(master.position == master.length);
    doAssert();

    // Test copyToAlloc, with a large buffer, and a non-static allocator.
    master.position = 0;
    slave.length    = 0;
    slave.position  = 0;
    master.copyToAlloc(slave, 8126, alloc);
    assert(master.position == master.length);
    doAssert();

    // Test copyToStack, with a small buffer, with the master stream not at the start.
    master.position = master.length / 2;
    slave.length    = 0;
    slave.position  = 0;
    master.copyToStack!64(slave);
    assert(master.position == master.length);
    assert(slave.length    == master.length / 2);
    doAssert(slave.data, master.data[master.length/2..$]);

    // Test copyToGC, with an average buffer, with the slave stream not at the start.
    master.position = 0;
    doAssert(slave.data, master.data[$/2..$]);
    master.copyToGC(slave, 512);
    assert(master.position == master.length);
    assert(slave.length    == master.length + (master.length / 2));
    doAssert(slave.data[0..master.length/2], master.data[$/2..$]);
    doAssert(slave.data[master.length/2..$], master.data);

    master.dispose();
    slave.dispose();

    // import std.stdio;
    // alloc.reportStatistics(stdout);
    assert(alloc.bytesUsed == 0);
}