//       ___           ___           ___           ___
//      /\__\         /\  \         /\  \         /\__\
//     /:/  /         \:\  \        \:\  \       /::|  |
//    /:/__/           \:\  \        \:\  \     /:|:|  |
//   /::\  \ ___       /::\  \       /::\  \   /:/|:|__|__
//  /:/\:\  /\__\     /:/\:\__\     /:/\:\__\ /:/ |::::\__\
//  \/__\:\/:/  /    /:/  \/__/    /:/  \/__/ \/__/~~/:/  /
//       \::/  /    /:/  /        /:/  /            /:/  /
//       /:/  /     \/__/         \/__/            /:/  /
//      /:/  /                                    /:/  /
//      \/__/                                     \/__/
//
// Copyright (c) 2023, Robert Swinford <robert.swinford<...at...>gmail.com>
//
// For the full copyright and license information, please view the LICENSE file
// that was distributed with this source code.

use crate::library::results::{HttmError, HttmResult};
use crate::library::utility::{find_common_path, fs_type_from_hidden_dir};
use crate::parse::snaps::MapOfSnaps;
use crate::{NILFS2_SNAPSHOT_ID_KEY, ROOT_DIRECTORY, TM_DIR, ZFS_HIDDEN_DIRECTORY};
use hashbrown::{HashMap, HashSet};
use once_cell::sync::Lazy;
use proc_mounts::MountIter;
use rayon::iter::Either;
use rayon::prelude::*;
use std::collections::BTreeMap;
use std::io::Read;
use std::ops::Deref;
use std::path::PathBuf;
use std::process::Command as ExecProcess;
use which::which;

