/*
 * Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

// Disk related functions. Original code was refactored and moved
// from psutil/_psutil_osx.c in 2023. This is the GIT blame before the move:
// https://github.com/giampaolo/psutil/blame/efd7ed3/psutil/_psutil_osx.c

#include <Python.h>
#include <sys/param.h>
#include <sys/ucred.h>
#include <sys/mount.h>
#include <CoreFoundation/CoreFoundation.h>
#include <IOKit/IOKitLib.h>
#include <IOKit/storage/IOBlockStorageDriver.h>
#include <IOKit/storage/IOMedia.h>
#include <IOKit/IOBSD.h>

#include "../../arch/all/init.h"


/*
 * Return a list of tuples including device, mount point and fs type
 * for all partitions mounted on the system.
 */
PyObject *
psutil_disk_partitions(PyObject *self, PyObject *args) {
    int num;
    int i;
    int len;
    uint64_t flags;
    char opts[400];
    struct statfs *fs = NULL;
    PyObject *py_dev = NULL;
    PyObject *py_mountp = NULL;
    PyObject *py_tuple = NULL;
    PyObject *py_retlist = PyList_New(0);

    if (py_retlist == NULL)
        return NULL;

    // get the number of mount points
    Py_BEGIN_ALLOW_THREADS
    num = getfsstat(NULL, 0, MNT_NOWAIT);
    Py_END_ALLOW_THREADS
    if (num == -1) {
        PyErr_SetFromErrno(PyExc_OSError);
        goto error;
    }

    len = sizeof(*fs) * num;
    fs = malloc(len);
    if (fs == NULL) {
        PyErr_NoMemory();
        goto error;
    }

    Py_BEGIN_ALLOW_THREADS
    num = getfsstat(fs, len, MNT_NOWAIT);
    Py_END_ALLOW_THREADS
    if (num == -1) {
        PyErr_SetFromErrno(PyExc_OSError);
        goto error;
    }

    for (i = 0; i < num; i++) {
        opts[0] = 0;
        flags = fs[i].f_flags;

        // see sys/mount.h
        if (flags & MNT_RDONLY)
            strlcat(opts, "ro", sizeof(opts));
        else
            strlcat(opts, "rw", sizeof(opts));
        if (flags & MNT_SYNCHRONOUS)
            strlcat(opts, ",sync", sizeof(opts));
        if (flags & MNT_NOEXEC)
            strlcat(opts, ",noexec", sizeof(opts));
        if (flags & MNT_NOSUID)
            strlcat(opts, ",nosuid", sizeof(opts));
        if (flags & MNT_UNION)
            strlcat(opts, ",union", sizeof(opts));
        if (flags & MNT_ASYNC)
            strlcat(opts, ",async", sizeof(opts));
        if (flags & MNT_EXPORTED)
            strlcat(opts, ",exported", sizeof(opts));
        if (flags & MNT_LOCAL)
            strlcat(opts, ",local", sizeof(opts));
        if (flags & MNT_QUOTA)
            strlcat(opts, ",quota", sizeof(opts));
        if (flags & MNT_ROOTFS)
            strlcat(opts, ",rootfs", sizeof(opts));
        if (flags & MNT_DOVOLFS)
            strlcat(opts, ",dovolfs", sizeof(opts));
        if (flags & MNT_DONTBROWSE)
            strlcat(opts, ",dontbrowse", sizeof(opts));
        if (flags & MNT_IGNORE_OWNERSHIP)
            strlcat(opts, ",ignore-ownership", sizeof(opts));
        if (flags & MNT_AUTOMOUNTED)
            strlcat(opts, ",automounted", sizeof(opts));
        if (flags & MNT_JOURNALED)
            strlcat(opts, ",journaled", sizeof(opts));
        if (flags & MNT_NOUSERXATTR)
            strlcat(opts, ",nouserxattr", sizeof(opts));
        if (flags & MNT_DEFWRITE)
            strlcat(opts, ",defwrite", sizeof(opts));
        if (flags & MNT_UPDATE)
            strlcat(opts, ",update", sizeof(opts));
        if (flags & MNT_RELOAD)
            strlcat(opts, ",reload", sizeof(opts));
        if (flags & MNT_FORCE)
            strlcat(opts, ",force", sizeof(opts));
        if (flags & MNT_CMDFLAGS)
            strlcat(opts, ",cmdflags", sizeof(opts));
        // requires macOS >= 10.5
#ifdef MNT_QUARANTINE
        if (flags & MNT_QUARANTINE)
            strlcat(opts, ",quarantine", sizeof(opts));
#endif
#ifdef MNT_MULTILABEL
        if (flags & MNT_MULTILABEL)
            strlcat(opts, ",multilabel", sizeof(opts));
#endif
#ifdef MNT_NOATIME
        if (flags & MNT_NOATIME)
            strlcat(opts, ",noatime", sizeof(opts));
#endif
        py_dev = PyUnicode_DecodeFSDefault(fs[i].f_mntfromname);
        if (! py_dev)
            goto error;
        py_mountp = PyUnicode_DecodeFSDefault(fs[i].f_mntonname);
        if (! py_mountp)
            goto error;
        py_tuple = Py_BuildValue(
            "(OOss)",
            py_dev,               // device
            py_mountp,            // mount point
            fs[i].f_fstypename,   // fs type
            opts);                // options
        if (!py_tuple)
            goto error;
        if (PyList_Append(py_retlist, py_tuple))
            goto error;
        Py_CLEAR(py_dev);
        Py_CLEAR(py_mountp);
        Py_CLEAR(py_tuple);
    }

    free(fs);
    return py_retlist;

error:
    Py_XDECREF(py_dev);
    Py_XDECREF(py_mountp);
    Py_XDECREF(py_tuple);
    Py_DECREF(py_retlist);
    if (fs != NULL)
        free(fs);
    return NULL;
}


