#include <Core/HAL/IO/IOStream.h>

#include <Core/Common/String.h>
#include <Core/HAL/IO/IOStream_Private.h>
#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/Memory/Allocator.h>
#include <Core/Thread/Thread.h>
#include <cstdarg>

namespace Juliet
{
    IOStream* IOFromFile(String filename, String mode)
    {
        if (!IsValid(filename))
        {
            Log(LogLevel::Error, LogCategory::Core, "Trying to open IOStream on invalid filename");
            return nullptr;
        }
        if (!IsValid(mode))
        {
            Log(LogLevel::Error, LogCategory::Core, "Trying to open IOStream with invalid mode");
            return nullptr;
        }

        return Internal::IOFromFile(filename, mode);
    }

    IOStream* IOFromInterface(NonNullPtr<const IOStreamInterface> streamInterface, NonNullPtr<IOStreamDataPayload> payload)
    {
        Assert(streamInterface->Version >= sizeof(*streamInterface.Get()));

        auto stream = static_cast<IOStream*>(Calloc(1, sizeof(IOStream)));
        if (stream)
        {
            IOStreamInterface*       dstInterface = &stream->Interface;
            const IOStreamInterface* srcInterface = streamInterface.Get();
            static_assert(sizeof(*(dstInterface)) == sizeof(*(srcInterface)), "Source and Destination type mismatch");
            MemCopy(dstInterface, srcInterface, sizeof(*srcInterface));
            stream->Data = payload.Get();
        }
        return stream;
    }

    size_t IOPrintf(NonNullPtr<IOStream> stream, _Printf_format_string_ const char* format, ...)
    {
        // TODO: Juliet format function should be able to allocate on scratch arena here.
        // This buffer should be big enough until then
        char    formattedBuffer[4096];
        va_list args;

        va_start(args, format);
        int writtenSize = vsprintf_s(formattedBuffer, format, args); // Cast to void to ignore the return type. TODO : Juliet format function
        va_end(args);

        Assert(writtenSize >= 0);

        ByteBuffer buffer = { .Data = reinterpret_cast<Byte*>(formattedBuffer), .Size = static_cast<size_t>(writtenSize) };
        return IOWrite(stream, buffer);
    }

    size_t IOWrite(NonNullPtr<IOStream> stream, ByteBuffer inBuffer)
    {
        if (!stream->Interface.Write)
        {
            stream->Status = IOStreamStatus::ReadOnly;
            Log(LogLevel::Error, LogCategory::Core, "Trying to write to a readonly IOStream");
            return 0;
        }

        stream->Status = IOStreamStatus::Ready;

        if (!IsValid(inBuffer))
        {
            return 0;
        }

        size_t writtenBytes = stream->Interface.Write(stream->Data, inBuffer, &stream->Status);
        if ((writtenBytes == 0) && (stream->Status == IOStreamStatus::Ready))
        {
            stream->Status = IOStreamStatus::Error;
        }
        return writtenBytes;
    }

    size_t IORead(NonNullPtr<IOStream> stream, void* ptr, size_t size)
    {
        if (!stream->Interface.Read)
        {
            Log(LogLevel::Error, LogCategory::Core, "Trying to read a writeonly IOStream");
            stream->Status = IOStreamStatus::WriteOnly;
            return 0;
        }

        stream->Status = IOStreamStatus::Ready;
        if (size == 0)
        {
            return 0;
        }
        size_t bytes = stream->Interface.Read(stream->Data, ptr, size, &stream->Status);
        if (bytes == 0 && stream->Status == IOStreamStatus::Ready)
        {
            // TODO: Detect error and do this stream->Status = IOStreamStatus::Error;
            stream->Status = IOStreamStatus::EndOfFile;
        }
        return bytes;
    }

    int64 IOSeek(NonNullPtr<IOStream> stream, int64 offset, IOStreamSeekPivot pivot)
    {
        if (!stream->Interface.Seek)
        {
            Log(LogLevel::Error, LogCategory::Core, "IOSeek: Stream interface cannot Seek");
            return -1;
        }
        return stream->Interface.Seek(stream->Data, offset, pivot);
    }

    int64 IOSize(NonNullPtr<IOStream> stream)
    {
        if (!stream->Interface.Size)
        {
            int64 pos = IOSeek(stream, 0, IOStreamSeekPivot::Current);
            if (pos < 0)
            {
                return -1;
            }
            int64 size = IOSeek(stream, 0, IOStreamSeekPivot::End);

            IOSeek(stream, pos, IOStreamSeekPivot::Begin);
            return size;
        }
        return stream->Interface.Size(stream->Data);
    }

    ByteBuffer LoadFile(String filename)
    {
        IOStream* stream = IOFromFile(filename, WrapString("rb"));
        if (!stream)
        {
            return {};
        }
        return LoadFile(stream, true);
    }

    ByteBuffer LoadFile(NonNullPtr<IOStream> stream, bool closeStreamWhenDone)
    {
        constexpr size_t kFileChunkSize = 1024;
        uint8*           data           = nullptr;
        uint8*           newData        = nullptr;
        size_t           totalSize      = 0;
        ByteBuffer       resultBuffer   = {};

        auto deferred = Defer(
            [&]()
            {
                if (closeStreamWhenDone)
                {
                    IOClose(stream);
                }
            });

        // Try reading the size from the stream, if failing we'll try to read it chunk by chunk
        bool  loadChunks = false;
        int64 ssize      = IOSize(stream);
        if (ssize < 0)
        {
            ssize      = kFileChunkSize;
            loadChunks = true;
        }
        size_t size = static_cast<size_t>(ssize);
        data        = static_cast<uint8*>(Malloc(static_cast<size_t>(size + 1)));
        if (!data)
        {
            return {};
        }

        while (true)
        {
            if (loadChunks)
            {
                if ((totalSize + kFileChunkSize) > size)
                {
                    size    = totalSize + kFileChunkSize;
                    newData = static_cast<uint8*>(Realloc(data, static_cast<size_t>(size + 1)));
                    if (!newData)
                    {
                        Free(data);
                        return {};
                    }
                    data = newData;
                }
            }

            size_t sizeRead = IORead(stream, data + totalSize, (size - totalSize));
            if (sizeRead > 0)
            {
                totalSize += sizeRead;
                continue;
            }
            if (stream->Status == IOStreamStatus::NotReady)
            {
                // Wait for the stream to be ready
                wait_ms(1);
                continue;
            }

            // The stream status will remain set for the caller to check
            break;
        }

        // Adding null terminator
        data[totalSize] = '\0';

        resultBuffer.Data = reinterpret_cast<Byte*>(data);
        resultBuffer.Size = totalSize;

        return resultBuffer;
    }

    bool IOClose(NonNullPtr<IOStream> stream)
    {
        bool result = true;
        if (stream->Interface.Close)
        {
            result = stream->Interface.Close(stream->Data);
        }
        Free(stream.Get());
        return result;
    }

} // namespace Juliet
