use crate::{find_files::*, Client, Downloaded, Error};
use std::collections::{BTreeSet, HashMap};
use std::ffi::CString;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
use tracing::*;

use crate::deb;

/// The result of downloading/extracting a package.
#[derive(Debug, Clone)]
pub struct Packages {
    /// Paths of each of the extracted packages, in reverse order from
    /// the root (i.e. the root is the last element).
    pub result: Vec<Arc<PathBuf>>,
    /// Paths to directories (i.e. */usr/bin, */usr/lib etc), useful
    /// to create $PATH or $CFLAGS from the client.
    pub paths: Vec<Arc<String>>,
}

/// Download a Debian package and its dependency DAG from package
/// indices and the package's name.
pub async fn download_extract_deps(
    index: &[deb::Index],
    client: &Client,
    package: &str,
) -> Result<Packages, Error> {
    let files = Arc::new(Mutex::new(HashMap::new()));
    let mut seen = HashMap::new();
    let pkg = multi_lookup(index, &package).await.unwrap();
    let mut stack = vec![StackElt {
        package: pkg,
        rx: None,
        deps: Vec::new(),
        transitive_deps: Vec::new(),
        transitive_deps_h: BTreeSet::new(),
        ld_path: Vec::new(),
        ld_path_h: BTreeSet::new(),
        files: BTreeSet::new(),
    }];
    let mut result = Vec::new();
    let mut paths = BTreeSet::new();
    while let Some(mut elt) = stack.pop() {
        info!("{:?}", elt);

        if let Some(rx) = elt.rx.take() {
            let downloaded = rx.await.unwrap()?;
            info!("downloaded {:?}", elt.package);
            let final_path = elt
                .finalize(client, &files, downloaded, &mut stack, &mut result)
                .await?;
            seen.insert(
                elt.package.package,
                (elt.ld_path, elt.transitive_deps, final_path),
            );
            paths = elt.files;
        } else {
            if let Some((ld_path, transitive_deps, final_path)) = seen.get(&elt.package.package) {
                debug!("already seen {:?}", elt.package.package);
                if let Some(last) = stack.iter_mut().rev().filter(|x| x.rx.is_some()).next() {
                    for ld in ld_path {
                        if last.ld_path_h.insert(ld.clone()) {
                            last.ld_path.push(ld.clone());
                        }
                    }
                    for dep in transitive_deps {
                        if last.transitive_deps_h.insert(dep.clone()) {
                            debug!(
                                "adding transitive dep: {:?} -> {:?}",
                                last.package.package, dep
                            );
                            last.transitive_deps.push(dep.clone());
                        }
                    }
                    last.deps.push(final_path.clone());
                }
                continue;
            } else {
                seen.insert(
                    elt.package.package,
                    (Vec::new(), Vec::new(), Arc::new(PathBuf::new())),
                );
            }

            elt.spawn_extract(client.clone()).await;
            elt.push_deps(index, &mut stack).await;
        }
    }
    Ok(Packages {
        result,
        paths: paths.into_iter().collect(),
    })
}

async fn add_subst(downloaded: PathBuf, target: &Path, files: Files) -> Result<(), std::io::Error> {
    let mut files = files.lock().unwrap();
    for (f, m) in find_files(downloaded.clone())? {
        if !m.is_dir() {
            if let Ok(p) = f.strip_prefix(&downloaded) {
                use std::collections::hash_map::Entry;
                if let Entry::Vacant(e) = files.entry(p.to_path_buf()) {
                    debug!("add_subst {:?} -> {:?}", p, target);
                    e.insert(target.to_path_buf());
                }
            }
        }
    }
    Ok(())
}

async fn hash_reader(
    mut reader: impl tokio::io::AsyncReadExt + Unpin,
    hasher: &mut blake3::Hasher,
) -> std::io::Result<u64> {
    let mut buffer = [0; 65536];
    let mut total = 0;
    loop {
        match reader.read(&mut buffer).await {
            Ok(0) => return Ok(total),
            Ok(n) => {
                hasher.update(&buffer[..n]);
                total += n as u64;
            }
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        }
    }
}

