#pragma once

#if !defined(MIJIN_IO_STREAM_HPP_INCLUDED)
#define MIJIN_IO_STREAM_HPP_INCLUDED 1

#include <cstdint>
#include <optional>
#include <ranges>
#include <span>
#include <stdexcept>
#include <string>
#include "../async/coroutine.hpp"
#include "../container/typeless_buffer.hpp"
#include "../internal/common.hpp"
#include "../types/result.hpp"
#include "../util/exception.hpp"

namespace mijin
{

//
// public defines
//

//
// public constants
//

//
// public types
//

enum class SeekMode
{
    ABSOLUTE,
    RELATIVE,
    RELATIVE_TO_END
};

struct ReadOptions
{
    bool partial : 1 = false;
    bool peek    : 1 = false;
    bool noBlock : 1 = false;
};

struct StreamFeatures
{
    bool read  : 1 = false;
    bool write : 1 = false;
    bool tell  : 1 = false;
    bool seek  : 1 = false;
    bool async : 1 = false;

    ReadOptions readOptions = {};
};

enum class FileOpenMode
{
    READ,
    WRITE,
    APPEND,
    READ_WRITE
};

enum class [[nodiscard]] StreamError
{
    SUCCESS            = 0,
    IO_ERROR           = 1,
    NOT_SUPPORTED      = 2,
    CONNECTION_CLOSED  = 3,
    PROTOCOL_ERROR     = 4,
    WOULD_BLOCK        = 5,
    CONNECTION_REFUSED = 6,
    UNKNOWN_ERROR      = -1
};

class Stream
{
public:
    virtual ~Stream() = default;
    
public:
    virtual StreamError readRaw(std::span<std::uint8_t> buffer, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr) = 0;
    [[deprecated("Partial parameter has been replaced with options.")]]
    StreamError readRaw(std::span<std::uint8_t> buffer, bool partial, std::size_t* outBytesRead = nullptr)
    {
        return readRaw(buffer, {.partial = partial}, outBytesRead);
    }
    virtual StreamError writeRaw(std::span<const std::uint8_t> buffer) = 0;
    virtual std::size_t tell() = 0;
    virtual StreamError seek(std::intptr_t pos, SeekMode seekMode = SeekMode::ABSOLUTE) = 0;
    virtual void flush();
    virtual bool isAtEnd() = 0;
    virtual StreamFeatures getFeatures() = 0;

    // async interface (requires getFeatures().async to be set)
    virtual mijin::Task<StreamError> c_readRaw(std::span<std::uint8_t> buffer, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr);
    virtual mijin::Task<StreamError> c_writeRaw(std::span<const std::uint8_t> buffer);

    StreamError readRaw(void* outData, std::size_t bytes, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        std::uint8_t* ptr = static_cast<std::uint8_t*>(outData);
        return readRaw(std::span(ptr, ptr + bytes), options, outBytesRead);
    }

    [[deprecated("Partial parameter has been replaced with options.")]]
    StreamError readRaw(void* outData, std::size_t bytes, bool partial, std::size_t* outBytesRead = nullptr)
    {
        std::uint8_t* ptr = static_cast<std::uint8_t*>(outData);
        return readRaw(std::span(ptr, ptr + bytes), {.partial = partial}, outBytesRead);
    }

    mijin::Task<StreamError> c_readRaw(void* outData, std::size_t bytes, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        std::uint8_t* ptr = static_cast<std::uint8_t*>(outData);
        co_return co_await c_readRaw(std::span(ptr, ptr + bytes), options, outBytesRead);
    }

    template<std::ranges::contiguous_range TRange>
    StreamError readRaw(TRange& range, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        const std::size_t bytes = std::distance(range.begin(), range.end()) * sizeof(std::ranges::range_value_t<TRange>);
        return readRaw(&*range.begin(), bytes, options, outBytesRead);
    }

    template<std::ranges::contiguous_range TRange>
    [[deprecated("Partial parameter has been replaced with options.")]]
    StreamError readRaw(TRange& range, bool partial, std::size_t* outBytesRead = nullptr)
    {
        const std::size_t bytes = std::distance(range.begin(), range.end()) * sizeof(std::ranges::range_value_t<TRange>);
        return readRaw(&*range.begin(), bytes, {.partial = partial}, outBytesRead);
    }