PyObject *
psutil_disk_usage_used(PyObject *self, PyObject *args) {
    PyObject *py_default_value;
    PyObject *py_mount_point_bytes = NULL;
    char* mount_point;

    if (!PyArg_ParseTuple(args, "O&O", PyUnicode_FSConverter, &py_mount_point_bytes, &py_default_value)) {
        return NULL;
    }
    mount_point = PyBytes_AsString(py_mount_point_bytes);
    if (NULL == mount_point) {
        Py_XDECREF(py_mount_point_bytes);
        return NULL;
    }

#ifdef ATTR_VOL_SPACEUSED
    /* Call getattrlist(ATTR_VOL_SPACEUSED) to get used space info. */
    int ret;
    struct {
        uint32_t size;
        uint64_t spaceused;
    } __attribute__((aligned(4), packed)) attrbuf = {0};
    struct attrlist attrs = {0};

    attrs.bitmapcount = ATTR_BIT_MAP_COUNT;
    attrs.volattr = ATTR_VOL_INFO | ATTR_VOL_SPACEUSED;
    Py_BEGIN_ALLOW_THREADS
    ret = getattrlist(mount_point, &attrs, &attrbuf, sizeof(attrbuf), 0);
    Py_END_ALLOW_THREADS
    if (ret == 0) {
        Py_XDECREF(py_mount_point_bytes);
        return PyLong_FromUnsignedLongLong(attrbuf.spaceused);
    }
    psutil_debug("getattrlist(ATTR_VOL_SPACEUSED) failed, fall-back to default value");
#endif
    Py_XDECREF(py_mount_point_bytes);
    Py_INCREF(py_default_value);
    return py_default_value;
}


/*
 * Return a Python dict of tuples for disk I/O information
 */
