//Created by Winson on 23/01/2025
//Open NVMe Device and Send NVMe Commands

#include 
#include 
#include 
#include 
#include 
#include 
#include 

class NVMeDeviceHandler {
private:
    std::string deviceName;
    io_service_t deviceService;
    io_connect_t ioServiceConnection;

public:
    NVMeDeviceHandler(const std::string& name) : deviceName(name), deviceService(0), ioServiceConnection(0) {}

    ~NVMeDeviceHandler() {
        if (ioServiceConnection) {
            IOServiceClose(ioServiceConnection);
        }
        if (deviceService) {
            IOObjectRelease(deviceService);
        }
    }

    bool openDevice() {
        
        //Step 1: Match the I/O service for IOMedia devices
        CFMutableDictionaryRef matchingDict = IOServiceMatching("IOMedia");
        if (!matchingDict) {
            std::cerr << "Failed to create matching dictionary for IOMedia" << std::endl;
            return false;
        }

        // Step 2: Add device name filter to cross check disk# enter from command prompt is valid disk num.
        CFStringRef devName = CFStringCreateWithCString(kCFAllocatorDefault, deviceName.c_str(), kCFStringEncodingUTF8);
        CFDictionarySetValue(matchingDict, CFSTR("BSD Name"), devName); //BSD Name are ties to IOMedia objects(disk0,1,2,etc)
        CFRelease(devName);

        //Step 3: Get the matching device service
        deviceService = IOServiceGetMatchingService(kIOMainPortDefault, matchingDict);
        if (!deviceService) {
            std::cerr << "Physical device " << deviceName << " not found." << std::endl;
            return false;
        }
        
        // Step 4: Re-match again the I/O service for NVMe controller devices
        CFMutableDictionaryRef matchingDict1 = IOServiceMatching("IONVMeController");
        if (!matchingDict1) {
            std::cerr << "Error: Could not create matching dictionary." << std::endl;
            return false;
        }

        // Step 5: Get the matching device service
        io_service_t service = IOServiceGetMatchingService(kIOMainPortDefault, matchingDict1);
        if (!service) {
            std::cerr << "Error: Could not find NVMe controller device." << std::endl;
            return false;
        }

        // Step 6: Open a connection to the NVMe device
        kern_return_t result = IOServiceOpen(service, mach_task_self(), 0, &ioServiceConnection);
        if (result != KERN_SUCCESS) {
            std::cerr << "Failed to open device with error: "
                      << result << " (" << mach_error_string(result) << ")" << std::endl;
            IOObjectRelease(service);
            return false;
                
        }

        std::cout << "Successfully opened device " << deviceName << "." << std::endl;
        return true;
    }

    void performActions() {
		//TODO Data Wipe command
        // Placeholder for NVMe-specific commands< "..." << std::endl;
    }
    
    bool sendNVMeAdminCommand() {

        // Construct the NVMe Identify Controller command
        struct NVMeIdentifyCommand {
            uint32_t opcode;   // NVMe Admin Command Opcode
            uint32_t nsid;     // Namespace Identifier
            uint64_t reserved; // Reserved
        };

        
        NVMeIdentifyCommand identifyCommand = {};
        identifyCommand.opcode = 0x06; // Identify command
        identifyCommand.nsid = 0;      // For Identify Controller

        // Prepare input and output buffers
        uint64_t input[] = {identifyCommand.opcode, identifyCommand.nsid};
        uint32_t inputCount = sizeof(input) / sizeof(uint64_t);

        uint64_t output[16] = {}; // Response buffer for Identify data
        uint32_t outputCount = sizeof(output) / sizeof(uint64_t);

        // Call IOConnectCallMethod
        kern_return_t result = IOConnectCallMethod(
            ioServiceConnection, // Connection
            1,                   // Method selector (verify the correct ID in documentation)
            input,               // Input data
            inputCount,          // Input data count
            nullptr,             // No additional input structure
            0,                   // Size of additional input
            output,              // Output buffer
            &outputCount,        // Output buffer size
            nullptr,             // No additional output structure
            nullptr              // Size of additional output
        );

        // Check the result
        if (result != KERN_SUCCESS) {
            std::cerr << "Failed to send NVMe Admin Command with error: "
                      << result << " (" << mach_error_string(result) << ")" << std::endl;
            return false;
        }

        // Successfully sent the command and received the response
        std::cout << "NVMe Admin Command sent successfully. Processing the response..." << std::endl;

        // Optionally, print the output response data
        std::cout << "Response Data: ";
        for (size_t i = 0; i < outputCount; ++i) {
            std::cout << output[i] << " ";
        }
        std::cout << std::endl;

        return true;
    }
};

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " " << std::endl;
        return 1;
    }

    std::string deviceName = argv[1];
    NVMeDeviceHandler deviceHandler(deviceName);

    if (!deviceHandler.openDevice()) {
        return 1;
    }

	//Todo
    deviceHandler.performActions();
    
	//Proof the the connection is worked 
    // Send NVMe Admin Command to test device interaction
    if (deviceHandler.sendNVMeAdminCommand()) {
        std::cout << "NVMe Admin Command was successful!" << std::endl;
    } else {
        std::cerr << "NVMe Admin Command failed." << std::endl;
    }

    return 0;
}