    template<template<typename> typename TAllocator>
    StreamError readRaw(BaseTypelessBuffer<TAllocator>& buffer, const ReadOptions& options = {})
    {
        return readRaw(buffer.data(), buffer.byteSize(), options);
    }

    template<std::ranges::contiguous_range TRange>
    mijin::Task<StreamError> c_readRaw(TRange& range, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        const std::size_t bytes = std::distance(range.begin(), range.end()) * sizeof(std::ranges::range_value_t<TRange>);
        return c_readRaw(&*range.begin(), bytes, options, outBytesRead);
    }

    template<template<typename> typename TAllocator>
    mijin::Task<StreamError> c_readRaw(BaseTypelessBuffer<TAllocator>& buffer, const ReadOptions& options = {})
    {
        return c_readRaw(buffer.data(), buffer.byteSize(), options);
    }

    StreamError writeRaw(const void* data, std::size_t bytes)
    {
        const std::uint8_t* ptr = static_cast<const std::uint8_t*>(data);
        return writeRaw(std::span(ptr, ptr + bytes));
    }

    mijin::Task<StreamError> c_writeRaw(const void* data, std::size_t bytes)
    {
        const std::uint8_t* ptr = static_cast<const std::uint8_t*>(data);
        return c_writeRaw(std::span(ptr, ptr + bytes));
    }

    template<std::ranges::contiguous_range TRange>
    StreamError writeRaw(const TRange& range)
    {
        const std::size_t bytes = std::distance(range.begin(), range.end()) * sizeof(std::ranges::range_value_t<TRange>);
        return writeRaw(&*range.begin(), bytes);
    }

    template<std::ranges::contiguous_range TRange>
    mijin::Task<StreamError> c_writeRaw(const TRange& range)
    {
        const std::size_t bytes = std::distance(range.begin(), range.end()) * sizeof(std::ranges::range_value_t<TRange>);
        return c_writeRaw(&*range.begin(), bytes);
    }

    template<typename T>
    StreamError read(T& value, const ReadOptions& options = {})
    {
        MIJIN_ASSERT(!options.partial, "Cannot partially read a value.");
        return readRaw(&value, sizeof(T), options);
    }

    template<typename T>
    mijin::Task<StreamError> c_read(T& value, const ReadOptions& options = {})
    {
        MIJIN_ASSERT(!options.partial, "Cannot partially read a value.");
        return c_readRaw(&value, sizeof(T), options);
    }

    template<typename T>
    StreamError readSpan(T& values, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        auto asSpan = std::span(values);
        return readRaw(asSpan.data(), asSpan.size_bytes(), options, outBytesRead);
    }

    template<typename T>
    mijin::Task<StreamError> c_readSpan(T& values, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        auto asSpan = std::span(values);
        return c_readRaw(asSpan.data(), asSpan.size_bytes(), options, outBytesRead);
    }

    template<typename TItBegin, typename TItEnd>
    StreamError readSpan(TItBegin&& begin, TItEnd&& end, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        auto asSpan = std::span(std::forward<TItBegin>(begin), std::forward<TItEnd>(end));
        return readRaw(asSpan.data(), asSpan.size_bytes(), options, outBytesRead);
    }

    template<typename TItBegin, typename TItEnd>
    mijin::Task<StreamError> c_readSpan(TItBegin&& begin, TItEnd&& end, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr)
    {
        auto asSpan = std::span(std::forward<TItBegin>(begin), std::forward<TItEnd>(end));
        return c_readRaw(asSpan.data(), asSpan.size_bytes(), options, outBytesRead);
    }

    template<typename T>
    StreamError write(const T& value)
    {
        return writeRaw(&value, sizeof(T));
    }

    template<typename T>
    mijin::Task<StreamError> c_write(const T& value) requires(std::is_trivial_v<T>)
    {
        return c_writeRaw(&value, sizeof(T));
    }

    template<typename T>
    StreamError writeSpan(const T& values)
    {
        auto asSpan = std::span(values);
        return writeRaw(asSpan.data(), asSpan.size_bytes());
    }