PyObject *
psutil_disk_io_counters(PyObject *self, PyObject *args) {
    CFDictionaryRef parent_dict;
    CFDictionaryRef props_dict;
    CFDictionaryRef stats_dict;
    io_registry_entry_t parent;
    io_registry_entry_t disk;
    io_iterator_t disk_list = 0;
    PyObject *py_disk_info = NULL;
    PyObject *py_retdict = PyDict_New();

    if (py_retdict == NULL)
        return NULL;

    if (IOServiceGetMatchingServices(
            kIOMasterPortDefault,
            IOServiceMatching(kIOMediaClass),
            &disk_list) != kIOReturnSuccess)
    {
        PyErr_SetString(PyExc_RuntimeError, "unable to get the list of disks");
        goto error;
    }

    while ((disk = IOIteratorNext(disk_list)) != 0) {
        py_disk_info = NULL;
        parent_dict = NULL;
        props_dict = NULL;
        stats_dict = NULL;
        parent = 0;

        if (IORegistryEntryGetParentEntry(disk, kIOServicePlane, &parent) != kIOReturnSuccess) {
            PyErr_SetString(PyExc_RuntimeError, "unable to get the disk's parent");
            IOObjectRelease(disk);
            goto error;
        }

        if (!IOObjectConformsTo(parent, "IOBlockStorageDriver")) {
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            continue;
        }

        if (IORegistryEntryCreateCFProperties(disk,
                (CFMutableDictionaryRef *)&parent_dict,
                kCFAllocatorDefault, kNilOptions) != kIOReturnSuccess)
        {
            PyErr_SetString(
                PyExc_RuntimeError, "unable to get the parent's properties"
            );
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        if (IORegistryEntryCreateCFProperties(parent,
                (CFMutableDictionaryRef *)&props_dict,
                kCFAllocatorDefault, kNilOptions) != kIOReturnSuccess)
        {
            PyErr_SetString(
                PyExc_RuntimeError, "unable to get the disk properties"
            );
            if (parent_dict) CFRelease(parent_dict);
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        CFStringRef disk_name_ref = (CFStringRef)CFDictionaryGetValue(
            parent_dict, CFSTR(kIOBSDNameKey)
        );
        if (disk_name_ref == NULL) {
            PyErr_SetString(PyExc_RuntimeError, "unable to get disk name");
            if (parent_dict)
                CFRelease(parent_dict);
            if (props_dict)
                CFRelease(props_dict);
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        const int kMaxDiskNameSize = 64;
        char disk_name[kMaxDiskNameSize];
        if (!CFStringGetCString(disk_name_ref, disk_name, kMaxDiskNameSize,
                                CFStringGetSystemEncoding())) {
            PyErr_SetString(
                PyExc_RuntimeError, "unable to convert disk name to C string"
            );
            if (parent_dict)
                CFRelease(parent_dict);
            if (props_dict)
                CFRelease(props_dict);
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        stats_dict = (CFDictionaryRef)CFDictionaryGetValue(
            props_dict, CFSTR(kIOBlockStorageDriverStatisticsKey)
        );
        if (stats_dict == NULL) {
            PyErr_SetString(PyExc_RuntimeError, "unable to get disk stats");
            if (parent_dict)
                CFRelease(parent_dict);
            if (props_dict)
                CFRelease(props_dict);
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        CFNumberRef number;
        int64_t reads = 0, writes = 0, read_bytes = 0, write_bytes = 0;
        int64_t read_time = 0, write_time = 0;

        if ((number = (CFNumberRef)CFDictionaryGetValue(
                stats_dict, CFSTR(kIOBlockStorageDriverStatisticsReadsKey))))
            CFNumberGetValue(number, kCFNumberSInt64Type, &reads);

        if ((number = (CFNumberRef)CFDictionaryGetValue(
                stats_dict, CFSTR(kIOBlockStorageDriverStatisticsWritesKey))))
            CFNumberGetValue(number, kCFNumberSInt64Type, &writes);

        if ((number = (CFNumberRef)CFDictionaryGetValue(
                stats_dict, CFSTR(kIOBlockStorageDriverStatisticsBytesReadKey))))
            CFNumberGetValue(number, kCFNumberSInt64Type, &read_bytes);

        if ((number = (CFNumberRef)CFDictionaryGetValue(
                stats_dict, CFSTR(kIOBlockStorageDriverStatisticsBytesWrittenKey))))
            CFNumberGetValue(number, kCFNumberSInt64Type, &write_bytes);

        if ((number = (CFNumberRef)CFDictionaryGetValue(
                stats_dict, CFSTR(kIOBlockStorageDriverStatisticsTotalReadTimeKey))))
            CFNumberGetValue(number, kCFNumberSInt64Type, &read_time);

        if ((number = (CFNumberRef)CFDictionaryGetValue(
                stats_dict, CFSTR(kIOBlockStorageDriverStatisticsTotalWriteTimeKey))))
            CFNumberGetValue(number, kCFNumberSInt64Type, &write_time);

        py_disk_info = Py_BuildValue(
            "(KKKKKK)",
            (unsigned long long)reads,
            (unsigned long long)writes,
            (unsigned long long)read_bytes,
            (unsigned long long)write_bytes,
            (unsigned long long)(read_time / 1000 / 1000),
            (unsigned long long)(write_time / 1000 / 1000)
        );

        if (!py_disk_info) {
            if (parent_dict)
                CFRelease(parent_dict);
            if (props_dict)
                CFRelease(props_dict);
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        if (PyDict_SetItemString(py_retdict, disk_name, py_disk_info)) {
            Py_CLEAR(py_disk_info);
            if (parent_dict)
                CFRelease(parent_dict);
            if (props_dict)
                CFRelease(props_dict);
            IOObjectRelease(parent);
            IOObjectRelease(disk);
            goto error;
        }

        Py_CLEAR(py_disk_info);

        if (parent_dict)
            CFRelease(parent_dict);
        IOObjectRelease(parent);
        if (props_dict)
            CFRelease(props_dict);
        IOObjectRelease(disk);
    }

    IOObjectRelease(disk_list);
    return py_retdict;

error:
    Py_XDECREF(py_disk_info);
    Py_DECREF(py_retdict);
    if (disk_list != 0)
        IOObjectRelease(disk_list);
    return NULL;
}
