// Copyright 2019 Microsoft Corporation. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#ifndef RANDOM_ACCESS_STREAM_H
#define RANDOM_ACCESS_STREAM_H

#include <wrl.h>
#include <wrl/client.h>

#include <windows.storage.streams.h>
#include <robuffer.h>

#include <istream>

namespace WinMLTest {

struct BufferBackedRandomAccessStreamReadAsync
    : public Microsoft::WRL::RuntimeClass<
    Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::WinRtClassicComMix | Microsoft::WRL::InhibitRoOriginateError>,
    __FIAsyncOperationWithProgress_2_Windows__CStorage__CStreams__CIBuffer_UINT32,
    ABI::Windows::Foundation::IAsyncInfo> {

    InspectableClass(L"WinMLTest.BufferBackedRandomAccessStreamReadAsync", BaseTrust)

    Microsoft::WRL::ComPtr<ABI::Windows::Storage::Streams::IBuffer> buffer_;

    Microsoft::WRL::ComPtr<ABI::Windows::Foundation::IAsyncOperationWithProgressCompletedHandler<ABI::Windows::Storage::Streams::IBuffer*, UINT32>> completed_handler_;
    Microsoft::WRL::ComPtr<ABI::Windows::Foundation::IAsyncOperationProgressHandler<ABI::Windows::Storage::Streams::IBuffer*, UINT32>> progress_handler_;

    AsyncStatus status_ = AsyncStatus::Started;

public:

    virtual /* [propget] */ HRESULT STDMETHODCALLTYPE get_Id(
        /* [retval][out] */ __RPC__out unsigned __int32* id) override {
        *id = 0; // Do we need to implement this?
        return S_OK;
    }

    virtual /* [propget] */ HRESULT STDMETHODCALLTYPE get_Status(
        /* [retval][out] */ __RPC__out AsyncStatus* status) override {
        *status = status_;
        return S_OK;
    }

    virtual /* [propget] */ HRESULT STDMETHODCALLTYPE get_ErrorCode(
        /* [retval][out] */ __RPC__out HRESULT* /*errorCode*/) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE Cancel(void) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE Close(void) override {
        return E_NOTIMPL;
    }



    HRESULT SetBuffer(ABI::Windows::Storage::Streams::IBuffer* buffer) {
        buffer_ = buffer;
        status_ = AsyncStatus::Completed;
        if (buffer_ != nullptr) {
            if (completed_handler_ != nullptr) {
                completed_handler_->Invoke(this, ABI::Windows::Foundation::AsyncStatus::Completed);
            }
        }
        return S_OK;
    }
    
    virtual HRESULT STDMETHODCALLTYPE put_Progress(
        ABI::Windows::Foundation::IAsyncOperationProgressHandler<ABI::Windows::Storage::Streams::IBuffer*, UINT32>* handler) override {
        progress_handler_ = handler;
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE get_Progress(ABI::Windows::Foundation::IAsyncOperationProgressHandler<ABI::Windows::Storage::Streams::IBuffer*, UINT32>** handler)override {
        progress_handler_.CopyTo(handler);
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE put_Completed(ABI::Windows::Foundation::IAsyncOperationWithProgressCompletedHandler<ABI::Windows::Storage::Streams::IBuffer*, UINT32>* handler) override {
        completed_handler_ = handler;
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE get_Completed(ABI::Windows::Foundation::IAsyncOperationWithProgressCompletedHandler<ABI::Windows::Storage::Streams::IBuffer*, UINT32>** handler) override {
        completed_handler_.CopyTo(handler); 
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE GetResults(ABI::Windows::Storage::Streams::IBuffer** results) override {
        if (buffer_ == nullptr) {
            return E_FAIL;
        }

        buffer_.CopyTo(results);
        return S_OK;
    }
};

struct RandomAccessStream
    : public Microsoft::WRL::RuntimeClass<
          Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::WinRtClassicComMix | Microsoft::WRL::InhibitRoOriginateError>,
          ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType,
          ABI::Windows::Storage::Streams::IContentTypeProvider,
          ABI::Windows::Storage::Streams::IRandomAccessStream,
          ABI::Windows::Storage::Streams::IInputStream,
          ABI::Windows::Storage::Streams::IOutputStream,
          ABI::Windows::Foundation::IClosable> {
  InspectableClass(L"WinMLTest.RandomAccessStream", BaseTrust)

private:
    Microsoft::WRL::ComPtr<ABI::Windows::Storage::Streams::IBuffer> buffer_ = nullptr;
    UINT64 position_ = 0;

public:
    HRESULT RuntimeClassInitialize(ABI::Windows::Storage::Streams::IBuffer* buffer) {
        buffer_ = buffer;
        position_ = 0;
        return S_OK;
    }

    HRESULT RuntimeClassInitialize(ABI::Windows::Storage::Streams::IBuffer* buffer, UINT64 position) {
        buffer_ = buffer;
        position_ = position;
        return S_OK;
    }

    // Content Provider

    /* [propget] */virtual HRESULT STDMETHODCALLTYPE get_ContentType(
        /* [retval, out] */__RPC__deref_out_opt HSTRING* value
    ) override {
        return WindowsCreateString(nullptr, 0, value);
    }

    // IRandomAccessStream

    /* [propget] */virtual HRESULT STDMETHODCALLTYPE get_Size(
        /* [retval, out] */__RPC__out UINT64* value
    ) override {
        *value = 0;
        uint32_t length;
        buffer_->get_Length(&length);
        *value = static_cast<uint64_t>(length);
        return S_OK;
    }

    /* [propput] */virtual HRESULT STDMETHODCALLTYPE put_Size(
        /* [in] */UINT64 /*value*/
    ) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE GetInputStreamAt(
        /* [in] */UINT64 position,
        /* [retval, out] */__RPC__deref_out_opt ABI::Windows::Storage::Streams::IInputStream** stream
    ) override {
        return Microsoft::WRL::MakeAndInitialize<RandomAccessStream>(stream, buffer_.Get(), position);
    }

    virtual HRESULT STDMETHODCALLTYPE GetOutputStreamAt(
        /* [in] */UINT64 /*position*/,
        /* [retval, out] */__RPC__deref_out_opt ABI::Windows::Storage::Streams::IOutputStream** /*stream*/
    ) override {
        return E_NOTIMPL;
    }

    /* [propget] */virtual HRESULT STDMETHODCALLTYPE get_Position(
        /* [retval, out] */__RPC__out UINT64* value
    ) override {
        *value = position_;
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE Seek(
        /* [in] */UINT64 position
    ) override {
        position_ = position;
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE CloneStream(
        /* [retval, out] */__RPC__deref_out_opt ABI::Windows::Storage::Streams::IRandomAccessStream** stream
    ) override {
        return Microsoft::WRL::MakeAndInitialize<RandomAccessStream>(stream, buffer_.Get(), 0);
    }

    /* [propget] */virtual HRESULT STDMETHODCALLTYPE get_CanRead(
        /* [retval, out] */__RPC__out::boolean* value
    ) override {
        UINT32 length;
        buffer_->get_Length(&length);
        *value = buffer_ != nullptr && position_ < static_cast<UINT64>(length);
        return S_OK;
    }

    /* [propget] */virtual HRESULT STDMETHODCALLTYPE get_CanWrite(
        /* [retval, out] */__RPC__out::boolean* value
    ) override {
        *value = false;
        return S_OK;
    }

    // IInputStream
    virtual HRESULT STDMETHODCALLTYPE ReadAsync(
        /* [in] */__RPC__in_opt ABI::Windows::Storage::Streams::IBuffer* buffer,
        /* [in] */UINT32 count,
        /* [in] */ABI::Windows::Storage::Streams::InputStreamOptions /*options*/,
        /* [retval, out] */__RPC__deref_out_opt __FIAsyncOperationWithProgress_2_Windows__CStorage__CStreams__CIBuffer_UINT32** operation
    ) override {
        auto read_async = Microsoft::WRL::Make<BufferBackedRandomAccessStreamReadAsync>();
        read_async.CopyTo(operation);

        // perform the "async work" which is actually synchronous atm
        Microsoft::WRL::ComPtr<ABI::Windows::Storage::Streams::IBuffer> spBuffer = buffer;
        Microsoft::WRL::ComPtr<Windows::Storage::Streams::IBufferByteAccess> out_buffer_byte_access;
        spBuffer.As<Windows::Storage::Streams::IBufferByteAccess>(&out_buffer_byte_access);
        byte* out_bytes = nullptr;
        out_buffer_byte_access->Buffer(&out_bytes);

        Microsoft::WRL::ComPtr<Windows::Storage::Streams::IBufferByteAccess> in_buffer_byte_access;
        buffer_.As<Windows::Storage::Streams::IBufferByteAccess>(&in_buffer_byte_access);
        byte* in_bytes = nullptr;
        in_buffer_byte_access->Buffer(&in_bytes);

        memcpy(out_bytes, in_bytes + static_cast<uint32_t>(position_), count);

        read_async->SetBuffer(buffer);

        return S_OK;
    }

    // IOutputStream
    virtual HRESULT STDMETHODCALLTYPE WriteAsync(
        /* [in] */__RPC__in_opt ABI::Windows::Storage::Streams::IBuffer* /*buffer*/,
        /* [retval, out] */__RPC__deref_out_opt __FIAsyncOperationWithProgress_2_UINT32_UINT32** /*operation*/
    ) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE FlushAsync(
        /* [retval, out] */__RPC__deref_out_opt __FIAsyncOperation_1_boolean** /*operation*/
    ) override {
        return E_NOTIMPL;
    }

    // IClosable
    virtual HRESULT STDMETHODCALLTYPE Close(void) override {
        buffer_ = nullptr;
        return S_OK;
    }

};

struct BufferBackedRandomAccessStreamReferenceOpenReadAsync
    : public Microsoft::WRL::RuntimeClass<
    Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::WinRtClassicComMix | Microsoft::WRL::InhibitRoOriginateError>,
    __FIAsyncOperation_1_Windows__CStorage__CStreams__CIRandomAccessStreamWithContentType,
    ABI::Windows::Foundation::IAsyncInfo> {

    InspectableClass(L"WinMLTest.BufferBackedRandomAccessStreamReferenceOpenReadAsync", BaseTrust)
public:
    Microsoft::WRL::ComPtr<ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType> ras_;
    Microsoft::WRL::ComPtr<ABI::Windows::Foundation::IAsyncOperationCompletedHandler<ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType*>> completed_handler_;
    AsyncStatus status_ = AsyncStatus::Started;

    HRESULT SetRandomAccessStream(ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType* ras) {
        ras_ = ras;
        status_ = AsyncStatus::Completed;
        if (ras_ != nullptr) {
            if (completed_handler_ != nullptr) {
                completed_handler_->Invoke(this, status_);
            }
        }
        return S_OK;
    }

    virtual /* [propget] */ HRESULT STDMETHODCALLTYPE get_Id(
        /* [retval][out] */ __RPC__out unsigned __int32* id) override {
        *id = 0; // Do we need to implement this?
        return S_OK;
    }

    virtual /* [propget] */ HRESULT STDMETHODCALLTYPE get_Status(
        /* [retval][out] */ __RPC__out AsyncStatus* status) override {
        *status = status_;
        return S_OK;
    }

    virtual /* [propget] */ HRESULT STDMETHODCALLTYPE get_ErrorCode(
        /* [retval][out] */ __RPC__out HRESULT* /*errorCode*/) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE Cancel(void) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE Close(void) override {
        return E_NOTIMPL;
    }

    virtual HRESULT STDMETHODCALLTYPE put_Completed(
        ABI::Windows::Foundation::IAsyncOperationCompletedHandler<ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType*>* handler) override
    {
        completed_handler_ = handler;
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE get_Completed(
        ABI::Windows::Foundation::IAsyncOperationCompletedHandler<ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType*>** handler) override {
        completed_handler_.CopyTo(handler);
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE GetResults(
        ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType** results) override {
        if (ras_ == nullptr) {
            return E_FAIL;
        }
        ras_.CopyTo(results);
        return S_OK;
    }
};

struct BufferBackedRandomAccessStreamReference
    : public Microsoft::WRL::RuntimeClass<
    Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::WinRtClassicComMix | Microsoft::WRL::InhibitRoOriginateError>,
    ABI::Windows::Storage::Streams::IRandomAccessStreamReference> {
    InspectableClass(L"WinMLTest.BufferBackedRandomAccessStreamReference", BaseTrust)

    Microsoft::WRL::ComPtr<ABI::Windows::Storage::Streams::IBuffer> buffer_ = nullptr;

public:
    HRESULT RuntimeClassInitialize(ABI::Windows::Storage::Streams::IBuffer* buffer) {
        buffer_ = buffer;
        return S_OK;
    }

    virtual HRESULT STDMETHODCALLTYPE OpenReadAsync(
        /* [retval, out] */__RPC__deref_out_opt __FIAsyncOperation_1_Windows__CStorage__CStreams__CIRandomAccessStreamWithContentType** operation
    ) override {
        auto open_read_async = Microsoft::WRL::Make<BufferBackedRandomAccessStreamReferenceOpenReadAsync>();
        open_read_async.CopyTo(operation);

        Microsoft::WRL::ComPtr<RandomAccessStream> ras;
        Microsoft::WRL::MakeAndInitialize<RandomAccessStream>(&ras, buffer_.Get());

        Microsoft::WRL::ComPtr<ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType> ras_interface = nullptr;
        ras.As<ABI::Windows::Storage::Streams::IRandomAccessStreamWithContentType>(&ras_interface);

        open_read_async.Get()->SetRandomAccessStream(ras_interface.Get());
        return S_OK;
    }

};

} // namespace WinMLTest

#endif // RANDOM_ACCESS_STREAM_H