use std::{fs::File, io::BufWriter, io::Write, mem::size_of, path::PathBuf};

use elf_utilities::{
    header::Ehdr64,
    section::{build_string_table, Shdr64, Type as ShType},
    segment::{Phdr64, Type as SeType, PF_R, PF_W, PF_X},
};

use crate::{
    common::{expand_path, pad_to_next_page, Loadable, Output, SegmentType},
    error::Error,
};

use super::ElfObject;

const SECTION_NAMES: [&str; 5] = [".text", ".rodata", ".data", ".bss", ".shstrtab"];

pub struct ElfOutput {
    destination: PathBuf,
    writer: BufWriter<File>,
}

impl ElfOutput {
    pub fn new(destination: PathBuf) -> Result<Self, Error> {
        let str_path = expand_path(&destination)?;
        let writer = make_file_writer(str_path, 0o755)?;

        let result = Self {
            destination,
            writer,
        };

        Ok(result)
    }

    fn write_shstrtab(
        &mut self,
        shstrtab: (Vec<u8>, Vec<usize>),
        mut offset: usize,
        name_idx: usize,
    ) -> Result<usize, Error> {
        const SHS: u16 = size_of::<Shdr64>() as u16;

        // write .shstrtab header
        let data_offset = (offset + usize::from(SHS)) as u64;
        let strtab_header = make_section_header(
            shstrtab.1[name_idx],
            ShType::StrTab,
            data_offset,
            shstrtab.0.len() as u64,
        );
        offset += self.writer.write(&strtab_header.to_le_bytes())?;
        // write .shstrtab data
        offset += self.writer.write(&shstrtab.0)?;

        Ok(offset)
    }
}

impl Output<ElfObject> for ElfOutput {
    fn finalize(mut self, objects: &[ElfObject], loadable: &Loadable) -> Result<PathBuf, Error> {
        const EHS: u64 = size_of::<Ehdr64>() as u64;
        const PHS: u16 = size_of::<Phdr64>() as u16;
        const SHS: u16 = size_of::<Shdr64>() as u16;

        let shstrtab = make_shstrtab();
        let mut ehdr = make_elf_header();
        ehdr.e_shnum = SECTION_NAMES.len() as u16;
        ehdr.e_phnum = 3u16; // .text, .rodata + .data, .bss
        ehdr.e_shstrndx = ehdr.e_shnum - 1; // .shstrab is always last
        ehdr.e_entry = loadable.start_offset()?;
        ehdr.e_phentsize = PHS;
        ehdr.e_shentsize = SHS;
        ehdr.e_phoff = EHS;
        ehdr.e_shoff =
            ehdr.e_phoff + (usize::from(ehdr.e_phentsize) * usize::from(ehdr.e_phnum)) as u64;

        let mut offset = 0;

        // write ELF header
        offset += self.writer.write(&ehdr.to_le_bytes())?;

        // write program header table
        // contains .text + .rodata as one segment header
        let ph = make_program_header(offset, loadable.program_size(), SegmentType::Text);
        offset += self.writer.write(&ph.to_le_bytes())?;
        // contains .data as one segment header
        let ph = make_program_header(offset, loadable.data_size(), SegmentType::Data);
        offset += self.writer.write(&ph.to_le_bytes())?;
        // contains .bss as one segment header
        let ph = make_program_header(offset, loadable.bss_size(), SegmentType::Bss);
        offset += self.writer.write(&ph.to_le_bytes())?;

        eprintln!("SH start: {}", offset);
        let mut name_idx = 0;
        // write section header table (text + rodata + data + bss)
        for sections in loadable.segment_sections() {
            let sh = make_section_header(
                shstrtab.1[name_idx],
                ShType::ProgBits,
                0,
                sections.data_size(),
            );
            name_idx += 1;
            offset += self.writer.write(&sh.to_le_bytes())?;
        }

        // .shstrtab as last section header written (+ the name data including)
        offset = self.write_shstrtab(shstrtab, offset, name_idx)?;

        // program data (segments)
        offset += pad_to_next_page(&mut self.writer, offset)?;
        eprintln!("Prog start: {}", offset);
        // write section/segment data
        for bytes in loadable.program_bytes(objects) {
            offset += self.writer.write(bytes?)?;
        }
        offset += pad_to_next_page(&mut self.writer, offset)?;

        eprintln!("Data start: {}", offset);
        for bytes in loadable.data_bytes(objects) {
            offset += self.writer.write(bytes?)?;
        }

        self.writer.flush()?;

        Ok(self.destination)
    }
}

fn make_file_writer(output_filename: &str, permission: u32) -> Result<BufWriter<File>, Error> {
    use std::os::unix::fs::OpenOptionsExt;

    let file = std::fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .mode(permission)
        .open(output_filename)?;

    Ok(BufWriter::new(file))
}

fn make_elf_header() -> Ehdr64 {
    use elf_utilities::header::{Class, Data, Machine, Type, Version, OSABI};

    let mut ehdr = Ehdr64::default();
    ehdr.set_elf_type(Type::Exec);
    ehdr.set_class(Class::Bit64);
    ehdr.set_machine(Machine::X8664);
    ehdr.set_object_version(Version::Current);
    ehdr.set_file_version(Version::Current);
    ehdr.set_osabi(OSABI::Linux);
    ehdr.set_data(Data::LSB2);
    ehdr.e_ehsize = size_of::<Ehdr64>() as u16;

    ehdr
}

fn make_section_header(name_idx: usize, stype: ShType, offset: u64, size: u64) -> Shdr64 {
    let mut h = Shdr64::default();
    h.set_type(stype);

    // section index used for name and section indexing (currently same)
    let idx = 1 + (name_idx as u32); // add 0 at start
                                     // name index
    h.sh_name = idx;
    // link index
    h.sh_link = 0;
    h.sh_size = size;
    // filled at finalize()
    h.sh_offset = offset;

    h
}

fn make_program_header(offset: usize, size: u64, which: SegmentType) -> Phdr64 {
    let mut header = Phdr64::default();

    header.set_type(SeType::Load);
    header.p_filesz = size;
    header.p_memsz = size;
    header.p_offset = offset as u64;

    match which {
        SegmentType::Text => header.p_flags = PF_R | PF_X,
        SegmentType::Data => header.p_flags = PF_R | PF_W,
        SegmentType::Bss => {
            header.p_flags = PF_R;
            header.p_filesz = 0;
        }
    }

    header
}

// strtab as bytes + indexes to individual strings
fn make_shstrtab() -> (Vec<u8>, Vec<usize>) {
    let strtab_bytes = build_string_table(Vec::from(SECTION_NAMES), true);

    let mut indexes = Vec::new();
    let mut on_string = false;
    for (i, byte) in strtab_bytes.iter().enumerate() {
        if *byte == 0 {
            on_string = false;
            continue;
        }

        if on_string {
            continue;
        }

        assert!(i > 0, "First byte of .shstrab not 0");
        indexes.push(i - 1);
        on_string = true;
    }

    (strtab_bytes, indexes)
}