#include "IoCtl.hpp"

#ifdef _WIN32

#define _NTSCSI_USER_MODE_ 1
#include <windows.h>
#include <winioctl.h>
#include <ntddcdrm.h>
#include <ntddcdvd.h>
#include <ntddmmc.h>
#include <ntddscsi.h>
#include <stddef.h>
#include <stdint.h>
#include <scsi.h>

#endif

#include <dvdcss.h>
#include <iostream>
#include <string>

namespace Li::Scsi {
    IoCtl::IoCtl(std::string drive) {
        this->dvdCssHandle = dvdcss_open(drive.c_str());
        this->osFileHandle = dvdcss_get_raw_fd(this->dvdCssHandle);
    }

    IoCtl::~IoCtl() {
        dvdcss_close(this->dvdCssHandle);
        this->dvdCssHandle = NULL;
        this->osFileHandle = NULL;
    }

    std::vector<uint32_t>* IoCtl::GetSupportedReadSpeeds() {
        std::vector<uint32_t>* speeds = new std::vector<uint32_t>();
#ifdef _WIN32
        CDROM_PERFORMANCE_REQUEST perfRequest;
        memset(&perfRequest, 0x00, sizeof(CDROM_PERFORMANCE_REQUEST));
        perfRequest.RequestType = CdromPerformanceRequest;
        perfRequest.PerformanceType = CdromReadPerformance;
        perfRequest.Exceptions = CdromNominalPerformance;
        perfRequest.Tolerance = Cdrom10Nominal20Exceptions;
        perfRequest.StaringLba = 0;

        uint64_t sz = sizeof(CDROM_PERFORMANCE_HEADER) * 0x1000;
        byte* buffer = new byte[sz];
        memset(buffer, 0x00, sz);

        CDROM_PERFORMANCE_HEADER* perfHeader = (CDROM_PERFORMANCE_HEADER*)buffer;
        

        DWORD unused;
        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, IOCTL_CDROM_GET_PERFORMANCE, &perfRequest, sizeof(CDROM_PERFORMANCE_REQUEST), perfHeader, sz, &unused, nullptr);
        if (!success) {
            std::cerr << "Error getting supported read speeds: " << std::to_string(GetLastError()) << std::endl;
        }

        uint64_t dataLen = _byteswap_ulong(*(uint32_t*)perfHeader->DataLength);
        uint64_t count = dataLen / sizeof(CDROM_NOMINAL_PERFORMANCE_DESCRIPTOR);
        CDROM_NOMINAL_PERFORMANCE_DESCRIPTOR* pptr = (CDROM_NOMINAL_PERFORMANCE_DESCRIPTOR*)perfHeader->Data;
        for (uint64_t i = 0; i < count; i++) {
            CDROM_NOMINAL_PERFORMANCE_DESCRIPTOR perfDesc = pptr[i];
            uint32_t startPerf = _byteswap_ulong(*(uint32_t*)perfDesc.StartPerformance);
            uint32_t endPerf = _byteswap_ulong(*(uint32_t*)perfDesc.EndPerformance);

            bool startAlreadyInList = false;
            bool endAlreadyInList = false;

            for (uint32_t speed : *speeds) {
                if (speed == startPerf) {
                    startAlreadyInList = true;
                }
                if (speed == startPerf) {
                    endAlreadyInList = true;
                }
            }

            if(!startAlreadyInList)
                speeds->push_back(startPerf);
    
            if(!endAlreadyInList)
                speeds->push_back(endPerf);
        }

        delete buffer;

#else
#error no way to get all supported read speeds on this platform
#endif
        return speeds;
    }

    uint64_t IoCtl::GetTotalSectors() {
#ifdef _WIN32
        DISK_GEOMETRY_EX geometry;
        memset(&geometry, 0x00, sizeof(DISK_GEOMETRY_EX));

        DWORD unused;
        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, nullptr, 0, &geometry, sizeof(DISK_GEOMETRY_EX), &unused, nullptr);
        if (!success) {
            std::cerr << "Error getting geometry: " << std::to_string(GetLastError()) << std::endl;
        }
        return (uint64_t)(geometry.DiskSize.QuadPart / DVDCSS_BLOCK_SIZE);
