// Take a look at the license at the top of the repository in the LICENSE file.

use crate::{DiskExt, DiskType};

use std::ffi::{OsStr, OsString};
use std::mem::size_of;
use std::path::Path;

use winapi::ctypes::c_void;
use winapi::shared::minwindef::{DWORD, MAX_PATH};
use winapi::um::fileapi::{
    CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW,
    OPEN_EXISTING,
};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::winbase::{DRIVE_FIXED, DRIVE_REMOVABLE};
use winapi::um::winioctl::{
    PropertyStandardQuery, StorageDeviceSeekPenaltyProperty, IOCTL_STORAGE_QUERY_PROPERTY,
    STORAGE_PROPERTY_QUERY,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE, ULARGE_INTEGER};

#[doc = include_str!("../../md_doc/disk.md")]
pub struct Disk {
    type_: DiskType,
    name: OsString,
    file_system: Vec<u8>,
    mount_point: Vec<u16>,
    s_mount_point: String,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        Path::new(&self.s_mount_point)
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        if self.total_space != 0 {
            unsafe {
                let mut tmp: ULARGE_INTEGER = std::mem::zeroed();
                if GetDiskFreeSpaceExW(
                    self.mount_point.as_ptr(),
                    std::ptr::null_mut(),
                    std::ptr::null_mut(),
                    &mut tmp,
                ) != 0
                {
                    self.available_space = *tmp.QuadPart();
                    return true;
                }
            }
        }
        false
    }
}

struct HandleWrapper(HANDLE);

impl HandleWrapper {
    unsafe fn new(drive_name: &[u16], open_rights: DWORD) -> Option<Self> {
        let handle = CreateFileW(
            drive_name.as_ptr(),
            open_rights,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            std::ptr::null_mut(),
            OPEN_EXISTING,
            0,
            std::ptr::null_mut(),
        );
        if handle == INVALID_HANDLE_VALUE {
            CloseHandle(handle);
            None
        } else {
            Some(Self(handle))
        }
    }
}

impl Drop for HandleWrapper {
    fn drop(&mut self) {
        unsafe {
            CloseHandle(self.0);
        }
    }
}

unsafe fn get_drive_size(mount_point: &[u16]) -> Option<(u64, u64)> {
    let mut total_size: ULARGE_INTEGER = std::mem::zeroed();
    let mut available_space: ULARGE_INTEGER = std::mem::zeroed();
    if GetDiskFreeSpaceExW(
        mount_point.as_ptr(),
        std::ptr::null_mut(),
        &mut total_size,
        &mut available_space,
    ) != 0
    {
        Some((
            *total_size.QuadPart() as _,
            *available_space.QuadPart() as _,
        ))
    } else {
        None
    }
}

// FIXME: To be removed once <https://github.com/retep998/winapi-rs/pull/1028> has been merged.
#[allow(non_snake_case)]
#[repr(C)]
struct DEVICE_SEEK_PENALTY_DESCRIPTOR {
    Version: DWORD,
    Size: DWORD,
    IncursSeekPenalty: BOOLEAN,
}

pub(crate) unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..DWORD::BITS)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            let is_removable = drive_type == DRIVE_REMOVABLE;

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = HandleWrapper::new(&drive_name, 0)?;
            let (total_space, available_space) = get_drive_size(&mount_point)?;
            if total_space == 0 {
                return None;
            }
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: StorageDeviceSeekPenaltyProperty,
                QueryType: PropertyStandardQuery,
                AdditionalParameters: [0],
            };
            let mut result: DEVICE_SEEK_PENALTY_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            let type_ = if DeviceIoControl(
                handle.0,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut result as *mut DEVICE_SEEK_PENALTY_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_SEEK_PENALTY_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_SEEK_PENALTY_DESCRIPTOR>() as DWORD
            {
                DiskType::Unknown(-1)
            } else {
                let is_ssd = result.IncursSeekPenalty == 0;
                if is_ssd {
                    DiskType::SSD
                } else {
                    DiskType::HDD
                }
            };
            Some(Disk {
                type_,
                name: name.to_owned(),
                file_system: file_system.to_vec(),
                mount_point: mount_point.to_vec(),
                s_mount_point: String::from_utf16_lossy(&mount_point[..mount_point.len() - 1]),
                total_space,
                available_space,
                is_removable,
            })
        })
        .collect::<Vec<_>>()
}