// SPDX-FileCopyrightText: Matteo Settenvini <matteo.settenvini@montecristosoftware.eu>
// SPDX-License-Identifier: EUPL-1.2

use super::Cleaner;
use crate::decision::{Action, Decision};
use anyhow::Result;
use async_trait::async_trait;
use goblin::elf::Elf;
use memmap2::Mmap;
use nix::{errno::Errno, libc::ino_t};
use petgraph::{prelude::DiGraphMap, visit::Dfs};
use std::{
    collections::{HashMap, HashSet},
    fs::File,
    io::{ErrorKind, Read, Seek},
    path::{Path, PathBuf},
};
use tokio::sync::{
    broadcast::{self, error::RecvError},
    mpsc,
};

type InodeMap = HashMap<ino_t, HashSet<PathBuf>>;
type InodeGraph = DiGraphMap<ino_t, ()>;

/// Cleans up unused shared libraries
/// and warns about broken dependencies as well
#[derive(Default)]
pub struct DsoCleaner {}

#[derive(Default)]
struct State {
    paths_map: InodeMap,
    graph: InodeGraph,
}

const INODE_ANY_EXECUTABLE: ino_t = 0;
const ELF_MAGIC_HEADER: &[u8; 4] = b"\x7fELF";

#[async_trait]
impl Cleaner for DsoCleaner {
    async fn run(
        &mut self,
        mut files: broadcast::Receiver<PathBuf>,
        decisions: mpsc::Sender<Decision>,
    ) -> Result<()> {
        let mut state = State::default();

        loop {
            match files.recv().await {
                Ok(file) => {
                    if let Err(e) = Self::process_file(&mut state, &file) {
                        log::warn!("{}: {}", file.display(), e);
                    }
                }
                Err(RecvError::Closed) => break,
                e => {
                    e?;
                }
            }
        }

        // DEBUGGING stuff you can uncomment for a quick and dirty graph:
        // println!(
        //     "{:?}",
        //     petgraph::dot::Dot::with_config(&state.graph, &[petgraph::dot::Config::EdgeNoLabel])
        // );

        let mut dfs = Dfs::empty(&state.graph);
        if state.graph.contains_node(INODE_ANY_EXECUTABLE) {
            dfs.move_to(INODE_ANY_EXECUTABLE);
        }
        while let Some(_) = dfs.next(&state.graph) {}

        for path in state
            .paths_map
            .into_iter()
            .filter_map(|(n, paths)| {
                if !dfs.discovered.contains(&n) {
                    Some(paths)
                } else {
                    None
                }
            })
            .flatten()
        {
            decisions
                .send(Decision {
                    path,
                    action: Action::Remove,
                })
                .await?;
        }

        Ok(())
    }
}

impl DsoCleaner {
    fn process_file(state: &mut State, path: &Path) -> Result<()> {
        let mut f = File::open(path)?;
        let mut hdr = [0u8; 4];
        if let Err(e) = f.read_exact(&mut hdr) {
            if e.kind() != ErrorKind::UnexpectedEof {
                anyhow::bail!("{}: {}", path.display(), e)
            }

            return Ok(()); // not ELF, ignore
        };

        let is_elf = &hdr == ELF_MAGIC_HEADER;
        if !is_elf {
            return Ok(());
        }

        f.rewind()?;
        let mmap = unsafe { Mmap::map(&f)? };
        let elf = Elf::parse(&mmap)?;

        if path.is_symlink() {
            if !elf.is_lib {
                // we don't care about symlinks to
                // executables in our graph, as we
                // are cleaning up only DSOs.
                Ok(())
            } else {
                Self::process_elf_symlink(state, path)
            }
        } else {
            Self::process_elf_file(state, path, &elf)
        }
    }

