mod ints;
mod output;
mod structs;

use crate::{
    output::Output,
    structs::{
        BigEndian32, BigEndian64, ElfConfig, ElfHeader1, ElfHeader2, LittleEndian32,
        LittleEndian64, SectionHeaderEntry, SymbolTableEntry,
    },
};
use bytemuck::from_bytes;
use clap::Parser;
use core::str::from_utf8;
use eyre::{bail, eyre, Context};
use ints::Decode;
use serde_json::{json, Value};
use std::{fs, mem::size_of, path::PathBuf};
use symbolic::demangle::demangle;

#[derive(Parser)]
struct Opt {
    elf: PathBuf,
}

const HEADER1_LEN: usize = size_of::<ElfHeader1>();

macro_rules! parse_elf {
    ($elfcfg:ty, $out:expr, $header1:expr, $elf:expr) => {{
        let out = $out;
        const HEADER2_LEN: usize = size_of::<ElfHeader2<$elfcfg>>();
        let elf = $elf;
        if elf.len() < HEADER1_LEN + HEADER2_LEN {
            bail!("not a valid elf file: file too small");
        }

        let header2: &[u8; HEADER2_LEN] = elf[HEADER1_LEN..][..HEADER2_LEN].try_into()?;
        let header2: ElfHeader2<$elfcfg> = *from_bytes(header2);
        out.write_append("header", &header2)?;

        let e_shentsize = usize::from(header2.e_shentsize.to_native());
        if e_shentsize != size_of::<SectionHeaderEntry<$elfcfg>>() {
            bail!("wrong e_shentsize (0x{e_shentsize:x})");
        }

        let e_shnum = usize::from(header2.e_shnum.to_native());
        let section_header_table_size = e_shentsize * e_shnum;
        let section_header_table_start =
            usize::try_from(header2.e_shoff.to_native()).wrap_err("e_shoff bigger than usize")?;
        let section_header_table_end = section_header_table_start + section_header_table_size;

        if elf.len() < section_header_table_end {
            bail!("not a valid elf file: section table goes past eof");
        } else if elf.len() > section_header_table_end {
            bail!("not a valid elf file: section table doesn't end at eof");
        }

        let mut section_header_entries = vec![];
        for entry_start in
            (section_header_table_start..section_header_table_end).step_by(e_shentsize)
        {
            let entry = &elf[entry_start..][..e_shentsize];
            let entry: SectionHeaderEntry<$elfcfg> = *from_bytes(entry);
            section_header_entries.push(entry);
        }
        out.write("sections", &section_header_entries)?;

        let names_section_i = usize::from(header2.e_shstrndx.to_native());
        let names_section_entry = &section_header_entries
            .get(names_section_i)
            .ok_or(eyre!("invalid e_shstrndx"))?;

        let names_section_start: usize = names_section_entry.start()?;
        let names_section_end = names_section_entry.end()?;
        let names_section = &elf[names_section_start..names_section_end];

        let mut symtab_i = None;
        let mut strtab = None;

        let mut section_header_entries_with_names = vec![];
        for (i, entry) in section_header_entries.iter().enumerate() {
            let name_start: usize = (entry.sh_name.to_native())
                .try_into()
                .wrap_err("sh_name bigger than usize")?;

            if name_start > names_section.len() {
                bail!("sh_name out of bounds for .names section");
            }

            let name_end = (name_start..names_section.len())
                .find(|&i| names_section[i] == b'\0')
                .ok_or(eyre!(".names entry missing null byte"))?;
            let name = &names_section[name_start..name_end];
            let name = from_utf8(name).wrap_err(eyre!(".names entry not valid utf-8"))?;
            let section_start = entry.start().wrap_err_with(|| name.to_string())?;
            let section_end = entry.end().wrap_err_with(|| name.to_string())?;
            if section_end > elf.len() {
                bail!("section {name:?} end is past eof");
            }

            match name {
                ".symtab" => symtab_i = Some(i),
                ".strtab" => {
                    strtab = Some(&elf[section_start..section_end]);
                }
                _ => {}
            }

            let mut entry_json =
                serde_json::to_value(entry).wrap_err("failed to serialize section header")?;

            let Value::Object(fields) = &mut entry_json else {
                unreachable!()
            };

            fields.insert("name".to_string(), Value::String(name.to_string()));
            section_header_entries_with_names.push(entry_json);
        }
        let _ = out.write("sections", &section_header_entries_with_names);

        if let Some(symtab_i) = symtab_i {
            let mut symbols = vec![];

            let symtab = &section_header_entries[symtab_i];
            const ENTRY_SIZE: usize = size_of::<<$elfcfg as ElfConfig>::SymtabEntry>();
            if (symtab.sh_entsize.to_native() as usize) != ENTRY_SIZE {
                bail!(".symtab sh_entsize not equal to symtab entry size");
            }

            let section_start: usize = symtab.start().wrap_err(".symtab")?;
            let section_end: usize = symtab.end().wrap_err(".symtab")?;

            for entry in elf[section_start..section_end].chunks_exact(ENTRY_SIZE) {
                let entry: &[u8; ENTRY_SIZE] = entry.try_into().unwrap();
                let entry: &<$elfcfg as ElfConfig>::SymtabEntry = from_bytes(entry);
                let entry = SymbolTableEntry::from(*entry);
                let name_start: usize = (entry.st_name.try_into())
                    .wrap_err("symtab entry st_name too big to fit in usize")?;

                let name = (name_start != 0)
                    .then(|| {
                        strtab.map(|strtab| {
                            let name_end = (name_start..strtab.len())
                                .find(|&i| strtab[i] == b'\0')
                                .ok_or(eyre!(".strtab entry missing null byte"))?;
                            let name = &strtab[name_start..name_end];
                            from_utf8(name).wrap_err(eyre!(
                                ".strtab entry at 0x{name_start:x} is not valid utf-8"
                            ))
                        })
                    })
                    // Option<Option<Result>> to Result<Option>
                    .flatten()
                    .transpose()?
                    .map(demangle)
                    .unwrap_or_default();

                let section = section_header_entries_with_names
                    .get(usize::from(entry.st_shndx))
                    .and_then(|v| v.get("name"))
                    .and_then(|v| v.as_str());

                symbols.push(json!({
                    "value": entry.st_value,
                    "size": entry.st_size,
                    "section": section,
                    "symbol": name,
                }));
            }

            out.write("symbols", &symbols)?;
        }
    }};
}

