#include <bzsddk/bzsddk.h>
#include <bzsddk/crossplatform/crossvol.h>
#define RAMDISK_SECTOR_SIZE 512

namespace BazisLib
{
    class RamDisk : public IBasicDisk
    {
    protected:
        ULONGLONG m_SectorCount;
        ULONGLONG m_TotalSize;
        char *m_pBuffer;

    public:
        unsigned GetSectorSize() {return RAMDISK_SECTOR_SIZE;}
        void Dispose();

        RamDisk(unsigned MegabyteCount)
        {
            m_SectorCount = MegabyteCount * (1024 * 1024 / RAMDISK_SECTOR_SIZE);
            m_TotalSize = MegabyteCount * 1024 * 1024;
            m_pBuffer = (char *)bulk_malloc((unsigned)m_TotalSize);
        }

        ~RamDisk()
        {
            if (m_pBuffer)
                bulk_free(m_pBuffer, m_TotalSize);
        }

        virtual ULONGLONG GetSectorCount() override {return m_SectorCount; }

        virtual unsigned Read(ULONGLONG ByteOffset, void *pBuffer, unsigned Length) override
        {
            if (!m_pBuffer)
                return 0;
            LONGLONG MaxSize = m_TotalSize - ByteOffset;
            if (MaxSize < 0)
                MaxSize = 0;
            if (Length > MaxSize)
                Length = MaxSize;
            memcpy(pBuffer, m_pBuffer + ByteOffset, Length);
            return Length;
        }

        virtual unsigned Write(ULONGLONG ByteOffset, const void *pBuffer, unsigned Length) override
        {
            if (!m_pBuffer)
                return 0;
            LONGLONG MaxSize = m_TotalSize - ByteOffset;
            if (MaxSize < 0)
                MaxSize = 0;
            if (Length > MaxSize)
                Length = MaxSize;
            memcpy(m_pBuffer + ByteOffset, pBuffer, Length);
            return Length;
        }

        virtual LPGUID GetStableGuid() override {return NULL;}

        virtual bool DeviceControl(unsigned CtlCode, void *pBuffer, unsigned InSize,
                                                 unsigned OutSize, unsigned *pBytesDone) override
        {
            return false;
        }

        virtual bool Initialize() override {return true;}
    };

    using namespace DDK;

    class RamDiskDriver : public Driver
    {
    public:
        RamDiskDriver() : DDK::Driver(true)
        {
        }

        virtual NTSTATUS AddDevice(IN PDEVICE_OBJECT PhysicalDeviceObject) override
        {
            UniversalVolume *pVol = new UniversalVolume(new RamDisk(256), true);
            if (!pVol)
                return STATUS_NO_MEMORY;
            NTSTATUS st = pVol->AddDevice(this, PhysicalDeviceObject);
            if (!NT_SUCCESS(st))
                return st;
    
            return STATUS_SUCCESS;
        }

        virtual ~RamDiskDriver()
        {
        }
    };

    DDK::Driver *_stdcall CreateMainDriverInstance()
    {
        return new RamDiskDriver();
    }
}