use winapi::{
	ctypes::c_void,
	shared::ntdef::LARGE_INTEGER,
    um::{
        fileapi::{self as fs, OPEN_EXISTING},
        handleapi::{CloseHandle, INVALID_HANDLE_VALUE},
		ioapiset::DeviceIoControl,
		winbase::{FILE_BEGIN, FILE_CURRENT, FILE_END, FILE_FLAG_NO_BUFFERING},
        winioctl::{DISK_GEOMETRY, IOCTL_DISK_GET_DRIVE_GEOMETRY},
        winnt::{FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, HANDLE}
    }
};
use std::{
	cmp::min,
	io::{self, Cursor, ErrorKind, Read, Seek, SeekFrom},
    mem::size_of,
    ptr::null_mut
};
use crate::win32;

/// Gives direct access to a physical drive.
///
/// NOTE: because of Win32 restrictions, you must read() whole sectors and must only seek to sector
/// boundaries. If you need random access, use [`BufferedPhysicalDrive`].
#[derive(Debug)]
pub struct PhysicalDrive {
	handle: HANDLE,
	/// The [`DiskGeometry`] of the drive.
	pub geometry: DiskGeometry
}

/// Gives buffered direct access to a physical drive to enable random access.
///
/// Also see [`PhysicalDrive`].
#[derive(Debug)]
pub struct BufferedPhysicalDrive {
	drive: PhysicalDrive,
	current_sector: Cursor<Vec<u8>>,
	current_sector_num: u64,
	/// The sector boundaries of the drive (e. g. to restrict access to a partition of the drive),
	/// where the first value of the tuple denotes the minimum and the second the maximum sector number
	/// (both inclusive). Calls to `seek()` will be considered relative to `sector_boundaries.0` if seeking
	/// from the start or `sector_boundaries.1` if seeking from the end.
	///
	/// NOTE: calls to `geometry.size()` will still return the size of the drive, regardless of the boundaries.
	pub sector_boundaries: (u64, u64),
	/// The [`DiskGeometry`] of the drive.
	pub geometry: DiskGeometry
}

/// Represents the geometry of a disk or physical drive.
#[derive(Copy, Clone, Debug)]
pub struct DiskGeometry {
	pub cylinders: i64,
	pub media_type: u32,
	pub tracks_per_cylinder: u32,
	pub sectors_per_track: u32,
	pub bytes_per_sector: u32
}

impl PhysicalDrive {
	/// Opens the physical drive with the given number.
	///
	/// NOTE: this requires administrator privileges.
	///
	/// Fails if an invalid drive number is given, the user has insufficient privileges
	/// or an error occures while opening the drive.
	/// If an error occurs, an error message containing the error number is returned, see
	/// [here](https://docs.microsoft.com/en-us/windows/win32/debug/system-error-codes) for a list.
	pub fn open(drive_num: u8) -> Result<Self, String> {
		let path = format!("\\\\.\\PhysicalDrive{}", drive_num);
		let path = win32::win32_string(&path);
		let handle = unsafe {
			fs::CreateFileW(path.as_ptr(), GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE, null_mut(),
			OPEN_EXISTING, FILE_FLAG_NO_BUFFERING,
			null_mut())
		};
		if handle == INVALID_HANDLE_VALUE {
			Err(
				format!("invalid handle: error code {} (do you have administrator privileges?)",
				win32::last_error())
			)
		} else {
			Ok(PhysicalDrive { handle, geometry: PhysicalDrive::geometry(&handle)? })
		}
	}
	
	/// Gets the disk geometry via the Win32 API.
	///
	/// If an error occurs, an error message containing the error number is returned, see
	/// [here](https://docs.microsoft.com/en-us/windows/win32/debug/system-error-codes) for a list.
	fn geometry(drive: &HANDLE) -> Result<DiskGeometry, String> {
		let mut geo = Default::default();
		let mut bytes_returned = 0u32;
		let geo_ptr: *mut DISK_GEOMETRY = &mut geo;
		let r = unsafe {
			DeviceIoControl(*drive, IOCTL_DISK_GET_DRIVE_GEOMETRY,
				null_mut(), 0, geo_ptr as *mut c_void,
				size_of::<DISK_GEOMETRY>() as u32, &mut bytes_returned,
				null_mut())
		};
		if r == 0 {
			Err(format!("something went wrong: error code {}", win32::last_error()))
		} else {
			Ok(DiskGeometry::from(geo))
		}
	}
}