async fn extract_task(client: &Client, download: &mut Downloaded) -> Result<(), Error> {
    debug!("extracting {:?}", download);
    let path = download.path.with_extension("");
    let lock = client.lock_store_path(&path).await;
    let mut f = std::io::BufReader::new(std::fs::File::open(&download.path)?);
    let d = match deb::Deb::read(&mut f) {
        Ok(d) => d,
        Err(e) => {
            std::fs::remove_file(&download.path).unwrap_or(());
            return Err(e.into());
        }
    };
    std::fs::create_dir_all(&path).unwrap_or(());
    match d.decompress(&mut f, &path) {
        Ok(()) => {
            download.path = path;
            drop(lock);
            Ok(())
        }
        Err(e) => {
            std::fs::remove_dir_all(&path).unwrap_or(());
            std::fs::remove_file(&download.path).unwrap_or(());
            drop(lock);
            Err(e.into())
        }
    }
}

async fn multi_lookup<'a>(index: &'a [deb::Index], package: &str) -> Option<deb::Stanza<'a>> {
    for ind in index {
        if let Some(i) = ind.lookup_async(package).await {
            return Some(i);
        }
    }
    for ind in index {
        if let Some(i) = ind.lookup_virtual_async(package).await.into_iter().next() {
            return Some(i);
        }
    }
    debug!("multi_lookup: package {:?} not found", package);
    None
}

type Files = Arc<Mutex<HashMap<PathBuf, PathBuf>>>;

impl<'a> StackElt<'a> {
    /// Some libexec files (e.g. gcc) are searched relative to their
    /// original package /bin path. Propagating the /usr/libexec
    /// folder avoids overrides and ad-hoc code in the client.
    async fn link_extra(&mut self, final_path: &Path, extra: &[&str]) -> Result<(), Error> {
        debug!("linking extra path in {:?}: {:?}", final_path, extra);
        for path in extra {
            let final_extra = final_path.join(path);

            for d in self.deps.iter() {
                let d_extra = d.join(path);
                let Ok(meta) = tokio::fs::metadata(&d_extra).await else {
                    continue;
                };
                let mut stack = vec![(d_extra.clone(), meta)];
                while let Some((elt, meta)) = stack.pop() {
                    if let Ok(mut dir) = tokio::fs::read_dir(&elt).await {
                        let p = elt.strip_prefix(&d_extra).unwrap();
                        tokio::fs::create_dir_all(&final_extra.join(&p)).await?;
                        while let Some(e) = dir.next_entry().await? {
                            stack.push((e.path(), e.metadata().await?));
                        }
                    } else if meta.is_file() {
                        let p = elt.strip_prefix(&d_extra).unwrap();
                        debug!("libexec hard link {:?} to {:?}", elt, final_extra.join(&p));
                        tokio::fs::hard_link(&elt, &final_extra.join(&p))
                            .await
                            .unwrap_or(());
                    }
                }
            }
        }
        Ok(())
    }

    fn add_ld_paths(&mut self, path: &Path) -> Result<(), std::io::Error> {
        let mut ld_so = path.join("etc");
        ld_so.push("ld.so.conf.d");
        if let Ok(dir) = std::fs::read_dir(&ld_so) {
            for f in dir {
                let f = f?;
                if let Ok(f) = std::fs::read_to_string(&f.path()) {
                    for path in f.lines().rev().filter(|x| !x.starts_with("#")) {
                        let path = Path::new(path);
                        if self.ld_path_h.insert(Arc::new(path.to_path_buf())) {
                            self.ld_path.push(Arc::new(path.to_path_buf()))
                        }
                    }
                }
            }
        }
        if !self.ld_path.is_empty() {
            info!("ld_path: {:?}", self.ld_path)
        }
        Ok(())
    }