    template<typename T>
    mijin::Task<StreamError> c_writeSpan(const T& values)
    {
        auto asSpan = std::span(values);
        return c_writeRaw(asSpan.data(), asSpan.size_bytes());
    }

    template<typename TItBegin, typename TItEnd>
    StreamError writeSpan(TItBegin&& begin, TItEnd&& end)
    {
        return writeSpan(std::span(std::forward<TItBegin>(begin), std::forward<TItEnd>(end)));
    }

    template<typename TItBegin, typename TItEnd>
    mijin::Task<StreamError> c_writeSpan(TItBegin&& begin, TItEnd&& end)
    {
        return c_writeSpan(std::span(std::forward<TItBegin>(begin), std::forward<TItEnd>(end)));
    }

    StreamError readBinaryString(std::string& outString);
    StreamError writeBinaryString(std::string_view str);

    StreamError readZString(std::string& outString);
    StreamError writeZString(std::string_view str);

    mijin::Task<StreamError> c_readBinaryString(std::string& outString);
    mijin::Task<StreamError> c_writeBinaryString(std::string_view str);

    [[deprecated("Use readBinaryString() or readAsString() instead.")]]
    inline StreamError readString(std::string& outString) { return readBinaryString(outString); }

    [[deprecated("Use writeBinaryString() or writeText() instead.")]]
    inline StreamError writeString(std::string_view str) { return writeBinaryString(str); }

    StreamError getTotalLength(std::size_t& outLength);

    StreamError copyTo(Stream& otherStream);

    template<template<typename> typename TAllocator>
    StreamError readRest(BaseTypelessBuffer<TAllocator>& outBuffer);

    template<template<typename> typename TAllocator>
    mijin::Task<StreamError> c_readRest(BaseTypelessBuffer<TAllocator>& outBuffer);

    StreamError readLine(std::string& outString);
    mijin::Task<StreamError> c_readLine(std::string& outString);

    template<typename TChar = char, typename TTraits = std::char_traits<TChar>, typename TAllocator = MIJIN_DEFAULT_ALLOCATOR<char>>
    StreamError readAsString(std::basic_string<TChar, TTraits, TAllocator>& outString);

    template<typename TChar = char, typename TTraits = std::char_traits<TChar>, typename TAllocator = MIJIN_DEFAULT_ALLOCATOR<char>>
    mijin::Task<StreamError> c_readAsString(std::basic_string<TChar, TTraits, TAllocator>& outString);

    StreamError writeText(std::string_view str)
    {
        return writeSpan(str);
    }

    mijin::Task<StreamError> c_writeText(std::string_view str)
    {
        return c_writeSpan(str);
    }
};

class FileStream : public Stream
{
private:
    std::FILE* handle = nullptr; // TODO: wrap in gsl::owner<>
    FileOpenMode mode;
    std::size_t length = 0;
public:
    ~FileStream() override;

    StreamError open(const char* path, FileOpenMode mode_);
    inline StreamError open(const std::string& path, FileOpenMode mode_) {
        return open(path.c_str(), mode_);
    }
    void close();
    [[nodiscard]] inline bool isOpen() const { return handle != nullptr; }

    // Stream overrides
    StreamError readRaw(std::span<std::uint8_t> buffer, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr) override;
    StreamError writeRaw(std::span<const std::uint8_t> buffer) override;
    std::size_t tell() override;
    StreamError seek(std::intptr_t pos, SeekMode seekMode = SeekMode::ABSOLUTE) override;
    void flush() override;
    bool isAtEnd() override;
    StreamFeatures getFeatures() override;
};

class MemoryStream : public Stream
{
private:
    std::span<std::uint8_t> data_;
    std::size_t pos_ = 0;
    bool canWrite_ = false;
public:
    void openRW(std::span<std::uint8_t> data);
    template<typename T>
    void openRO(std::span<T> data) {
        openROImpl(data.data(), data.size_bytes());
    }
    template<typename TChar>
    inline void openRO(std::basic_string_view<TChar> stringView) {
        openROImpl(stringView.data(), stringView.size() * sizeof(TChar));
    }
    void close();
    [[nodiscard]] inline bool isOpen() const { return data_.data() != nullptr; }
    [[nodiscard]] inline std::size_t availableBytes() const
    {
        MIJIN_ASSERT(isOpen(), "MemoryStream is not open.");
        return data_.size() - pos_;
    }