impl Read for PhysicalDrive {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
		let sector_size = self.geometry.bytes_per_sector;
		if buf.len() < sector_size as usize {
			return Err(io::Error::new(
				io::ErrorKind::InvalidInput,
				format!("buffer too small, must be at least {} bytes", sector_size)
			));
		}

		let mut bytes_read = 0u32;
	
		let res = unsafe {
			fs::ReadFile(self.handle, buf.as_mut_ptr() as *mut c_void,
			(buf.len() as u32 / sector_size) * sector_size,
			&mut bytes_read, null_mut())
		};
		if res == 0 {
			Err(io::Error::new(
				io::ErrorKind::Other,
				format!(
					"Could not read: error code {} (did you respect sector boundaries?)", win32::last_error()
				)
			))
		} else {
			Ok(bytes_read as usize)
		}
    }
}

impl Seek for PhysicalDrive {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
		fn check_pos(p: u64) -> io::Result<()> {
			if p > i64::MAX as u64 {
				Err(io::Error::new(
					io::ErrorKind::InvalidInput,
					format!("can only seek to pos <= {}", i64::MAX)
				))
			} else {
				Ok(())
			}
		}

		let mut dist = LARGE_INTEGER::default();
		let mut new_pos = LARGE_INTEGER::default();

        let res = match pos {
			SeekFrom::Start(p) => {
				check_pos(p)?;
				unsafe {
					*dist.QuadPart_mut() = p as i64;
					fs::SetFilePointerEx(self.handle, dist,
					&mut new_pos, FILE_BEGIN)
				}
			},
			SeekFrom::Current(p) => {
				unsafe {
					*dist.QuadPart_mut() = p;
					fs::SetFilePointerEx(self.handle, dist,
					&mut new_pos, FILE_CURRENT)
				}
			},
			SeekFrom::End(p) => {
				unsafe {
					*dist.QuadPart_mut() = p;
					fs::SetFilePointerEx(self.handle, dist,
					&mut new_pos, FILE_END)
				}
			}
		};

		if res == 0 {
			Err(io::Error::new(
				io::ErrorKind::Other,
				format!("Could not seek: error code {} (did you respect sector boundaries?)", win32::last_error())
			))
		} else {
			Ok(unsafe { *new_pos.QuadPart() } as u64)
		}
    }
}

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

impl BufferedPhysicalDrive {
	/// Opens the physical drive with the given number.
	///
	/// NOTE: this requires administrator privileges.
	///
	/// Fails if an invalid drive number is given, the user has insufficient privileges
	/// or an error occures while opening the drive.
	/// If an error occurs, an error message containing the error number is returned, see
	/// [here](https://docs.microsoft.com/en-us/windows/win32/debug/system-error-codes) for a list.
	pub fn open(drive_num: u8) -> Result<Self, String> {
		let mut drive = PhysicalDrive::open(drive_num)?;
		let mut sector = vec![0; drive.geometry.bytes_per_sector as usize];
		match drive.read(&mut sector) {
			Ok(_) => (),
			Err(e) => return Err(format!("{}", e))
		}
		let geo = drive.geometry;
		Ok(BufferedPhysicalDrive {
			drive,
			current_sector: Cursor::new(sector),
			current_sector_num: 0,
			sector_boundaries: (0, u64::MAX),
			geometry: geo
		})
	}

	pub fn open_bounded(drive_num: u8, boundaries: (u64, u64)) -> Result<Self, String> {
		if boundaries.1 <= boundaries.0 {
			return Err("Maximum sector boundary must be greater than minimum boundary".to_string());
		}
		let mut d = Self::open(drive_num)?;
		d.sector_boundaries = boundaries;
		// seeks are now relative to boundaries.0
		if let Err(e) = d.seek(SeekFrom::Start(0)) {
			return Err(format!("{}", e));
		}
		Ok(d)
	}
}

