use crate::backend::{LibImport, pe::data_dir::DataDir};

use super::ByteEncoder;

pub fn encode(data: &mut ByteEncoder, imports: &[LibImport], code_start: usize) -> DataDir {
    data.align(4);
    let start = data.pos() as u32;
    let idt = data.reserve_arr::<ImportDirTable>(imports.len());
    // null entry to mark end
    data.pad(size_of::<ImportDirTable>());
    let end = data.pos() as u32;

    for (i, import) in imports.iter().enumerate() {
        // name
        let name_rva = data.pos() as u32;
        data.extend(import.name.as_bytes());
        data.push(0);

        // lookup table
        data.align(size_of::<ImportLookupEntry>());
        let lookup_start = data.pos();
        let lookup = data.reserve_arr::<ImportLookupEntry>(import.syms.len());
        data.pad(size_of::<ImportLookupEntry>());

        for (i, sym) in import.syms.iter().enumerate() {
            let rva = hint_name_entry(data, 0, &sym.name);
            data[lookup][i] = ImportLookupEntry::name(rva);
        }

        // address table
        data.align(size_of::<ImportLookupEntry>());
        let addr_start = data.pos();
        for (i, sym) in import.syms.iter().enumerate() {
            let here = data.pos() as i32;
            for &usage in &sym.usages {
                // NOTE: sets relative offet rn
                let code_pos = code_start + usage;
                data.set_at::<i32>(code_pos, here - code_pos as i32 - 4);
            }
            let entry = data[lookup][i];
            data.val(&entry);
        }
        data.pad(size_of::<ImportLookupEntry>());

        // entry
        data[idt][i] = ImportDirTable {
            lookup_table_rva: lookup_start as u32,
            time_date_stamp: 0,
            forwarder_chain: 0,
            name_rva,
            address_table_rva: addr_start as u32,
        };
    }
    DataDir {
        virt_addr_rva: start,
        size: end - start,
    }
}

#[repr(C)]
pub struct ImportDirTable {
    pub lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub address_table_rva: u32,
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct ImportLookupEntry(u64);

impl ImportLookupEntry {
    pub const NULL: Self = Self(0);
    pub fn name(hint_name_table_rva: u32) -> Self {
        assert!(hint_name_table_rva >> 30 == 0);
        Self(hint_name_table_rva as u64)
    }
    pub fn ordinal(ordinal: u16) -> Self {
        Self(ordinal as u64 | (1 << 63))
    }
    pub fn bytes(&self) -> [u8; 8] {
        self.0.to_le_bytes()
    }
}

pub fn hint_name_entry(data: &mut ByteEncoder, hint: u16, name: &str) -> u32 {
    let pos = data.pos() as u32;
    data.extend(hint.to_le_bytes());
    data.extend(name.as_bytes());
    data.push(0);
    data.align(2);
    pos
}