pub const ZFS_FSTYPE: &str = "zfs";
pub const NILFS2_FSTYPE: &str = "nilfs2";
pub const BTRFS_FSTYPE: &str = "btrfs";
pub const SMB_FSTYPE: &str = "smbfs";
pub const NFS_FSTYPE: &str = "nfs";
pub const AFP_FSTYPE: &str = "afpfs";

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum FilesystemType {
    Zfs,
    Btrfs,
    Nilfs2,
    Apfs,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum MountType {
    Local,
    Network,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatasetMetadata {
    pub source: PathBuf,
    pub fs_type: FilesystemType,
    pub mount_type: MountType,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FilterDirs {
    inner: HashSet<PathBuf>,
}

impl Deref for FilterDirs {
    type Target = HashSet<PathBuf>;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

pub trait MaxLen {
    fn max_len(&self) -> usize;
}

impl MaxLen for FilterDirs {
    fn max_len(&self) -> usize {
        self.inner
            .iter()
            .map(|dir| dir.components().count())
            .max()
            .unwrap_or(usize::MAX)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MapOfDatasets {
    inner: HashMap<PathBuf, DatasetMetadata>,
}

impl Deref for MapOfDatasets {
    type Target = HashMap<PathBuf, DatasetMetadata>;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl MaxLen for MapOfDatasets {
    fn max_len(&self) -> usize {
        self.inner
            .keys()
            .map(|mount| mount.components().count())
            .max()
            .unwrap_or(usize::MAX)
    }
}

pub struct BaseFilesystemInfo {
    pub map_of_datasets: MapOfDatasets,
    pub map_of_snaps: MapOfSnaps,
    pub filter_dirs: FilterDirs,
}

static PROC_MOUNTS: Lazy<PathBuf> = Lazy::new(|| PathBuf::from("/proc/mounts"));
static ETC_MNTTAB: Lazy<PathBuf> = Lazy::new(|| PathBuf::from("/etc/mnttab"));

impl BaseFilesystemInfo {
    // divide by the type of system we are on
    // Linux allows us the read proc mounts
    pub fn new() -> HttmResult<Self> {
        let (raw_datasets, filter_dirs_set) = if PROC_MOUNTS.exists() {
            Self::from_proc_mounts()?
        } else if ETC_MNTTAB.exists() {
            Self::from_mnttab()?
        } else {
            Self::from_mount_cmd()?
        };

        let map_of_snaps = MapOfSnaps::new(&raw_datasets)?;

        let map_of_datasets = {
            MapOfDatasets {
                inner: raw_datasets,
            }
        };

        let filter_dirs = {
            FilterDirs {
                inner: filter_dirs_set,
            }
        };

        Ok(BaseFilesystemInfo {
            map_of_datasets,
            map_of_snaps,
            filter_dirs,
        })
    }

    // parsing from proc mounts is both faster and necessary for certain btrfs features
    // for instance, allows us to read subvolumes mounts, like "/@" or "/@home"
    fn from_proc_mounts() -> HttmResult<(HashMap<PathBuf, DatasetMetadata>, HashSet<PathBuf>)> {
        let (map_of_datasets, filter_dirs): (HashMap<PathBuf, DatasetMetadata>, HashSet<PathBuf>) =
            MountIter::new()?
                .par_bridge()
                .flatten()
                // but exclude snapshot mounts.  we want only the raw filesystems
                .filter(|mount_info| {
                    if mount_info.fstype.as_str() == ZFS_FSTYPE
                        && mount_info
                            .dest
                            .to_string_lossy()
                            .contains(ZFS_HIDDEN_DIRECTORY)
                    {
                        return false;
                    }

                    if mount_info.fstype.as_str() == NILFS2_FSTYPE
                        && mount_info
                            .options
                            .iter()
                            .any(|opt| opt.contains(NILFS2_SNAPSHOT_ID_KEY))
                    {
                        return false;
                    }

                    true
                })
                .partition_map(|mount_info| match mount_info.fstype.as_str() {
                    ZFS_FSTYPE => Either::Left((
                        mount_info.dest,
                        DatasetMetadata {
                            source: mount_info.source,
                            fs_type: FilesystemType::Zfs,
                            mount_type: MountType::Local,
                        },
                    )),
                    SMB_FSTYPE | AFP_FSTYPE | NFS_FSTYPE => {
                        match fs_type_from_hidden_dir(&mount_info.dest) {
                            Some(FilesystemType::Zfs) => Either::Left((
                                mount_info.dest,
                                DatasetMetadata {
                                    source: mount_info.source,
                                    fs_type: FilesystemType::Zfs,
                                    mount_type: MountType::Network,
                                },
                            )),
                            Some(FilesystemType::Btrfs) => Either::Left((
                                mount_info.dest,
                                DatasetMetadata {
                                    source: mount_info.source,
                                    fs_type: FilesystemType::Btrfs,
                                    mount_type: MountType::Network,
                                },
                            )),
                            _ => Either::Right(mount_info.dest),
                        }
                    }
                    BTRFS_FSTYPE => {
                        let keyed_options: BTreeMap<&str, &str> = mount_info
                            .options
                            .iter()
                            .filter(|line| line.contains('='))
                            .filter_map(|line| line.split_once('='))
                            .collect();

                        let source = match keyed_options.get("subvol") {
                            Some(subvol) => PathBuf::from(subvol),
                            None => mount_info.source,
                        };

                        Either::Left((
                            mount_info.dest,
                            DatasetMetadata {
                                source,
                                fs_type: FilesystemType::Btrfs,
                                mount_type: MountType::Local,
                            },
                        ))
                    }
                    NILFS2_FSTYPE => Either::Left((
                        mount_info.dest,
                        DatasetMetadata {
                            source: mount_info.source,
                            fs_type: FilesystemType::Nilfs2,
                            mount_type: MountType::Local,
                        },
                    )),
                    _ => Either::Right(mount_info.dest),
                });

        if map_of_datasets.is_empty() {
            Err(HttmError::new("httm could not find any valid datasets on the system.").into())
        } else {
            Ok((map_of_datasets, filter_dirs))
        }
    }

    fn from_mnttab() -> HttmResult<(HashMap<PathBuf, DatasetMetadata>, HashSet<PathBuf>)> {
        let file = std::fs::File::open(&*ETC_MNTTAB)?;

        let mut reader = std::io::BufReader::new(file);
        let mut buffer = String::new();

        reader.read_to_string(&mut buffer)?;

        let (map_of_datasets, filter_dirs): (HashMap<PathBuf, DatasetMetadata>, HashSet<PathBuf>) =
            buffer
                .par_lines()
                // exclude snapshot mounts.  we want the raw filesystem names.
                .filter(|line| !line.contains(ZFS_HIDDEN_DIRECTORY))
                .map(|line| line.split("\t"))
                .filter_map(|mut iter| {
                    let opt_file_system = iter.next();
                    let opt_mount = iter.next();
                    let opt_fs_type = iter.next();

                    opt_file_system.and_then(|file_system| {
                        opt_mount.and_then(|mount| {
                            opt_fs_type.map(|fs_type| (file_system, mount, fs_type))
                        })
                    })
                })
                .partition_map(|(filesystem_str, mount_str, fs_type_str)| {
                    let mount = PathBuf::from(mount_str);

                    match fs_type_str {
                        ZFS_FSTYPE => Either::Left((
                            mount,
                            DatasetMetadata {
                                source: PathBuf::from(filesystem_str),
                                fs_type: FilesystemType::Zfs,
                                mount_type: MountType::Local,
                            },
                        )),
                        SMB_FSTYPE | AFP_FSTYPE | NFS_FSTYPE => {
                            match fs_type_from_hidden_dir(&mount) {
                                Some(fs_type) => Either::Left((
                                    mount,
                                    DatasetMetadata {
                                        source: PathBuf::from(filesystem_str),
                                        fs_type: fs_type,
                                        mount_type: MountType::Network,
                                    },
                                )),
                                None => Either::Right(mount),
                            }
                        }
                        _ => Either::Right(mount),
                    }
                });

        if map_of_datasets.is_empty() {
            Err(HttmError::new("httm could not find any valid datasets on the system.").into())
        } else {
            Ok((map_of_datasets, filter_dirs))
        }
    }

    // old fashioned parsing for non-Linux systems, nearly as fast, works everywhere with a mount command
    // both methods are much faster than using zfs command
    fn from_mount_cmd() -> HttmResult<(HashMap<PathBuf, DatasetMetadata>, HashSet<PathBuf>)> {
        // do we have the necessary commands for search if user has not defined a snap point?
        // if so run the mount search, if not print some errors
        let mount_command = which("mount").map_err(|_err| {
            HttmError::new(
                "'mount' command not be found. Make sure the command 'mount' is in your path.",
            )
        })?;

        let command_output = &ExecProcess::new(mount_command).output()?;

        let stderr_string = std::str::from_utf8(&command_output.stderr)?;

        if !stderr_string.is_empty() {
            return Err(HttmError::new(stderr_string).into());
        }

        let stdout_string = std::str::from_utf8(&command_output.stdout)?;

        // parse "mount" for filesystems and mountpoints
        let (mut map_of_datasets, filter_dirs): (
            HashMap<PathBuf, DatasetMetadata>,
            HashSet<PathBuf>,
        ) = stdout_string
            .par_lines()
            // but exclude snapshot mounts.  we want the raw filesystem names.
            .filter(|line| !line.contains(ZFS_HIDDEN_DIRECTORY))
            .filter(|line| !line.contains(TM_DIR))
            // mount cmd includes and " on " between src and rest
            .filter_map(|line| line.split_once(" on "))
            // where to split, to just have the src and dest of mounts
            .filter_map(|(filesystem, rest)| {
                // GNU Linux mount output
                if rest.contains("type") {
                    let opt_mount = rest.split_once(" type");
                    opt_mount.map(|mount| (filesystem, mount.0))
                // Busybox and BSD mount output
                } else if rest.contains(" (") {
                    let opt_mount = rest.split_once(" (");
                    opt_mount.map(|mount| (filesystem, mount.0))
                // illumos has no clear delimiter except the next space
                // and swaps the mount and filesystem in the output!
                } else {
                    let mount = filesystem;
                    let opt_filesystem = rest.split_once(" ");
                    opt_filesystem.map(|fs| (fs.0, mount))
                }
            })
            .map(|(filesystem, mount)| (PathBuf::from(filesystem), PathBuf::from(mount)))
            // sanity check: does the filesystem exist and have a ZFS hidden dir? if not, filter it out
            // and flip around, mount should key of key/value
            .partition_map(|(source, mount)| match fs_type_from_hidden_dir(&mount) {
                Some(FilesystemType::Zfs) => Either::Left((
                    mount,
                    DatasetMetadata {
                        source,
                        fs_type: FilesystemType::Zfs,
                        mount_type: MountType::Local,
                    },
                )),
                Some(FilesystemType::Btrfs) => Either::Left((
                    mount,
                    DatasetMetadata {
                        source,
                        fs_type: FilesystemType::Btrfs,
                        mount_type: MountType::Local,
                    },
                )),
                _ => Either::Right(mount),
            });

        Self::from_tm_dir(&mut map_of_datasets)?;

        if map_of_datasets.is_empty() {
            Err(HttmError::new("httm could not find any valid datasets on the system.").into())
        } else {
            Ok((map_of_datasets, filter_dirs))
        }
    }

    fn from_tm_dir(map_of_datasets: &mut HashMap<PathBuf, DatasetMetadata>) -> HttmResult<()> {
        let tm_dir_path = std::path::Path::new(TM_DIR);

        if tm_dir_path.exists() {
            std::fs::read_dir(tm_dir_path)?
                .flatten()
                // this is done because sometimes smb shares
                // are mounted in this path too, but they have extensions
                .filter(|mount| mount.path().extension().is_none())
                .map(|mount| {
                    (
                        PathBuf::from(ROOT_DIRECTORY),
                        DatasetMetadata {
                            source: mount.path(),
                            fs_type: FilesystemType::Apfs,
                            mount_type: MountType::Local,
                        },
                    )
                })
                .for_each(|(k, v)| {
                    let _ = map_of_datasets.insert(k, v);
                });
        }

        Ok(())
    }

    // if we have some btrfs mounts, we check to see if there is a snap directory in common
    // so we can hide that common path from searches later
    pub fn common_snap_dir(&self) -> Option<PathBuf> {
        let map_of_datasets: &MapOfDatasets = &self.map_of_datasets;
        let map_of_snaps: &MapOfSnaps = &self.map_of_snaps;

        if map_of_datasets
            .par_iter()
            .any(|(_mount, dataset_info)| dataset_info.fs_type == FilesystemType::Btrfs)
        {
            let vec_snaps: Vec<&PathBuf> = map_of_datasets
                .par_iter()
                .filter(|(_mount, dataset_info)| {
                    if dataset_info.fs_type == FilesystemType::Btrfs {
                        return true;
                    }

                    false
                })
                .filter_map(|(mount, _dataset_info)| map_of_snaps.get(mount))
                .flatten()
                .collect();

            return find_common_path(vec_snaps);
        }

        None
    }
}
