win-drives 0.1.0

use heapless::{String, format};
use winapi::shared::ntdef::HANDLE;

use crate::{DiskGeometry, DriverError, PartitionInfo, api::*};

#[derive(Debug)]
pub struct HarddiskVolume {
    handle: HANDLE,
    pub geometry: DiskGeometry,
    pub partition_info: PartitionInfo,
}

impl HarddiskVolume {

    pub fn enumerate() -> HarddiskVolumeIter {
        HarddiskVolumeIter {
            next_num: 0,
            consecutive_misses: 0,
            done: false,
        }
    }

    pub fn open(volume_num: u8) -> Result<Self, DriverError> {
        let obj_name: String<_> = format!(20; r#"\??\HarddiskVolume{}"#, volume_num).unwrap();
        let handle = open_handle(obj_name)?;

        Ok(HarddiskVolume {
            handle,
            geometry: geometry(handle)?.into(),
            partition_info: partition_info(handle)?.into(),
        })
    }

    pub fn read_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<usize, DriverError> {
        let sector_size = self.geometry.bytes_per_sector() as usize;
        if buf.len() < sector_size || buf.len() % sector_size != 0 {
            return Err(DriverError::InvalidParameter);
        }

        let read = read_file(self.handle, buf, offset)?;

        Ok(read)
    }

    /// Logical size of the volume in bytes (partition extent).
    #[inline]
    pub fn size(&self) -> u64 {
        self.partition_info.length()
    }

    /// Access the partition info (style, offset, length, MBR/GPT details).
    #[inline]
    pub fn partition_info(&self) -> &PartitionInfo {
        &self.partition_info
    }

    /// Returns the number of bytes per sector reported by the drive geometry.
    #[inline]
    pub fn bytes_per_sector(&self) -> u64 {
        self.geometry.bytes_per_sector()
    }

    /// Returns the number of cylinders reported by the drive geometry.
    #[inline]
    pub fn cylinders(&self) -> u64 {
        self.geometry.cylinders()
    }

    #[inline]
    pub fn tracks_per_cylinder(&self) -> u32 {
        self.geometry.tracks_per_cylinder()
    }

    /// Returns the number of sectors per track reported by the drive geometry.
    #[inline]
    pub fn sectors_per_track(&self) -> u32 {
        self.geometry.sectors_per_track()
    }

    /// Returns the total number of sectors on the drive.
    #[inline]
    pub fn sectors(&self) -> u64 {
        self.geometry.sectors()
    }

    /// Returns the media type reported by the drive geometry.
    #[inline]
    pub fn media_type(&self) -> u32 {
        self.geometry.media_type()
    }
}

pub struct HarddiskVolumeIter {
    next_num: u8,
    consecutive_misses: u8,
    done: bool,
}

impl HarddiskVolumeIter {
    /// Highest volume number to probe before giving up.
    const MAX_VOLUMES: u8 = 64;
    /// Number of consecutive `NotFound`/`PathNotFound` results before stopping.
    const MISS_THRESHOLD: u8 = 8;
}

impl Iterator for HarddiskVolumeIter {
    type Item = HarddiskVolume;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }

        while self.next_num < Self::MAX_VOLUMES {
            let num = self.next_num;
            self.next_num += 1;

            match HarddiskVolume::open(num) {
                Ok(vol) => {
                    self.consecutive_misses = 0;
                    return Some(vol);
                }
                Err(DriverError::NotFound) | Err(DriverError::PathNotFound) => {
                    self.consecutive_misses += 1;
                    if self.consecutive_misses >= Self::MISS_THRESHOLD {
                        self.done = true;
                        return None;
                    }
                }
                Err(DriverError::Permission) => {
                    // Volume exists but we can't open it — keep going.
                    self.consecutive_misses = 0;
                    continue;
                }
                Err(_) => {
                    self.done = true;
                    return None;
                }
            }
        }

        self.done = true;
        None
    }
}