// Based on the gimli 0.16.1 dwarfdump.rs example
use anyhow::{anyhow, Error, Result};
use fallible_iterator::{convert, FallibleIterator};
use gimli::{AttributeValue, Endianity, Reader};
use object::Object;
use serde::ser::SerializeMap;
use serde::{Serialize, Serializer};
use std::borrow::{Borrow, Cow};
use std::fs;
use std::path::PathBuf;
use structopt::StructOpt;
use typed_arena::Arena;

fn list_file<E: Endianity>(file: &object::File, endian: E) -> Result<Vec<Unit>> {
    let arena = Arena::new();

    fn load_section<'a, 'file, 'input, S, E>(
        arena: &'a Arena<Cow<'file, [u8]>>,
        file: &'file object::File<'input>,
        endian: E,
    ) -> S
    where
        S: gimli::Section<gimli::EndianSlice<'a, E>>,
        E: gimli::Endianity,
        'file: 'input,
        'a: 'file,
    {
        let data = file
            .section_data_by_name(S::section_name())
            .unwrap_or(Cow::Borrowed(&[]));
        let data_ref = (*arena.alloc(data)).borrow();
        S::from(gimli::EndianSlice::new(data_ref, endian))
    }

    // Variables representing sections of the file. The type of each is inferred from its use in the
    // functions below.
    let debug_abbrev = &load_section(&arena, file, endian);
    let debug_info = &load_section(&arena, file, endian);
    let debug_str = &load_section(&arena, file, endian);

    list_info(debug_info, debug_abbrev, debug_str)
}

fn list_info<R: Reader>(
    debug_info: &gimli::DebugInfo<R>,
    debug_abbrev: &gimli::DebugAbbrev<R>,
    debug_str: &gimli::DebugStr<R>,
) -> Result<Vec<Unit>> {
    let mut v = Vec::new();
    let units = debug_info.units().collect::<Vec<_>>().unwrap();
    for u in units {
        let abbrevs = u.abbreviations(debug_abbrev)?;
        v.append(&mut list_entries(u.entries(&abbrevs), debug_str)?);
    }
    Ok(v)
}

#[derive(Debug, Serialize)]
struct Unit {
    comp_dir: String,
    comp_name: String,
}

fn list_entries<R: Reader>(
    mut entries: gimli::EntriesCursor<R>,
    debug_str: &gimli::DebugStr<R>,
) -> Result<Vec<Unit>> {
    let mut depth = 0;
    let mut v = Vec::new();
    while let Some((delta_depth, entry)) = entries.next_dfs()? {
        depth += delta_depth;
        assert!(depth >= 0);
        if depth > 0 {
            continue;
        }

        if entry.tag() == gimli::DW_TAG_compile_unit || entry.tag() == gimli::DW_TAG_type_unit {
            let comp_dir = entry
                .attr(gimli::DW_AT_comp_dir)?
                .and_then(|attr| attr.string_value(debug_str))
                .map::<Result<String>, _>(|s| Ok(s.to_string()?.into_owned()))
                .transpose()?
                .ok_or_else(|| anyhow!("Missing DW_AT_comp_dir"))?;

            let at_name = if let Some(it) = entry.attr(gimli::DW_AT_name)? {
                it
            } else {
                eprintln!("Warning: unit without name, skipping it");
                continue;
            };

            if let Some(r) = at_name.string_value(debug_str) {
                let comp_name = r.to_string()?;
                if comp_name == "<artificial>" {
                    eprintln!("Warning: Artificial name in compile unit, probably DWARF debug information has been generated with LTO")
                } else {
                    v.push(Unit {
                        comp_dir,
                        comp_name: comp_name.into_owned(),
                    });
                }
            } else {
                match at_name.raw_value() {
                    AttributeValue::DebugStrRefSup(_) => {
                        eprintln!("Warning: compilation unit name in supplemental file")
                    }
                    _ => eprintln!("Warning: compilation unit has unexpected name type"),
                }
            }
        }
    }
    Ok(v)
}

#[derive(Debug, Serialize)]
struct Data {
    units: Vec<Unit>,
}

#[derive(StructOpt, Debug)]
struct Opt {
    #[structopt(short, long, parse(from_os_str))]
    /// File to write json output to
    output: Option<PathBuf>,
    #[structopt(long = "strip-prefix", short, long)]
    /// Strip given prefix from comp_dir paths
    strip_prefix: Option<String>,
    file: Vec<String>,
}

fn process_file<E: Endianity>(
    file: &object::File,
    endian: E,
    strip_prefix: Option<&String>,
) -> Result<Vec<Unit>> {
    let units = list_file(&file, endian)?;
    let units = if let Some(strip) = strip_prefix {
        units
            .iter()
            .map(|u| {
                let comp_dir = if u.comp_dir.starts_with(strip) {
                    (&u.comp_dir[strip.len()..]).to_string()
                } else {
                    u.comp_dir.clone()
                };
                Unit {
                    comp_dir,
                    comp_name: u.comp_name.clone(),
                }
            })
            .collect()
    } else {
        units
    };

    Ok(units)
}

fn serialize<S, K, V, I>(s: S, mut iter: I, len: Option<usize>) -> Result<()>
where
    S: Serializer,
    S::Error: 'static + Send + Sync,
    K: Serialize,
    V: Serialize,
    I: FallibleIterator<Item = (K, V), Error = Error>,
{
    let mut map = s.serialize_map(len)?;
    while let Some((k, v)) = iter.next()? {
        map.serialize_entry(&k, &v)?;
    }
    map.end()?;
    Ok(())
}

fn main() -> Result<()> {
    let opt = Opt::from_args();

    let data = convert(opt.file.iter().map(|f| -> Result<_> {
        let file = fs::File::open(&f)?;
        let file = unsafe { memmap::Mmap::map(&file) }?;
        let file = object::File::parse(&*file).map_err(|e| anyhow!(e))?;

        let endian = if file.is_little_endian() {
            gimli::RunTimeEndian::Little
        } else {
            gimli::RunTimeEndian::Big
        };

        let units = process_file(&file, endian, opt.strip_prefix.as_ref())?;
        Ok((f.clone(), Data { units }))
    }));

    if let Some(output) = &opt.output {
        let file = fs::OpenOptions::new()
            .write(true)
            .create_new(true)
            .open(output)?;
        let mut s = serde_json::Serializer::new(&file);
        if let Err(e) = serialize(&mut s, data, Some(opt.file.len())) {
            fs::remove_file(output)?;
            return Err(e);
        }
    } else {
        let mut s = serde_json::Serializer::pretty(std::io::stdout());
        serialize(&mut s, data, Some(opt.file.len()))?;
    }

    Ok(())
}