    /// Let `path` be the path to the raw extracted deb package, and
    /// `dest` be the final destination of this package in the store
    /// (when taking the package's context into account).
    ///
    /// Then, `find_libs` finds all paths in `path` that match the
    /// patterns in self.ld_path (`self.ld_path` paths are of the form
    /// `/usr/lib/x86_64-linux`), and add their equivalent in `dest` to
    /// `self.transitive_deps`.
    ///
    /// This makes sure that packages depending on `self` will be able
    /// to find the correct paths to these libs.
    fn find_libs(&mut self, path: &Path, dest: &Path) {
        for ld in self.ld_path.iter() {
            let path = path.join(ld.strip_prefix("/").unwrap());
            debug!("checking {:?}", path);
            if std::fs::metadata(&path).is_ok() {
                let path = Arc::new(dest.join(ld.strip_prefix("/").unwrap()));
                if self.transitive_deps_h.insert(path.clone()) {
                    self.transitive_deps.push(path);
                }
            }
        }
    }

    async fn patch_elf(
        &mut self,
        f: &Path,
        dest_path: &Path,
        files: &Files,
    ) -> Result<bool, Error> {
        use elfedit::*;
        info!("patch_elf {:?}", f);
        let file = std::fs::OpenOptions::new()
            .read(true)
            .write(true)
            .open(&f)?;

        let mut elf = match Elf::open(&file) {
            Ok(elf) => elf,
            Err(e) => {
                info!("error opening {:?}: {:?}", file, e);
                return Ok(false);
            }
        };
        info!("patching {:?}", f);
        let Some(parsed) = elf.parse().unwrap() else {
            info!("No dynamic section");
            return Ok(false);
        };
        let needed: Vec<_> = parsed
            .needed()
            .map(|x| x.unwrap().to_str().unwrap().to_string())
            .collect();
        let interp = parsed.interpreter();
        if let Some(interp) = interp.unwrap() {
            let interp = interp.to_str().unwrap();
            let files = files.lock().unwrap();

            let p = Path::new(interp).strip_prefix("/").unwrap();

            let subst = if let Some(subst) = files.get(p) {
                subst.join(p)
            } else if interp.starts_with("/usr")
                || interp.starts_with("/lib")
                || interp.starts_with("/bin")
            {
                // https://www.freedesktop.org/wiki/Software/systemd/TheCaseForTheUsrMerge/
                let p2 = "usr".to_string() + interp;
                let p2 = Path::new(&p2);
                debug!("{:?}", p2);
                if let Some(subst) = files.get(p2) {
                    subst.join(p2)
                } else {
                    error!("No subst for {:?}", p2);
                    let mut p2 = p2.to_path_buf();
                    while p2.pop() {
                        debug!("p2 = {:?} {:?}", p2, files.get(&p2));
                    }
                    return Ok(false);
                }
            } else {
                // TODO: not sure what else to do here, we
                // might want to set the interpreter to a
                // different value (equivalent to "recompiling
                // everything" on Nix).
                info!("Interpreter is {interp}. Already patched?");
                return Ok(false);
            };
            let subst = CString::new(subst.to_str().unwrap()).unwrap();
            info!("set interpreter {:?}", subst);
            elf.set_interpreter(subst.to_bytes_with_nul());
        } else if needed.is_empty() {
            // No need to patch
            return Ok(false);
        }

        let mut deps_h = BTreeSet::new();
        let mut path = String::new();
        for dep in self.transitive_deps.iter().rev() {
            debug!("dep of {:?}: {:?}", f, dep);
            if !deps_h.insert(dep) {
                continue;
            }
            let mut dep = dep.to_path_buf();
            let mut is_needed = false;
            for n in needed.iter() {
                dep.push(&n);
                is_needed |= tokio::fs::metadata(&dep).await.is_ok();
                dep.pop();
            }
            if !is_needed {
                continue;
            }
            if !path.is_empty() {
                path.push(':')
            }
            path.push_str(dep.to_str().unwrap());
        }
        path.push('\0');
        info!("Setting path {:?}", path);
        if path.len() > 1 {
            elf.set_runpath(&path.as_bytes());
        }

        debug!("patching into {:?}", dest_path);
        Ok(elf.update(Some(dest_path)).unwrap()) // map_err(From::from)
    }