fn main() -> eyre::Result<()> {
    let opt = Opt::parse();
    color_eyre::install()?;

    let elf = fs::read(opt.elf).wrap_err("failed to read elf file")?;

    if elf.len() < HEADER1_LEN {
        bail!("not a valid elf file: file too small");
    }
    let header: &[u8; HEADER1_LEN] = elf[..HEADER1_LEN].try_into()?;
    let header: ElfHeader1 = *from_bytes(header);

    let mut out = Output::default();
    out.write("header", &header)?;

    if &header.ei_magic != b"\x7FELF" {
        bail!("not a valid elf file: invalid magic");
    }

    if header.ei_version != 1 {
        bail!("unknown elf version: 0x{:x}", header.ei_version);
    }

    match (header.ei_class, header.ei_data) {
        (1, 1) => parse_elf!(LittleEndian32, &mut out, header1, elf),
        (2, 1) => parse_elf!(LittleEndian64, &mut out, header1, elf),
        (1, 2) => parse_elf!(BigEndian32, &mut out, header1, elf),
        (2, 2) => parse_elf!(BigEndian64, &mut out, header1, elf),
        (1 | 2, ei_data) => bail!("unknown e_ident[EI_DATA]: 0x{ei_data:x}"),
        (ei_class, ..) => bail!("unknown e_ident[EI_CLASS]: 0x{ei_class:x}"),
    }

    Ok(())
}