impl Read for BufferedPhysicalDrive {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
		// respect sector boundaries when reading
		if (self.current_sector.position() == self.geometry.bytes_per_sector as u64) &&
			(self.current_sector_num < self.sector_boundaries.1) {
			let mut sector = vec![0; self.geometry.bytes_per_sector as usize];
			match self.drive.read(&mut sector) {
				Ok(_) => (),
				Err(e) => return Err(io::Error::new(
					io::ErrorKind::Other, e
				))
			}
			self.current_sector = Cursor::new(sector);
			self.current_sector_num += 1;
		}

		self.current_sector.read(buf)
    }
}

impl Seek for BufferedPhysicalDrive {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
		// moves to the specified sector, if it lies within the sector boundaries (if not, moves to last allowed sector
		// or returns an error)
		fn go_to_sector(drive: &mut BufferedPhysicalDrive, sector_num: u64) -> io::Result<()> {
			if sector_num < drive.sector_boundaries.0 {
				return Err(io::Error::new(ErrorKind::InvalidInput, "invalid seek to position before boundary"));
			}
			let sector_num = min(sector_num, drive.sector_boundaries.1);
			let mut sector = vec![0; drive.geometry.bytes_per_sector as usize];
			drive.drive.seek(SeekFrom::Start(sector_num * drive.geometry.bytes_per_sector as u64))?;
			drive.drive.read(&mut sector)?;
			drive.current_sector = Cursor::new(sector);
			drive.current_sector_num = sector_num;

			Ok(())
		}

        match pos {
            SeekFrom::Start(p) => {
				let sector = (p / self.geometry.bytes_per_sector as u64) + self.sector_boundaries.0;
				if sector != self.current_sector_num {
					go_to_sector(self, sector)?;
				}
				self.current_sector.seek(SeekFrom::Start(p % self.geometry.bytes_per_sector as u64))?;
			},
            SeekFrom::End(p) => {
				if p < 0 {
					let end = min(
						self.geometry.size(),
						(self.sector_boundaries.1 + 1) * self.geometry.bytes_per_sector as u64 - 1
					) as i64;
					let sector = (end + p) / self.geometry.bytes_per_sector as i64;
					if sector != self.current_sector_num as i64 {
						go_to_sector(self, sector as u64)?;
					}
					let target_pos = (end + p) - sector * self.geometry.bytes_per_sector as i64;
					self.current_sector.seek(SeekFrom::Start(target_pos as u64))?;
				}
			},
            SeekFrom::Current(p) => {
				let current = (self.current_sector_num * self.geometry.bytes_per_sector as u64 + self.current_sector.position()) as i64;
				let sector = (current + p) / self.geometry.bytes_per_sector as i64;
				if sector != self.current_sector_num as i64 {
					go_to_sector(self, sector as u64)?;
				}

				let target_pos = (current + p) - sector * self.geometry.bytes_per_sector as i64;
				self.current_sector.seek(SeekFrom::Start(target_pos as u64))?;
			}
		}

		Ok(self.current_sector_num * self.geometry.bytes_per_sector as u64 + self.current_sector.position())
    }
}

impl DiskGeometry {
	/// Returns the size of the disk in bytes.
	pub fn size(&self) -> u64 {
		self.cylinders as u64 * self.tracks_per_cylinder as u64 * self.sectors_per_track as u64
			* self.bytes_per_sector as u64
	}
}

impl From<DISK_GEOMETRY> for DiskGeometry {
    fn from(geo: DISK_GEOMETRY) -> Self {
        DiskGeometry {
			cylinders: unsafe { *geo.Cylinders.QuadPart() },
			media_type: geo.MediaType,
			tracks_per_cylinder: geo.TracksPerCylinder,
			sectors_per_track: geo.SectorsPerTrack,
			bytes_per_sector: geo.BytesPerSector
		}
    }
}