    async fn finalize(
        &mut self,
        client: &Client,
        files: &Files,
        downloaded: Downloaded,
        stack: &mut Vec<StackElt<'a>>,
        result: &mut Vec<Arc<PathBuf>>,
    ) -> Result<Arc<PathBuf>, Error> {
        let mut context_hasher = blake3::Hasher::new();
        context_hasher.update(self.package.sha256.unwrap().as_bytes());
        self.deps.sort();
        for d in self.deps.iter() {
            context_hasher.update(d.to_str().unwrap().as_bytes());
        }
        debug!(
            "finalize {:?} {:#?}",
            self.package.package, self.transitive_deps
        );
        let dest = client
            .store_path
            .join(data_encoding::HEXLOWER.encode(context_hasher.finalize().as_bytes()));

        let lock = client.lock_store_path(&dest).await;

        // Find the extra ld paths to look for.
        self.add_ld_paths(&downloaded.path).unwrap();

        let initial_deps_len = self.transitive_deps.len();
        // Find the libs in this package.
        self.find_libs(&downloaded.path, &dest);

        let base_package_name = self.package.file_name.unwrap().split('/').last().unwrap();
        let base_package_name = Path::new(&base_package_name).with_extension("");
        let base_package_name = base_package_name.to_str().unwrap();

        let final_path = if std::fs::metadata(&dest).is_err() {
            info!("create final path for {dest:?}");
            match self
                .create_final_path(client, &files, &downloaded.path, &dest, &base_package_name)
                .await
            {
                Ok(x) => x,
                Err(e) => {
                    tokio::fs::remove_dir_all(&dest).await.unwrap_or(());
                    return Err(e);
                }
            }
        } else {
            info!("found, no patching: {:?}", dest);
            let mut output_hasher = blake3::Hasher::new();
            let blakesums = dest.join("blake3sums");
            let file = match tokio::fs::File::open(&blakesums).await {
                Ok(file) => file,
                Err(e) => {
                    error!("Error {:?} {:?}: {:?}", blakesums, downloaded.path, e);
                    return Err(e.into());
                }
            };
            hash_reader(file, &mut output_hasher).await?;

            let r = tokio::io::BufReader::new(
                tokio::fs::File::open(&dest.join("paths")).await.unwrap(),
            );
            let mut l = r.lines();
            while let Some(l) = l.next_line().await? {
                self.files.insert(Arc::new(l));
            }

            client.store_path.join(&format!(
                "{}-{}",
                data_encoding::HEXLOWER.encode(output_hasher.finalize().as_bytes()),
                base_package_name,
            ))
        };
        let final_path = Arc::new(final_path);

        add_subst(downloaded.path.clone(), &dest, files.clone()).await?;

        // Replace prefix for all library deps we've just added,
        // in order to get a Merkle tree.
        for dep in &mut self.transitive_deps[initial_deps_len..] {
            let end = dep.strip_prefix(&dest).unwrap();
            *dep = Arc::new(final_path.join(&end));
        }
        info!("symlink {:?} {:?}", downloaded.path, final_path);
        match std::os::unix::fs::symlink(&dest, &*final_path) {
            Ok(()) => (),
            Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
                let got = std::fs::read_link(&*final_path)?;
                debug!("Path already exists, previous value {:?}", got);
                // This situation means that we've come to the same
                // result via different build recipes.
                /*
                    if dest != got {
                        return Err(Error::WrongResultSymlink {
                            expected: dest,
                            got,
                        });
                }
                     */
            }
            Err(e) => return Err(e.into()),
        }
        result.push(final_path.clone());

