//! Uses DWARF debug info, if available, to find file and line number information for a particular
//! offset in an input section.

use crate::elf::Elf;
use crate::elf::ElfClass;
use crate::elf::ElfCrel;
use crate::elf::ElfRela;
use crate::elf::File;
use crate::elf::SectionHeader;
use crate::error::Result;
use crate::fs::path_from_bytes;
use crate::platform::Arch;
use crate::platform::ObjectFile as _;
use crate::platform::Relocation;
use crate::platform::SourceInfo;
use crate::platform::SourceInfoDetails;
use anyhow::Context;
use linker_utils::elf::RelocationKind;
use object::LittleEndian;
use object::read::elf::RelocationSections;
use object::read::elf::SectionHeader as _;
use object::read::elf::Sym as _;
use std::borrow::Cow;
use std::fmt::Display;
use std::path::PathBuf;

/// The address at which we'll pretend that we loaded the section we're interested in. This value is
/// arbitrary, but should be larger than the largest input section we expect to encounter and small
/// enough to fit comfortably in a u32.
const SECTION_LOAD_ADDRESS: u64 = 0x8000_0000;

/// Attempts to locate source info for `offset_in_section` within `section`.
pub(crate) fn get_source_info<C: ElfClass, A: Arch<Platform = Elf<C>>>(
    object: &File<C>,
    relocations: &RelocationSections,
    section: &SectionHeader<C>,
    offset_in_section: u64,
) -> Result<SourceInfo> {
    let dwarf_sections =
        gimli::DwarfSections::load(&|id: gimli::SectionId| -> Result<Cow<[u8]>> {
            section_data_with_relocations::<C, A>(object, relocations, id, section)
        })?;

    let borrow_section: &dyn for<'a> Fn(
        &'a Cow<[u8]>,
    ) -> gimli::EndianSlice<'a, gimli::LittleEndian> =
        &|section| gimli::EndianSlice::new(section, gimli::LittleEndian);

    let dwarf = dwarf_sections.borrow(borrow_section);

    let mut details = None;

    let address_of_interest = SECTION_LOAD_ADDRESS + offset_in_section;

    let mut iter = dwarf.units();
    while let Some(header) = iter.next()? {
        let unit = dwarf.unit(header)?;

        let Some(program) = unit.line_program.clone() else {
            continue;
        };

        let comp_dir = unit
            .comp_dir
            .as_ref()
            .map(|dir| path_from_bytes(dir))
            .unwrap_or_default();

        let mut rows = program.rows();

        while let Some((header, row)) = rows.next_row()? {
            if row.address() > address_of_interest {
                break;
            }

            // Computing the path for every row seems a bit wasteful. If it turns out that this is
            // actually a problem, then we could experiment with iterating through the rows twice.
            // Once to determine which row we want then a second time to get the relevant properties
            // from the row.
            let mut path = PathBuf::new();
            if let Some(file) = row.file(header) {
                path = comp_dir.clone();

                path.push(path_from_bytes(
                    &dwarf.attr_string(&unit, file.path_name())?,
                ));
            }

            let line = row.line().map_or(0, |l| l.get());

            details = Some(SourceInfoDetails { path, line });
        }
    }

    Ok(SourceInfo(details))
}

/// Gets the data for section `id` from `object` and applies relocations to it.
fn section_data_with_relocations<C: ElfClass, A: Arch<Platform = Elf<C>>>(
    object: &File<C>,
    relocations: &RelocationSections,
    id: gimli::SectionId,
    section_of_interest: &SectionHeader<C>,
) -> Result<Cow<'static, [u8]>> {
    let data = match object.section_by_name(id.name()) {
        Some((index, section)) => {
            let mut section_data = object.section_data_cow(section)?.into_owned();

            // Apply relocations.
            match object.relocations(index, relocations)? {
                crate::elf::RelocationList::Rela(relocations) => {
                    apply_section_relocations::<C, A, ElfRela<C>>(
                        object,
                        section_of_interest,
                        &mut section_data,
                        relocations.iter().copied().map(ElfRela::new),
                    )?;
                }
                crate::elf::RelocationList::Crel(relocations) => {
                    apply_section_relocations::<C, A, ElfCrel<C>>(
                        object,
                        section_of_interest,
                        &mut section_data,
                        relocations.filter_map(|r| r.ok().map(ElfCrel::new)),
                    )?;
                }
            }

            Cow::Owned(section_data)
        }
        None => Cow::Borrowed(&[][..]),
    };

    Ok(data)
}

fn apply_section_relocations<
    C: ElfClass,
    A: Arch<Platform = Elf<C>>,
    R: Relocation<Platform = Elf<C>>,
>(
    object: &File<C>,
    section_of_interest: &SectionHeader<C>,
    section_data: &mut [u8],
    relocations: impl Iterator<Item = R>,
) -> Result {
    for rel in relocations {
        let sym_index = rel.symbol().context("Relocation has no symbol")?;
        let symbol = object.symbol(sym_index)?;

        let mut value = symbol
            .st_value(LittleEndian)
            .into()
            .wrapping_add(rel.addend() as u64);

        let Some(section_index) = object.symbol_section(symbol, sym_index)? else {
            // Ignore undefined symbols.
            continue;
        };
        let symbol_section = object.section(section_index)?;

        let data_offset = rel.offset() as usize;

        let r_type = A::relocation_from_raw(rel.raw_type())?;

        let linker_utils::elf::RelocationSize::ByteSize(num_bytes) = r_type.size else {
            continue;
        };

        let rel_out = section_data
            .get_mut(data_offset..data_offset + num_bytes)
            .context("Invalid relocation offset")?;

        match r_type.kind {
            RelocationKind::Absolute => {
                let in_section_of_interest = symbol_section.sh_offset(LittleEndian).into()
                    == section_of_interest.sh_offset(LittleEndian).into();

                if in_section_of_interest {
                    value += SECTION_LOAD_ADDRESS;
                }

                r_type.write_to_buffer(value, rel_out)?;
            }
            RelocationKind::AbsoluteAddition | RelocationKind::AbsoluteSubtraction => {
                r_type.write_to_buffer(value, rel_out)?;
            }
            _ => {}
        }
    }
    Ok(())
}

impl Display for SourceInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(details) = self.0.as_ref() {
            let SourceInfoDetails { path, line } = details;
            write!(f, "\n    {}:{}", path.display(), line)?;
        }
        Ok(())
    }
}