    // Stream overrides
    StreamError readRaw(std::span<std::uint8_t> buffer, const ReadOptions& options = {}, std::size_t* outBytesRead = nullptr) override;
    StreamError writeRaw(std::span<const std::uint8_t> buffer) override;
    std::size_t tell() override;
    StreamError seek(std::intptr_t pos, SeekMode seekMode = SeekMode::ABSOLUTE) override;
    bool isAtEnd() override;
    StreamFeatures getFeatures() override;
private:
    void openROImpl(const void* data, std::size_t bytes);
};

template<typename TSuccess>
using StreamResult = ResultBase<TSuccess, StreamError>;

//
// public functions
//

template<template<typename> typename TAllocator>
StreamError Stream::readRest(BaseTypelessBuffer<TAllocator>& outBuffer)
{
    // first try to allocate everything at once
    std::size_t length = 0;
    if (const StreamError lengthError = getTotalLength(length); lengthError == StreamError::SUCCESS)
    {
        MIJIN_ASSERT(getFeatures().tell, "How did you find the length if you cannot tell()?");
        length -= tell();
        outBuffer.resize(length);
        if (const StreamError error = readRaw(outBuffer.data(), length); error != StreamError::SUCCESS)
        {
            return error;
        }
        return StreamError::SUCCESS;
    }

    // could not determine the size, read chunk-wise
    static constexpr std::size_t CHUNK_SIZE = 4096;
    std::array<std::byte, CHUNK_SIZE> chunk = {};

    while (!isAtEnd())
    {
        std::size_t bytesRead = 0;
        if (const StreamError error = readRaw(chunk, {.partial = true}, &bytesRead); error != StreamError::SUCCESS)
        {
            return error;
        }

        outBuffer.resize(outBuffer.byteSize() + bytesRead);
        const std::span<std::byte> bufferBytes = outBuffer.template makeSpan<std::byte>();
        std::copy_n(chunk.begin(), bytesRead, bufferBytes.end() - static_cast<long>(bytesRead));
    }
    return StreamError::SUCCESS;
}

template<template<typename> typename TAllocator>
mijin::Task<StreamError> Stream::c_readRest(BaseTypelessBuffer<TAllocator>& outBuffer)
{
    // first try to allocate everything at once
    std::size_t length = 0;
    if (StreamError lengthError = getTotalLength(length); lengthError == StreamError::SUCCESS)
    {
        MIJIN_ASSERT(getFeatures().tell, "How did you find the length if you cannot tell()?");
        length -= tell();
        outBuffer.resize(length);
        if (StreamError error = co_await c_readRaw(outBuffer.data(), length); error != StreamError::SUCCESS)
        {
            co_return error;
        }
        co_return StreamError::SUCCESS;
    }

    // could not determine the size, read chunk-wise
    static constexpr std::size_t CHUNK_SIZE = 4096;
    std::array<std::byte, CHUNK_SIZE> chunk = {};

    while (!isAtEnd())
    {
        std::size_t bytesRead = 0;
        if (StreamError error = co_await c_readRaw(chunk, {.partial = true}, &bytesRead); error != StreamError::SUCCESS)
        {
            co_return error;
        }

        outBuffer.resize(outBuffer.byteSize() + bytesRead);
        std::span<std::byte> bufferBytes = outBuffer.template makeSpan<std::byte>();
        std::copy_n(chunk.begin(), bytesRead, bufferBytes.end() - static_cast<long>(bytesRead));
    }
    co_return StreamError::SUCCESS;
}

template<typename TChar, typename TTraits, typename TAllocator>
StreamError Stream::readAsString(std::basic_string<TChar, TTraits, TAllocator>& outString)
{
    static_assert(sizeof(TChar) == 1, "Can only read to 8-bit character types (char, unsigned char or char8_t");

    // first try to allocate everything at once
    std::size_t length = 0;
    if (StreamError lengthError = getTotalLength(length); lengthError == StreamError::SUCCESS)
    {
        MIJIN_ASSERT(getFeatures().tell, "How did you find the length if you cannot tell()?");
        length -= tell();
        outString.resize(length);
        if (StreamError error = readRaw(outString.data(), length); error != StreamError::SUCCESS)
        {
            return error;
        }
        return StreamError::SUCCESS;
    }

    // could not determine the size, read chunk-wise
    static constexpr std::size_t CHUNK_SIZE = 4096;
    std::array<TChar, CHUNK_SIZE> chunk;

    outString.clear();
    while (!isAtEnd())
    {
        std::size_t bytesRead = 0;
        if (StreamError error = readRaw(chunk, {.partial = true}, &bytesRead); error != StreamError::SUCCESS)
        {
            return error;
        }
        outString.append(chunk.data(), chunk.data() + bytesRead);
    }
    return StreamError::SUCCESS;
}

template<typename TChar, typename TTraits, typename TAllocator>
mijin::Task<StreamError> Stream::c_readAsString(std::basic_string<TChar, TTraits, TAllocator>& outString)
{
    static_assert(sizeof(TChar) == 1, "Can only read to 8-bit character types (char, unsigned char or char8_t");

    // first try to allocate everything at once
    std::size_t length = 0;
    if (StreamError lengthError = getTotalLength(length); lengthError == StreamError::SUCCESS)
    {
        MIJIN_ASSERT(getFeatures().tell, "How did you find the length if you cannot tell()?");
        length -= tell();
        outString.resize(length);
        if (StreamError error = co_await c_readRaw(outString.data(), length); error != StreamError::SUCCESS)
        {
            co_return error;
        }
        co_return StreamError::SUCCESS;
    }

    // could not determine the size, read chunk-wise
    static constexpr std::size_t CHUNK_SIZE = 4096;
    std::array<TChar, CHUNK_SIZE> chunk;

    outString.clear();
    while (!isAtEnd())
    {
        std::size_t bytesRead = 0;
        if (StreamError error = co_await c_readRaw(chunk, true, &bytesRead); error != StreamError::SUCCESS)
        {
            co_return error;
        }
        outString.append(chunk.data(), chunk.data() + bytesRead);
    }
    co_return StreamError::SUCCESS;
}


inline const char* errorName(StreamError error) MIJIN_NOEXCEPT
{
    switch (error)
    {
        case StreamError::SUCCESS:
            return "success";
        case StreamError::IO_ERROR:
            return "IO error";
        case StreamError::NOT_SUPPORTED:
            return "not supported";
        case StreamError::CONNECTION_CLOSED:
            return "connection closed";
        case StreamError::PROTOCOL_ERROR:
            return "protocol error";
        case StreamError::WOULD_BLOCK:
            return "would block";
        case StreamError::UNKNOWN_ERROR:
            return "unknown error";
        case StreamError::CONNECTION_REFUSED:
            return "connection refused";
    }
    return "<invalid error>";
}

#if MIJIN_ENABLE_EXCEPTIONS // these functions don't make sense with exceptions disabled
inline void throwOnError(mijin::StreamError error)
{
    if (error == mijin::StreamError::SUCCESS) {
        return;
    }
    throw Exception(errorName(error));
}

inline void throwOnError(mijin::StreamError error, std::string message)
{
    if (error == mijin::StreamError::SUCCESS) {
        return;
    }
    throw Exception(message + ": " + errorName(error));
}

template<typename TSuccess>
inline decltype(auto) throwOnError(StreamResult<TSuccess>&& result)
{
    if (result.isError())
    {
        throw Exception(errorName(result.getError()));
    }
    else if (!result.isSuccess())
    {
        throw Exception("result is empty");
    }
    return *result;
}

template<typename TSuccess>
inline decltype(auto) throwOnError(StreamResult<TSuccess>&& result, const std::string& message)
{
    if (result.isError())
    {
        throw Exception(message + ": " + errorName(result.getError()));
    }
    else if (!result.isSuccess())
    {
        throw Exception(message + ": result is empty");
    }
    return *result;
}
#endif // MIJIN_ENABLE_EXCEPTIONS
} // namespace mijin

#endif // !defined(MIJIN_IO_STREAM_HPP_INCLUDED)