        // Maintenant, faire remonter les deps.
        if let Some(last) = stack.iter_mut().rev().filter(|x| x.rx.is_some()).next() {
            for ld in self.ld_path.iter() {
                if last.ld_path_h.insert(ld.clone()) {
                    last.ld_path.push(ld.clone());
                }
            }
            for dep in self.transitive_deps.iter() {
                if last.transitive_deps_h.insert(dep.clone()) {
                    debug!(
                        "adding transitive dep: {:?} -> {:?}",
                        last.package.package, dep
                    );
                    last.transitive_deps.push(dep.clone());
                }
            }
            for f in self.files.iter() {
                last.files.insert(f.clone());
            }
            last.deps.push(final_path.clone());
        }
        drop(lock);
        info!("done with {:?}", self.package.package);
        Ok(final_path)
    }

    async fn spawn_extract(&mut self, client: Client) {
        let client = client.clone();
        let url = client.url(self.package.file_name.as_deref());
        let sha256 = self.package.sha256.unwrap().to_string();
        let (tx, rx) = tokio::sync::oneshot::channel();
        tokio::spawn(async move {
            let permit = client.download_sem.clone().acquire_owned().await.unwrap();
            info!("downloading {:?}", url);
            let (mut task, _) = match client.download_url(&url, &sha256).await {
                Ok(x) => x,
                Err(e) => {
                    tx.send(Err(e)).unwrap_or(());
                    return Ok(());
                }
            };
            let is_extracted = std::fs::metadata(&task.path.with_extension("")).is_ok();
            info!("finished downloading {:?}", url);
            if !is_extracted {
                // Sets extension
                if let Err(e) = extract_task(&client, &mut task).await {
                    info!("finished extracting {:?} {:?}", url, e);
                    tx.send(Err(e)).unwrap_or(());
                } else {
                    info!("finished extracting {:?}, Ok", url);
                    tx.send(Ok(task)).unwrap_or(());
                }
                info!("sent {:?}", url);
            } else {
                task.path.set_extension("");
                tx.send(Ok(task)).unwrap_or(());
            }
            drop(permit);
            Ok::<_, Error>(())
        });
        self.rx = Some(rx);
    }

    async fn push_deps(self, index: &'a [deb::Index], stack: &mut Vec<StackElt<'a>>) {
        let depends = self.package.depends.clone();
        stack.push(self);
        for dep in depends.iter() {
            match dep {
                deb::Dep::Simple(s) => {
                    debug!("dep {:?}", s);
                    let Some(dep) = multi_lookup(index, &s.name).await else {
                        panic!("could not find {:?}", s.name)
                    };
                    stack.push(StackElt {
                        package: dep,
                        rx: None,
                        deps: Vec::new(),
                        transitive_deps: Vec::new(),
                        transitive_deps_h: BTreeSet::new(),
                        ld_path: Vec::new(),
                        ld_path_h: BTreeSet::new(),
                        files: BTreeSet::new(),
                    })
                }
                deb::Dep::Alternatives { alt } => {
                    debug!("alt {:?}", alt);
                    let stack_len = stack.len();
                    for dep in alt {
                        if let Some(dep_) = multi_lookup(index, &dep.name).await {
                            stack.push(StackElt {
                                package: dep_,
                                rx: None,
                                deps: Vec::new(),
                                transitive_deps: Vec::new(),
                                transitive_deps_h: BTreeSet::new(),
                                ld_path: Vec::new(),
                                ld_path_h: BTreeSet::new(),
                                files: BTreeSet::new(),
                            });
                            break;
                        }
                    }
                    if stack.len() == stack_len {
                        panic!("Not found: {:?}", alt);
                    }
                }
            }
        }
    }

    async fn create_final_path(
        &mut self,
        client: &Client,
        files: &Files,
        downloaded_path: &Path,
        dest: &Path,
        base_package_name: &str,
    ) -> Result<PathBuf, Error> {
        // Link the required libexec before hashing.
        let tmp = async_tempfile::TempDir::new_in(dest.parent().unwrap()).await?;
        self.link_extra(&tmp.dir_path(), &["usr/libexec", "usr/lib/gcc"])
            .await?;

        // Patch the ELFs, now that we have all the deps.
        let mut hashing = Vec::new();
        let mut hashes = Vec::with_capacity(hashing.len());
        debug!("create_final_path {:?}", downloaded_path);
        for (f, meta) in find_files(downloaded_path.to_path_buf())? {
            debug!("f = {:?}", f);
            let rel = f.strip_prefix(&downloaded_path).unwrap();
            let dest_path = tmp.dir_path().join(&rel);

            if meta.is_dir() {
                tokio::fs::create_dir_all(dest_path).await.unwrap_or(());
                continue;
            }

            if meta.is_symlink() {
                // Relink the file to the subst.
                let target = tokio::fs::read_link(&f).await?;
                debug!("relink {:?} -> {:?} {:?}", f, target, target.is_relative());
                let subst = {
                    let l = files.lock().unwrap();
                    let target_rel = rel.parent().unwrap().join(&target);
                    if let Some(subst) = l.get(&target_rel) {
                        Some(subst.join(&target_rel))
                    } else {
                        None
                    }
                };
                if let Some(subst) = subst {
                    debug!("relink to {:?}", dest_path);
                    tokio::fs::symlink(&subst, &dest_path).await?;
                    hashes.push((f, subst.to_str().unwrap().to_string()))
                } else {
                    // Leave the symlink untouched
                    // tokio::fs::create_dir_all(dest_path.parent().unwrap()).await.unwrap_or(());
                    debug!("symlink {:?} {:?} {:?}", target, dest_path, f);
                    tokio::fs::symlink(&target, &dest_path).await?;
                    hashes.push((f, target.to_str().unwrap().to_string()));
                }
                continue;
            }

            if !self
                .patch_elf(&f, &dest_path, &files)
                .await
                .unwrap_or(false)
            {
                // Hard link
                debug!("hard link {:?} {:?}", f, dest_path);
                tokio::fs::hard_link(&f, &dest_path).await.unwrap_or(());
            }

            hashing.push(tokio::spawn(async move {
                // hash + write
                info!("hashing {:?}", f);
                if let Ok(file) = tokio::fs::File::open(&dest_path).await {
                    let mut hasher = blake3::Hasher::new();
                    hash_reader(file, &mut hasher).await.unwrap();
                    let hex = data_encoding::HEXLOWER.encode(hasher.finalize().as_bytes());
                    Ok::<_, Error>(Some((f, hex)))
                } else {
                    Ok(None)
                }
            }));
        }
        info!("patched all");

        for h in hashing.into_iter() {
            if let Some(h) = h.await.unwrap().unwrap() {
                hashes.push(h)
            }
        }
        hashes.sort_by(|a, b| a.0.cmp(&b.0));
        info!("hashed all");

        let mut output_hasher = blake3::Hasher::new();
        {
            let blakesums = tmp.dir_path().join("blake3sums");
            let mut file = tokio::fs::File::create(&blakesums).await?;
            for (path, hash) in hashes {
                let path = path
                    .strip_prefix(&downloaded_path)
                    .unwrap()
                    .to_str()
                    .unwrap();
                file.write_all(hash.as_bytes()).await?;
                file.write_all(b" ").await?;
                file.write_all(path.as_bytes()).await?;
                file.write_all(b"\n").await?;
                writeln!(output_hasher, "{} {}", hash, path)?;
            }
        }
        let final_path = client.store_path.join(&format!(
            "{}-{}",
            data_encoding::HEXLOWER.encode(output_hasher.finalize().as_bytes()),
            base_package_name,
        ));

        {
            let transitive = tmp.dir_path().join("paths");
            let mut file = tokio::fs::File::create(&transitive).await?;
            for path in self.files.iter() {
                file.write_all(path.as_bytes()).await?;
                file.write_all(b"\n").await?;
            }
        }

        for (f, _) in find_dirs(downloaded_path.to_path_buf())? {
            let rel = f.strip_prefix(&downloaded_path).unwrap();
            self.files
                .insert(Arc::new(final_path.join(rel).to_str().unwrap().to_string()));
        }

        info!("rename {:?} to {:?}", tmp.dir_path(), dest);
        tokio::fs::rename(tmp.dir_path(), dest).await?;
        std::mem::forget(tmp);

        Ok(final_path)
    }
}

#[derive(Debug)]
struct StackElt<'a> {
    package: deb::Stanza<'a>,
    rx: Option<tokio::sync::oneshot::Receiver<Result<Downloaded, Error>>>,
    deps: Vec<Arc<PathBuf>>,
    transitive_deps: Vec<Arc<PathBuf>>,
    transitive_deps_h: BTreeSet<Arc<PathBuf>>,
    ld_path: Vec<Arc<PathBuf>>,
    ld_path_h: BTreeSet<Arc<PathBuf>>,
    files: BTreeSet<Arc<String>>,
}