use core::ops::{Add, Range, Sub};
use sub_core::cast_traits::Upcast;
use crate::structs::*;

#[repr(C)]
pub struct ImageBase(());

extern "C" {
    #[link_name = "__ImageBase"]
    static ImageBasePtr: [u8;0];
}

impl ImageBase {

    #[inline(always)]
    pub fn get() -> &'static ImageBase {
        unsafe { &*(ImageBasePtr.as_ptr() as *const ImageBase) }
    }


    #[inline(always)]
    pub fn as_ptr(&self) -> usize {
        self as *const _ as usize
    }

    #[inline(always)]
    pub(crate) fn offset<T: Upcast<usize>>(&self, offset: T) -> *const u8 {
        unsafe { (self as *const _ as *const u8).add(offset.upcast()) }
    }

    pub unsafe fn dos(&self) -> &ImageDOSHeader {
        &*(self as *const _ as *const ImageDOSHeader)
    }

    pub unsafe fn dos_mut(&mut self) -> &mut ImageDOSHeader {
        &mut *(self as *mut _ as *mut ImageDOSHeader)
    }

    pub unsafe fn nt_header(&self) -> &'static ImageNTHeaders64 {
        &*(self.offset(self.dos().e_lfanew) as *const ImageNTHeaders64)
    }

    pub unsafe fn nt_header_mut(&mut self) -> &'static mut ImageNTHeaders64 {
        &mut *(self.offset(self.dos().e_lfanew) as *mut ImageNTHeaders64)
    }

    pub unsafe fn sections(&self) -> &[ImageSectionHeader] {
        self.nt_header().sections()
    }

    pub unsafe fn as_range(&self) -> Range<usize> {
        self.as_ptr()..self.as_ptr() + self.nt_header().optional_header.size_of_image as usize
    }

    pub unsafe fn as_slice(&self) -> &[u8] {
        let ptr = self.as_ptr() as *const u8;
        let size = self.nt_header().optional_header.size_of_image as usize;
        core::slice::from_raw_parts(ptr, size)
    }

    pub unsafe fn as_slice_mut(&self) -> &mut [u8] {
        let ptr = self.as_ptr() as *mut u8;
        let size = self.nt_header().optional_header.size_of_image as usize;
        core::slice::from_raw_parts_mut(ptr, size)
    }

}

impl Add<usize> for &ImageBase {
    type Output = usize;
    fn add(self, rhs: usize) -> Self::Output {
        (self as *const _ as usize) + rhs
    }
}

impl Sub<usize> for &ImageBase {
    type Output = usize;
    fn sub(self, rhs: usize) -> Self::Output {
        (self as *const _ as usize) + rhs
    }
}