    fn process_elf_symlink(state: &mut State, path: &Path) -> Result<()> {
        let src = nix::sys::stat::lstat(path)?;
        let dst = nix::sys::stat::stat(path)?;
        if src.st_dev != dst.st_dev {
            log::warn!(
                "dso: {} points outside of the sysroot filesystem, check if this is intended",
                path.display()
            );
            return Ok(());
        }
        let current_dir = std::env::current_dir()?;
        let dst_path = std::fs::canonicalize(path)?
            .strip_prefix(current_dir)?
            .to_path_buf();

        log::trace!(
            "dso: adding to graph symlink: '{}' to '{}'",
            path.display(),
            dst_path.display()
        );

        Self::update_graph(state, path.into(), src.st_ino, dst_path, dst.st_ino);
        Ok(())
    }

    fn process_elf_file(state: &mut State, path: &Path, elf: &Elf) -> Result<()> {
        log::trace!("dso: adding to graph elf file '{}'", path.display());

        let search_paths = Self::determine_lib_search_paths(path, elf)?;

        let src_stat = nix::sys::stat::stat(path)?;
        let src_inode = if elf.is_lib {
            src_stat.st_ino
        } else {
            // We put all executables in the same node
            INODE_ANY_EXECUTABLE
        };

        'next_lib: for &library in elf.libraries.iter() {
            for lib_path in search_paths.iter() {
                let tentative_path = PathBuf::from(lib_path).strip_prefix("/")?.join(library);
                let dst = match nix::sys::stat::stat(&tentative_path) {
                    Ok(dst) => dst,
                    Err(Errno::ENOENT) => continue,
                    Err(e) => anyhow::bail!(
                        "got errno {} while accessing {}",
                        e,
                        tentative_path.display()
                    ),
                };

                if src_stat.st_dev != dst.st_dev {
                    continue; // These are not the droids you are looking for.
                }

                Self::update_graph(state, path.into(), src_inode, tentative_path, dst.st_ino);
                continue 'next_lib;
            }

            anyhow::bail!("{}: unable to find library {}", path.display(), library);
        }

        Ok(())
    }

    fn determine_lib_search_paths(path: &Path, elf: &Elf<'_>) -> Result<Vec<String>> {
        let mut search_paths = vec![];

        let current_dir = std::env::current_dir()?;
        let origin = std::fs::canonicalize(path)?
            .parent()
            .unwrap()
            .strip_prefix(current_dir)?
            .to_path_buf()
            .into_os_string()
            .into_string()
            .map_err(|s| anyhow::anyhow!("cannot represent {:?} as a UTF-8 string", s))?;

        if elf.rpaths != vec![""] {
            if elf.runpaths != vec![""] {
                let mut rpaths = elf
                    .rpaths
                    .iter()
                    .map(|p| p.replace("$ORIGIN", &origin))
                    .collect::<Vec<_>>();
                search_paths.append(&mut rpaths);
            }

            search_paths.append(&mut Self::get_env_library_paths());
        }

        if elf.runpaths != vec![""] {
            let mut runpaths = elf
                .runpaths
                .iter()
                .map(|p| p.replace("$ORIGIN", &origin))
                .collect::<Vec<_>>();
            search_paths.append(&mut runpaths);
        }

        search_paths.push("/usr/local/lib".into());
        search_paths.push("/lib".into());
        search_paths.push("/usr/lib".into());
        Ok(search_paths)
    }

    fn get_env_library_paths() -> Vec<String> {
        let ld_config_path = std::env::var("LD_LIBRARY_PATH");
        ld_config_path
            .as_ref()
            .map(|env| {
                env.split(':')
                    .filter(|s| s.is_empty())
                    .map(|s| s.into())
                    .collect()
            })
            .unwrap_or_default()
    }

    fn update_graph(
        state: &mut State,
        src_path: PathBuf,
        src_inode: ino_t,
        dst_path: PathBuf,
        dst_inode: ino_t,
    ) {
        state
            .paths_map
            .entry(src_inode)
            .or_default()
            .insert(src_path);

        state
            .paths_map
            .entry(dst_inode)
            .or_default()
            .insert(dst_path);

        state.graph.add_edge(src_inode, dst_inode, ());
    }
}