use std::collections::BTreeSet;
use std::fmt::Write as _;
use std::io;
use std::path::Path;

use rustix::process::{getgid, getuid};
use serde::{Deserialize, Serialize};

pub const REQUIRED_SUBORDINATE_IDS: u32 = 65_536;

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct UserNamespaceIdRange {
    pub container_id: u32,
    pub parent_id: u32,
    pub size: u32,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct UserNamespaceMap {
    pub uid: Vec<UserNamespaceIdRange>,
    pub gid: Vec<UserNamespaceIdRange>,
}

impl UserNamespaceMap {
    /// Resolves the invoking identity and the lowest sufficient configured subordinate ranges.
    ///
    /// # Errors
    /// Refuses missing, malformed, overlapping, or insufficient `/etc/subuid` and `/etc/subgid`
    /// allocations.
    pub fn resolve_current() -> io::Result<Self> {
        let uid = getuid().as_raw();
        let gid = getgid().as_raw();
        let names = account_names(uid)?;
        let subordinate_user_start = select_subordinate_range(Path::new("/etc/subuid"), &names)?;
        let subordinate_group_start = select_subordinate_range(Path::new("/etc/subgid"), &names)?;
        let map = Self {
            uid: vec![
                UserNamespaceIdRange {
                    container_id: 0,
                    parent_id: uid,
                    size: 1,
                },
                UserNamespaceIdRange {
                    container_id: 1,
                    parent_id: subordinate_user_start,
                    size: REQUIRED_SUBORDINATE_IDS,
                },
            ],
            gid: vec![
                UserNamespaceIdRange {
                    container_id: 0,
                    parent_id: gid,
                    size: 1,
                },
                UserNamespaceIdRange {
                    container_id: 1,
                    parent_id: subordinate_group_start,
                    size: REQUIRED_SUBORDINATE_IDS,
                },
            ],
        };
        map.validate()?;
        Ok(map)
    }

    /// # Errors
    /// Refuses empty, zero-sized, overflowing, overlapping, or rootless-incompatible maps.
    pub fn validate(&self) -> io::Result<()> {
        validate_ranges("UID", &self.uid, getuid().as_raw())?;
        validate_ranges("GID", &self.gid, getgid().as_raw())
    }

    #[must_use]
    pub fn uid_helper_arguments(&self, pid: u32) -> Vec<String> {
        helper_arguments(pid, &self.uid)
    }

    #[must_use]
    pub fn gid_helper_arguments(&self, pid: u32) -> Vec<String> {
        helper_arguments(pid, &self.gid)
    }

    #[must_use]
    pub fn uid_proc_text(&self) -> String {
        proc_text(&self.uid)
    }

    #[must_use]
    pub fn gid_proc_text(&self) -> String {
        proc_text(&self.gid)
    }

    /// # Errors
    /// Refuses an unreadable process or a live UID/GID map that differs from this map.
    pub fn verify_process(&self, pid: u32) -> io::Result<()> {
        verify_proc_map(&format!("/proc/{pid}/uid_map"), &self.uid, "UID")?;
        verify_proc_map(&format!("/proc/{pid}/gid_map"), &self.gid, "GID")
    }
}

fn account_names(uid: u32) -> io::Result<BTreeSet<String>> {
    let mut names = BTreeSet::from([uid.to_string()]);
    for line in std::fs::read_to_string("/etc/passwd")?.lines() {
        let fields = line.split(':').collect::<Vec<_>>();
        if fields.len() >= 3 && fields[2].parse::<u32>().ok() == Some(uid) {
            names.insert(fields[0].to_owned());
        }
    }
    Ok(names)
}

fn select_subordinate_range(path: &Path, names: &BTreeSet<String>) -> io::Result<u32> {
    select_subordinate_range_contents(
        &path.display().to_string(),
        &std::fs::read_to_string(path)?,
        names,
    )
}

fn select_subordinate_range_contents(
    label: &str,
    contents: &str,
    names: &BTreeSet<String>,
) -> io::Result<u32> {
    let mut candidates = Vec::new();
    for (line_number, line) in contents.lines().enumerate() {
        if line.trim().is_empty() {
            continue;
        }
        let fields = line.split(':').collect::<Vec<_>>();
        if fields.len() != 3 {
            return Err(io::Error::other(format!(
                "{} row {} is malformed",
                label,
                line_number + 1
            )));
        }
        if !names.contains(fields[0]) {
            continue;
        }
        let start = fields[1].parse::<u32>().map_err(|_| {
            io::Error::other(format!(
                "{} row {} has an invalid start",
                label,
                line_number + 1
            ))
        })?;
        let count = fields[2].parse::<u32>().map_err(|_| {
            io::Error::other(format!(
                "{} row {} has an invalid count",
                label,
                line_number + 1
            ))
        })?;
        if count >= REQUIRED_SUBORDINATE_IDS
            && start.checked_add(REQUIRED_SUBORDINATE_IDS - 1).is_some()
        {
            candidates.push(start);
        }
    }
    candidates.into_iter().min().ok_or_else(|| {
        io::Error::other(format!(
            "{label} has no contiguous {REQUIRED_SUBORDINATE_IDS}-ID allocation for the invoking user"
        ))
    })
}

fn validate_ranges(
    name: &str,
    ranges: &[UserNamespaceIdRange],
    invoking_id: u32,
) -> io::Result<()> {
    if ranges.len() != 2
        || ranges[0]
            != (UserNamespaceIdRange {
                container_id: 0,
                parent_id: invoking_id,
                size: 1,
            })
        || ranges[1].container_id != 1
        || ranges[1].size != REQUIRED_SUBORDINATE_IDS
    {
        return Err(io::Error::other(format!(
            "{name} map does not match the canonical rootless shape"
        )));
    }
    for range in ranges {
        range
            .container_id
            .checked_add(range.size - 1)
            .and_then(|_| range.parent_id.checked_add(range.size - 1))
            .ok_or_else(|| io::Error::other(format!("{name} map overflows")))?;
    }
    let subordinate = ranges[1];
    if invoking_id >= subordinate.parent_id
        && invoking_id - subordinate.parent_id < subordinate.size
    {
        return Err(io::Error::other(format!(
            "{name} subordinate range contains the invoking identity"
        )));
    }
    Ok(())
}

fn helper_arguments(pid: u32, ranges: &[UserNamespaceIdRange]) -> Vec<String> {
    std::iter::once(pid.to_string())
        .chain(ranges.iter().flat_map(|range| {
            [
                range.container_id.to_string(),
                range.parent_id.to_string(),
                range.size.to_string(),
            ]
        }))
        .collect()
}

fn proc_text(ranges: &[UserNamespaceIdRange]) -> String {
    ranges.iter().fold(String::new(), |mut text, range| {
        writeln!(
            text,
            "{} {} {}",
            range.container_id, range.parent_id, range.size
        )
        .expect("writing to String cannot fail");
        text
    })
}

fn verify_proc_map(path: &str, expected: &[UserNamespaceIdRange], name: &str) -> io::Result<()> {
    let observed = std::fs::read_to_string(path)?
        .lines()
        .map(|line| {
            let fields = line.split_whitespace().collect::<Vec<_>>();
            if fields.len() != 3 {
                return Err(io::Error::other(format!("{name} map row is malformed")));
            }
            Ok(UserNamespaceIdRange {
                container_id: fields[0]
                    .parse()
                    .map_err(|_| io::Error::other(format!("{name} container ID is invalid")))?,
                parent_id: fields[1]
                    .parse()
                    .map_err(|_| io::Error::other(format!("{name} parent ID is invalid")))?,
                size: fields[2]
                    .parse()
                    .map_err(|_| io::Error::other(format!("{name} range size is invalid")))?,
            })
        })
        .collect::<io::Result<Vec<_>>>()?;
    if observed != expected {
        return Err(io::Error::other(format!(
            "live {name} map differs from the canonical map"
        )));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn helper_arguments_preserve_canonical_range_order() {
        let map = UserNamespaceMap {
            uid: vec![
                UserNamespaceIdRange {
                    container_id: 0,
                    parent_id: 1000,
                    size: 1,
                },
                UserNamespaceIdRange {
                    container_id: 1,
                    parent_id: 524_288,
                    size: 65_536,
                },
            ],
            gid: Vec::new(),
        };
        assert_eq!(
            map.uid_helper_arguments(42),
            ["42", "0", "1000", "1", "1", "524288", "65536"]
        );
        assert_eq!(map.uid_proc_text(), "0 1000 1\n1 524288 65536\n");
    }

    #[test]
    fn subordinate_selection_uses_the_lowest_complete_allocation() {
        let names = BTreeSet::from(["1000".to_owned(), "developer".to_owned()]);
        assert_eq!(
            select_subordinate_range_contents(
                "/etc/subuid",
                "other:100000:65536\ndeveloper:700000:131072\n1000:524288:65536\n",
                &names,
            )
            .unwrap(),
            524_288
        );
    }

    #[test]
    fn subordinate_selection_refuses_partial_and_overflowing_allocations() {
        let names = BTreeSet::from(["developer".to_owned()]);
        let error = select_subordinate_range_contents(
            "/etc/subgid",
            "developer:524288:65535\ndeveloper:4294967295:65536\n",
            &names,
        )
        .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("no contiguous 65536-ID allocation")
        );
    }
}