#else
#error no way to get total sector count of disc on this platform ..
#endif
    }
    bool IoCtl::AllowReadingPastDisc() {
#ifdef _WIN32
        // On windows, you cannot read anything outside the last session of the disc, unless you enable this
        // this means you cannot read the Lead-Out, *or* double densitity DVDs
        DWORD unused;
        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, FSCTL_ALLOW_EXTENDED_DASD_IO, nullptr, 0, nullptr, 0, &unused, nullptr);
        if (!success) {
            std::cerr << "Error enabling DASD I/O: " << std::to_string(GetLastError()) << std::endl;
        }
        return success;
#else
#error no way to disable bounds checks (DASHD) i/o on this platform (maybe not required?)
#endif
    }

    bool IoCtl::UnmountVolume() {
#ifdef _WIN32
        DWORD unused;
        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, FSCTL_LOCK_VOLUME, nullptr, 0, nullptr, 0, &unused, nullptr);
        if (!success) {
            std::cerr << "Error unmounting drive: " << std::to_string(GetLastError()) << std::endl;
        }
        return success;
#else
#error no way to unmount drive on this platform
#endif
    }

    bool IoCtl::MountVolume() {
#ifdef _WIN32
        DWORD unused;
        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, FSCTL_UNLOCK_VOLUME, nullptr, 0, nullptr, 0, &unused, nullptr);
        if (!success) {
            std::cerr << "Error mounting drive: " << std::to_string(GetLastError()) << std::endl;
        }
        return success;
#else
#error no way to mount drive on this platform
#endif
    }

    bool IoCtl::ExclusiveLockDrive() {
#ifdef _WIN32

        this->UnmountVolume();

        DWORD unused;
        CDROM_EXCLUSIVE_LOCK exclusiveLockRequest;
        memset(&exclusiveLockRequest, 0x00, sizeof(CDROM_EXCLUSIVE_LOCK));
        
        exclusiveLockRequest.Access.RequestType = ExclusiveAccessLockDevice;
        exclusiveLockRequest.Access.Flags = CDROM_LOCK_IGNORE_VOLUME;

        strncpy((char*)(exclusiveLockRequest.CallerName), "LiDeeVeeDeeDumperProgramThing", CDROM_EXCLUSIVE_CALLER_LENGTH - 1);

        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, IOCTL_CDROM_EXCLUSIVE_ACCESS, &exclusiveLockRequest, sizeof(CDROM_EXCLUSIVE_LOCK), nullptr, 0, &unused, nullptr);
        if (!success) {
            std::cerr << "Error Locking the drive: " << std::to_string(GetLastError()) << std::endl;
        }
        return success;
#else
#error no way to exclusive lock drive on this platform.
#endif
    }

    bool IoCtl::ExclusiveUnlockDrive() {
#ifdef _WIN32
        DWORD unused;
        CDROM_EXCLUSIVE_ACCESS exclusiveUnlockRequest;
        memset(&exclusiveUnlockRequest, 0x00, sizeof(CDROM_EXCLUSIVE_ACCESS));
        exclusiveUnlockRequest.RequestType = ExclusiveAccessUnlockDevice;
        
        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, IOCTL_CDROM_EXCLUSIVE_ACCESS, &exclusiveUnlockRequest, sizeof(CDROM_EXCLUSIVE_LOCK), nullptr, 0, &unused, nullptr);
        if (!success) {
            std::cerr << "Error unlocking the drive: " << std::to_string(GetLastError()) << std::endl;
        }

        this->MountVolume();
        return success;
#else
#error no way to exclusive unlock drive on this platform.
#endif
    }

    bool IoCtl::SetDriveSpeed(int readSpeed, int writeSpeed)
    {
#ifdef _WIN32
        DWORD unused;
        CDROM_SET_SPEED cdromSetSpeed;
        memset(&cdromSetSpeed, 0x00, sizeof(CDROM_SET_SPEED));

        cdromSetSpeed.RequestType = CdromSetSpeed;
        cdromSetSpeed.ReadSpeed = (uint16_t)readSpeed;
        cdromSetSpeed.WriteSpeed = (uint16_t)writeSpeed;
        cdromSetSpeed.RotationControl = CdromDefaultRotation;

        BOOL success = DeviceIoControl((HANDLE)this->osFileHandle, IOCTL_CDROM_SET_SPEED, &cdromSetSpeed, sizeof(CDROM_SET_SPEED), nullptr, 0, &unused, nullptr);
        if (!success) {
            std::cerr << "Error setting speed: " << std::to_string(GetLastError()) << std::endl;
        }
        return success;
#else
#error no way to set drive speed folder on this platform.
#endif
    }


    dvdcss_t IoCtl::GetDvdCssHandle() {
        return this->dvdCssHandle;
    }

}