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

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

#[derive(Debug, Default, Serialize)]
struct ObjectInfo {
    units: Vec<Unit>,
    // Note that we use a BTreeSet so that the output order is consistent across
    // runs.
    external_files: BTreeSet<String>,
}

fn list_file<E: Endianity>(file: &object::File, endian: E) -> Result<ObjectInfo> {
    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_by_name(S::section_name())
            .and_then(|section| section.uncompressed_data().ok())
            .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);
    let debug_line_str = &load_section(&arena, file, endian);
    let debug_line = &load_section(&arena, file, endian);

    list_info(
        debug_info,
        debug_abbrev,
        debug_str,
        debug_line_str,
        debug_line,
    )
}

fn list_info<R: Reader>(
    debug_info: &gimli::DebugInfo<R>,
    debug_abbrev: &gimli::DebugAbbrev<R>,
    debug_str: &gimli::DebugStr<R>,
    debug_line: &gimli::DebugLine<R>,
    debug_line_str: &gimli::DebugLineStr<R>,
) -> Result<ObjectInfo> {
    let mut info = ObjectInfo::default();
    let dw_units = debug_info.units().collect::<Vec<_>>().unwrap();
    for u in dw_units {
        let abbrevs = u.abbreviations(debug_abbrev)?;
        info.units.append(&mut list_entries(
            &mut info.external_files,
            u.entries(&abbrevs),
            u.address_size(),
            debug_str,
            debug_line_str,
            debug_line,
        )?);
    }
    Ok(info)
}

fn get_attr_string<R: Reader>(
    dw_attr_value: &gimli::AttributeValue<R>,
    debug_str: &gimli::DebugStr<R>,
    debug_line_str: &gimli::DebugLineStr<R>,
) -> Result<(R, String)> {
    let reader = if let Some(r) = dw_attr_value.string_value(debug_str) {
        r
    } else {
        match dw_attr_value {
            AttributeValue::DebugLineStrRef(offset) => debug_line_str
                .get_str(*offset)
                .expect("attribute value should be set"),
            AttributeValue::DebugStrRefSup(_) => {
                bail!("attribute in supplemental file");
            }
            _ => bail!("attribute with unexpected type"),
        }
    };

    let string = reader.to_string()?.into_owned();
    Ok((reader, string))
}

fn entry_attr_value<R: Reader>(
    entry: &gimli::DebuggingInformationEntry<R>,
    attr_name: gimli::constants::DwAt,
    debug_str: &gimli::DebugStr<R>,
    debug_line_str: &gimli::DebugLineStr<R>,
) -> Result<Option<(R, String)>> {
    let dw_attr = if let Some(it) = entry.attr(attr_name)? {
        it
    } else {
        return Ok(None);
    };

    get_attr_string(&dw_attr.value(), debug_str, debug_line_str).map(Some)
}

fn collect_external_files<R: Reader>(
    external_files: &mut BTreeSet<String>,
    entry: &gimli::DebuggingInformationEntry<R>,
    address_size: u8,
    debug_str: &gimli::DebugStr<R>,
    debug_line_str: &gimli::DebugLineStr<R>,
    debug_line: &gimli::DebugLine<R>,
    comp_dir_reader: R,
    comp_name_reader: R,
) -> Result<()> {
    let program = if let Some(gimli::AttributeValue::DebugLineRef(offs)) =
        entry.attr_value(gimli::DW_AT_stmt_list)?
    {
        debug_line.program(
            offs,
            address_size,
            Some(comp_dir_reader),
            Some(comp_name_reader),
        )?
    } else {
        return Ok(());
    };

    let header = program.header();
    for file_entry in header.file_names() {
        if file_entry.directory_index() == 0 {
            // Directory index of 0 means it's from the current CU's directory,
            // so skip it.
            continue;
        }

        let mut name =
            PathBuf::from(get_attr_string(&file_entry.path_name(), debug_str, debug_line_str)?.1);

        if let Some(dir_attr) = file_entry.directory(header) {
            let (_, dir) = get_attr_string(&dir_attr, debug_str, debug_line_str)?;
            name = PathBuf::from(dir).join(name);
        }

        if name.is_absolute() {
            // This shouldn't be invalid UTF-8, because it just came from some
            // String instances already.
            external_files.insert(name.into_os_string().into_string().unwrap());
        }
    }

    Ok(())
}

fn list_entries<R: Reader>(
    external_files: &mut BTreeSet<String>,
    mut entries: gimli::EntriesCursor<R>,
    address_size: u8,
    debug_str: &gimli::DebugStr<R>,
    debug_line_str: &gimli::DebugLineStr<R>,
    debug_line: &gimli::DebugLine<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_reader, comp_dir) =
                entry_attr_value(entry, gimli::DW_AT_comp_dir, debug_str, debug_line_str)?
                    .ok_or_else(|| anyhow!("Unable to parse or missing DW_AT_comp_dir"))?;

            let at_name = entry_attr_value(entry, gimli::DW_AT_name, debug_str, debug_line_str);
            let (comp_name_reader, comp_name) = match at_name {
                Ok(Some(result)) => result,
                Ok(None) => {
                    eprintln!("Warning: unit without name, skipping it");
                    continue;
                }
                Err(err) => {
                    eprintln!("Warning: {}", err);
                    continue;
                }
            };

            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,
                });
            }

            if let Err(err) = collect_external_files(
                external_files,
                entry,
                address_size,
                debug_str,
                debug_line_str,
                debug_line,
                comp_dir_reader,
                comp_name_reader,
            ) {
                eprintln!("Warning: collecting external files: {}", err);
            }
        }
    }
    Ok(v)
}

#[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<ObjectInfo> {
    let mut info = list_file(file, endian)?;
    if let Some(strip) = strip_prefix {
        for u in info.units.iter_mut() {
            if u.comp_dir.starts_with(strip) {
                u.comp_dir = (&u.comp_dir[strip.len()..]).to_string();
            }
        }
    }

    Ok(info)
}

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 info = process_file(&file, endian, opt.strip_prefix.as_ref())?;
        Ok((f.clone(), info))
    }));

    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(())
}
