use std::collections::{BTreeMap, HashMap};
use std::ffi::{OsStr, OsString};
use std::fs::File;
use std::io::{self, Read};
use std::os::fd::AsFd;
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};

use rustix::fs::{Mode, OFlags, RenameFlags, Timestamps, fchmod, fchown, flistxattr, futimens};
use rustix::process::{Gid, Uid};
use rustix::time::Timespec;
use sha2::{Digest, Sha256};

use crate::backing::BackingRoot;
use crate::classifier::{Classification, ClassificationSnapshot, GitClassifier, PolicyRecord};
use crate::migration::{TRANSACTIONS, reserved};

const ROUTING: &str = "routing";
const POLICY: &str = "policy.json";
const POLICY_TEMPORARY: &str = "policy.tmp";
const CANDIDATE_PREFIX: &[u8] = b"candidate-";
const LEGACY_STAGE_PREFIX: &[u8] = b"stage-";
static CANDIDATE_SEQUENCE: AtomicU64 = AtomicU64::new(0);

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct ObjectKey {
    device: u64,
    inode: u64,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum EntryKind {
    Directory,
    Regular,
    Symlink,
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct AdmissionRecord {
    kind: EntryKind,
    key: ObjectKey,
    links: u64,
    mode: u32,
    uid: u32,
    gid: u32,
    size: i64,
    access_seconds: i64,
    access_nanoseconds: u64,
    modification_seconds: i64,
    modification_nanoseconds: u64,
    change_seconds: i64,
    change_nanoseconds: u64,
    digest: Option<[u8; 32]>,
    link_target: Option<Vec<u8>>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct Transition {
    path: PathBuf,
    directory: bool,
    old: Classification,
    new: Classification,
}

pub(crate) fn admit_policy(
    workspace: &Path,
    host: &BackingRoot,
    isolated: &BackingRoot,
    classifier: &GitClassifier,
) -> io::Result<()> {
    let routing = routing_directory();
    match isolated.mkdir(&routing, Mode::from_raw_mode(0o700)) {
        Ok(()) => {}
        Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {}
        Err(error) => return Err(error),
    }
    isolated.sync_directory(Path::new(TRANSACTIONS))?;
    clean_candidates(isolated)?;

    let previous = read_policy(isolated)?;
    let current_record = classifier.policy_record();
    if previous.as_ref() == Some(&current_record) {
        return Ok(());
    }
    if let Some(previous) = previous {
        let previous = GitClassifier::snapshot_from_record(workspace, previous)?;
        reconcile(host, isolated, &previous, &classifier.snapshot())?;
    }
    persist_policy(isolated, &current_record)
}

fn reconcile(
    host: &BackingRoot,
    isolated: &BackingRoot,
    previous: &ClassificationSnapshot,
    current: &ClassificationSnapshot,
) -> io::Result<()> {
    let transitions = transitions(host, isolated, previous, current)?;
    let admissions = capture_admissions(host, &transitions)?;
    let candidate = create_candidate(isolated)?;
    let result = reconcile_candidate(
        host,
        isolated,
        previous,
        current,
        &transitions,
        &admissions,
        &candidate,
    );
    if result.is_ok() {
        remove_tree(isolated, &candidate)?;
        isolated.sync_directory(&routing_directory())?;
    }
    result
}

fn reconcile_candidate(
    host: &BackingRoot,
    isolated: &BackingRoot,
    previous: &ClassificationSnapshot,
    current: &ClassificationSnapshot,
    expected_transitions: &[Transition],
    admissions: &BTreeMap<PathBuf, Option<AdmissionRecord>>,
    candidate: &Path,
) -> io::Result<()> {
    let mut releases = expected_transitions
        .iter()
        .filter(|transition| {
            transition.old == Classification::Isolated && transition.new == Classification::Host
        })
        .collect::<Vec<_>>();
    releases.sort_by_key(|transition| usize::MAX - transition.path.components().count());
    for transition in releases {
        remove_released_path(isolated, current, &transition.path, transition.directory)?;
    }

    let mut directories = admissions
        .iter()
        .filter_map(|(path, record)| {
            record
                .as_ref()
                .filter(|record| record.kind == EntryKind::Directory)
                .map(|record| (path, record))
        })
        .collect::<Vec<_>>();
    directories.sort_by_key(|(path, _)| path.components().count());
    for (path, record) in &directories {
        ensure_directory(isolated, path, record)?;
    }

    let anchors = create_regular_anchors(host, isolated, admissions, candidate)?;
    let names = candidate.join("names");
    let mut name_sequence = 0_u64;
    for (path, record) in admissions {
        let Some(record) = record else {
            if stat_if_exists(isolated, path)?.is_some() {
                remove_tree(isolated, path)?;
            }
            continue;
        };
        match record.kind {
            EntryKind::Directory => {}
            EntryKind::Regular => {
                let staging = names.join(format!("name-{name_sequence}"));
                name_sequence += 1;
                isolated.link(&anchors[&record.key], &staging)?;
                install_name(isolated, &staging, path)?;
            }
            EntryKind::Symlink => {
                let staging = names.join(format!("name-{name_sequence}"));
                name_sequence += 1;
                let target = PathBuf::from(OsString::from_vec(
                    record.link_target.clone().ok_or_else(invalid_manifest)?,
                ));
                isolated.symlink(&target, &staging)?;
                isolated.set_symlink_owner(
                    &staging,
                    Uid::from_raw(record.uid),
                    Gid::from_raw(record.gid),
                )?;
                isolated.set_symlink_times(&staging, &record.timestamps())?;
                install_name(isolated, &staging, path)?;
            }
        }
    }

    directories.sort_by_key(|(path, _)| usize::MAX - path.components().count());
    for (path, record) in directories {
        let directory = isolated.open_directory(path)?;
        set_metadata(&directory, record)?;
        directory.sync_all()?;
    }

    require_admissions(host, admissions)?;
    verify_isolated(isolated, admissions)?;
    let current_transitions = transitions(host, isolated, previous, current)?;
    let current_admissions = capture_admissions(host, &current_transitions)?;
    if !same_admissions(&current_admissions, admissions) {
        return Err(io::Error::from_raw_os_error(libc::EAGAIN));
    }
    require_admissions(host, admissions)
}

fn create_regular_anchors(
    host: &BackingRoot,
    isolated: &BackingRoot,
    admissions: &BTreeMap<PathBuf, Option<AdmissionRecord>>,
    candidate: &Path,
) -> io::Result<HashMap<ObjectKey, PathBuf>> {
    let mut sources = BTreeMap::new();
    for (path, record) in admissions {
        if let Some(record) = record
            && record.kind == EntryKind::Regular
        {
            sources.entry(record.key).or_insert((path, record));
        }
    }
    let objects = candidate.join("objects");
    let mut anchors = HashMap::new();
    for (sequence, (key, (source, record))) in sources.into_iter().enumerate() {
        let anchor = objects.join(format!("object-{sequence}"));
        let mut input = host.open_entry(source, OFlags::RDONLY, Mode::empty())?;
        let mut output = isolated.create_file(&anchor, Mode::from_raw_mode(0o600))?;
        io::copy(&mut input, &mut output)?;
        set_metadata(&output, record)?;
        output.sync_all()?;
        anchors.insert(key, anchor);
    }
    isolated.sync_directory(&objects)?;
    Ok(anchors)
}

fn ensure_directory(
    isolated: &BackingRoot,
    path: &Path,
    record: &AdmissionRecord,
) -> io::Result<()> {
    if let Some(stat) = stat_if_exists(isolated, path)?
        && stat.st_mode & libc::S_IFMT != libc::S_IFDIR
    {
        remove_tree(isolated, path)?;
    }
    isolated.ensure_parent_directories(path)?;
    match isolated.mkdir(path, Mode::from_raw_mode(record.mode & 0o7777)) {
        Ok(()) => {}
        Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {}
        Err(error) => return Err(error),
    }
    Ok(())
}

fn install_name(isolated: &BackingRoot, staging: &Path, destination: &Path) -> io::Result<()> {
    if stat_if_exists(isolated, destination)?.is_some() {
        remove_tree(isolated, destination)?;
    }
    isolated.ensure_parent_directories(destination)?;
    isolated.rename(staging, destination, RenameFlags::NOREPLACE)?;
    isolated.sync_directory(destination.parent().unwrap_or_else(|| Path::new("")))
}

fn capture_admissions(
    host: &BackingRoot,
    transitions: &[Transition],
) -> io::Result<BTreeMap<PathBuf, Option<AdmissionRecord>>> {
    transitions
        .iter()
        .filter(|transition| {
            transition.old == Classification::Host && transition.new == Classification::Isolated
        })
        .map(|transition| {
            capture_record(host, &transition.path, transition.directory)
                .map(|record| (transition.path.clone(), record))
        })
        .collect()
}

fn capture_record(
    root: &BackingRoot,
    path: &Path,
    directory: bool,
) -> io::Result<Option<AdmissionRecord>> {
    let Some(before) = stat_if_exists(root, path)? else {
        return Ok(None);
    };
    let kind = EntryKind::from_mode(before.st_mode)?;
    if (kind == EntryKind::Directory) != directory {
        return Err(invalid_manifest());
    }
    let (digest, link_target) = match kind {
        EntryKind::Directory => (None, None),
        EntryKind::Regular => {
            let file = root.open_entry(path, OFlags::RDONLY, Mode::empty())?;
            require_copyable_xattrs(&file)?;
            (Some(digest(file)?), None)
        }
        EntryKind::Symlink => {
            if root
                .symlink_xattrs(path)?
                .iter()
                .any(|name| unsupported_xattr_name(name.as_bytes()))
            {
                return Err(io::Error::from_raw_os_error(libc::EXDEV));
            }
            (None, Some(root.readlink(path)?))
        }
    };
    let record = AdmissionRecord::new(&before, kind, digest, link_target);
    let after = root.stat(path)?;
    if !record.same_source(&after) {
        return Err(io::Error::from_raw_os_error(libc::EAGAIN));
    }
    Ok(Some(record))
}

fn require_admissions(
    host: &BackingRoot,
    expected: &BTreeMap<PathBuf, Option<AdmissionRecord>>,
) -> io::Result<()> {
    for (path, expected) in expected {
        let actual = capture_record(
            host,
            path,
            expected
                .as_ref()
                .is_some_and(|record| record.kind == EntryKind::Directory),
        )?;
        if !same_admission(actual.as_ref(), expected.as_ref()) {
            return Err(io::Error::from_raw_os_error(libc::EAGAIN));
        }
    }
    Ok(())
}

fn verify_isolated(
    isolated: &BackingRoot,
    expected: &BTreeMap<PathBuf, Option<AdmissionRecord>>,
) -> io::Result<()> {
    let mut source_to_destination = HashMap::new();
    let mut destination_to_source = HashMap::new();
    let mut digests = HashMap::new();
    for (path, expected) in expected {
        let Some(expected) = expected else {
            if stat_if_exists(isolated, path)?.is_some() {
                return Err(invalid_manifest());
            }
            continue;
        };
        let actual = isolated.stat(path)?;
        if !expected.equivalent_destination(&actual) {
            return Err(invalid_manifest());
        }
        match expected.kind {
            EntryKind::Regular => {
                let destination = ObjectKey::from_stat(&actual);
                if destination == expected.key {
                    return Err(invalid_manifest());
                }
                if source_to_destination
                    .insert(expected.key, destination)
                    .is_some_and(|known| known != destination)
                    || destination_to_source
                        .insert(destination, expected.key)
                        .is_some_and(|known| known != expected.key)
                {
                    return Err(invalid_manifest());
                }
                let actual_digest = if let Some(digest) = digests.get(&destination) {
                    *digest
                } else {
                    let digest =
                        digest(isolated.open_entry(path, OFlags::RDONLY, Mode::empty())?)?;
                    digests.insert(destination, digest);
                    digest
                };
                if Some(actual_digest) != expected.digest {
                    return Err(invalid_manifest());
                }
            }
            EntryKind::Symlink => {
                if Some(isolated.readlink(path)?) != expected.link_target {
                    return Err(invalid_manifest());
                }
            }
            EntryKind::Directory => {}
        }
    }
    Ok(())
}

fn transitions(
    host: &BackingRoot,
    isolated: &BackingRoot,
    previous: &ClassificationSnapshot,
    current: &ClassificationSnapshot,
) -> io::Result<Vec<Transition>> {
    let mut paths = BTreeMap::new();
    collect_paths(host, Path::new(""), &mut paths)?;
    collect_paths(isolated, Path::new(""), &mut paths)?;
    paths.retain(|path, _| !reserved(path));
    Ok(paths
        .into_iter()
        .filter_map(|(path, directory)| {
            let old = previous.classify(&path, directory);
            let new = current.classify(&path, directory);
            (old != new).then_some(Transition {
                path,
                directory,
                old,
                new,
            })
        })
        .collect())
}

fn collect_paths(
    root: &BackingRoot,
    directory: &Path,
    paths: &mut BTreeMap<PathBuf, bool>,
) -> io::Result<()> {
    for (name, kind) in root.typed_entries(directory)? {
        let path = directory.join(name);
        if reserved(&path) {
            continue;
        }
        let is_directory = kind == rustix::fs::FileType::Directory
            || (kind == rustix::fs::FileType::Unknown
                && root
                    .stat(&path)
                    .is_ok_and(|stat| stat.st_mode & libc::S_IFMT == libc::S_IFDIR));
        paths
            .entry(path.clone())
            .and_modify(|known| *known |= is_directory)
            .or_insert(is_directory);
        if is_directory {
            collect_paths(root, &path, paths)?;
        }
    }
    Ok(())
}

fn remove_released_path(
    isolated: &BackingRoot,
    current: &ClassificationSnapshot,
    path: &Path,
    directory: bool,
) -> io::Result<()> {
    if stat_if_exists(isolated, path)?.is_none() {
        return Ok(());
    }
    if directory && has_isolated_descendant(isolated, current, path)? {
        return Ok(());
    }
    remove_tree(isolated, path)
}

fn has_isolated_descendant(
    isolated: &BackingRoot,
    current: &ClassificationSnapshot,
    directory: &Path,
) -> io::Result<bool> {
    for (name, kind) in isolated.typed_entries(directory)? {
        let child = directory.join(name);
        let child_directory = kind == rustix::fs::FileType::Directory
            || (kind == rustix::fs::FileType::Unknown
                && isolated
                    .stat(&child)
                    .is_ok_and(|stat| stat.st_mode & libc::S_IFMT == libc::S_IFDIR));
        if current.classify(&child, child_directory) == Classification::Isolated
            || (child_directory && has_isolated_descendant(isolated, current, &child)?)
        {
            return Ok(true);
        }
    }
    Ok(false)
}

fn create_candidate(isolated: &BackingRoot) -> io::Result<PathBuf> {
    let sequence = CANDIDATE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
    let candidate =
        routing_directory().join(format!("candidate-{}-{sequence}", std::process::id()));
    isolated.mkdir(&candidate, Mode::from_raw_mode(0o700))?;
    isolated.mkdir(&candidate.join("objects"), Mode::from_raw_mode(0o700))?;
    isolated.mkdir(&candidate.join("names"), Mode::from_raw_mode(0o700))?;
    isolated.sync_directory(&candidate)?;
    isolated.sync_directory(&routing_directory())?;
    Ok(candidate)
}

fn clean_candidates(isolated: &BackingRoot) -> io::Result<()> {
    let directory = routing_directory();
    for (name, kind) in isolated.typed_entries(&directory)? {
        let path = directory.join(&name);
        if name == OsStr::new(POLICY) && kind == rustix::fs::FileType::RegularFile {
            continue;
        }
        if name == OsStr::new(POLICY_TEMPORARY)
            || name.as_bytes().starts_with(CANDIDATE_PREFIX)
            || name.as_bytes().starts_with(LEGACY_STAGE_PREFIX)
        {
            remove_tree(isolated, &path)?;
            continue;
        }
        return Err(io::Error::from_raw_os_error(libc::EIO));
    }
    isolated.sync_directory(&directory)
}

fn read_policy(isolated: &BackingRoot) -> io::Result<Option<PolicyRecord>> {
    match isolated.read_file(&routing_directory().join(POLICY), 64 * 1024 * 1024) {
        Ok(bytes) => serde_json::from_slice(&bytes)
            .map(Some)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)),
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
        Err(error) => Err(error),
    }
}

fn persist_policy(isolated: &BackingRoot, policy: &PolicyRecord) -> io::Result<()> {
    let directory = routing_directory();
    let temporary = directory.join(POLICY_TEMPORARY);
    let destination = directory.join(POLICY);
    if stat_if_exists(isolated, &temporary)?.is_some() {
        isolated.unlink(&temporary, false)?;
    }
    let bytes = serde_json::to_vec(policy).map_err(io::Error::other)?;
    isolated.write_file(&temporary, &bytes, Mode::from_raw_mode(0o600))?;
    isolated.rename(&temporary, &destination, RenameFlags::empty())?;
    isolated.sync_directory(&directory)
}

fn routing_directory() -> PathBuf {
    Path::new(TRANSACTIONS).join(ROUTING)
}

fn stat_if_exists(root: &BackingRoot, path: &Path) -> io::Result<Option<rustix::fs::Stat>> {
    match root.stat(path) {
        Ok(stat) => Ok(Some(stat)),
        Err(error) if error.raw_os_error() == Some(libc::ENOENT) => Ok(None),
        Err(error) => Err(error),
    }
}

fn remove_tree(root: &BackingRoot, path: &Path) -> io::Result<()> {
    let stat = root.stat(path)?;
    if stat.st_mode & libc::S_IFMT == libc::S_IFDIR {
        for (name, _) in root.typed_entries(path)? {
            remove_tree(root, &path.join(name))?;
        }
        root.unlink(path, true)
    } else {
        root.unlink(path, false)
    }
}

fn require_copyable_xattrs(file: &File) -> io::Result<()> {
    let mut xattrs = vec![0; 64 * 1024];
    let length = flistxattr(file.as_fd(), &mut xattrs)?;
    xattrs.truncate(length);
    if xattrs.split(|byte| *byte == 0).any(unsupported_xattr_name) {
        return Err(io::Error::from_raw_os_error(libc::EXDEV));
    }
    Ok(())
}

fn unsupported_xattr_name(name: &[u8]) -> bool {
    name.starts_with(b"user.")
        || matches!(
            name,
            b"system.posix_acl_access" | b"system.posix_acl_default"
        )
}

fn digest(mut file: File) -> io::Result<[u8; 32]> {
    let mut hasher = Sha256::new();
    let mut buffer = vec![0; 64 * 1024];
    loop {
        let length = file.read(&mut buffer)?;
        if length == 0 {
            break;
        }
        hasher.update(&buffer[..length]);
    }
    Ok(hasher.finalize().into())
}

fn set_metadata(file: &File, record: &AdmissionRecord) -> io::Result<()> {
    fchmod(file.as_fd(), Mode::from_raw_mode(record.mode & 0o7777))?;
    fchown(
        file.as_fd(),
        Some(Uid::from_raw(record.uid)),
        Some(Gid::from_raw(record.gid)),
    )?;
    futimens(file.as_fd(), &record.timestamps())?;
    Ok(())
}

impl ObjectKey {
    fn from_stat(stat: &rustix::fs::Stat) -> Self {
        Self {
            device: stat.st_dev,
            inode: stat.st_ino,
        }
    }
}

impl EntryKind {
    fn from_mode(mode: u32) -> io::Result<Self> {
        match mode & libc::S_IFMT {
            libc::S_IFDIR => Ok(Self::Directory),
            libc::S_IFREG => Ok(Self::Regular),
            libc::S_IFLNK => Ok(Self::Symlink),
            _ => Err(io::Error::from_raw_os_error(libc::EXDEV)),
        }
    }
}

impl AdmissionRecord {
    fn new(
        stat: &rustix::fs::Stat,
        kind: EntryKind,
        digest: Option<[u8; 32]>,
        link_target: Option<Vec<u8>>,
    ) -> Self {
        Self {
            kind,
            key: ObjectKey::from_stat(stat),
            links: stat.st_nlink,
            mode: stat.st_mode,
            uid: stat.st_uid,
            gid: stat.st_gid,
            size: stat.st_size,
            access_seconds: stat.st_atime,
            access_nanoseconds: stat.st_atime_nsec,
            modification_seconds: stat.st_mtime,
            modification_nanoseconds: stat.st_mtime_nsec,
            change_seconds: stat.st_ctime,
            change_nanoseconds: stat.st_ctime_nsec,
            digest,
            link_target,
        }
    }

    fn same_source(&self, stat: &rustix::fs::Stat) -> bool {
        self.key == ObjectKey::from_stat(stat)
            && self.links == stat.st_nlink
            && self.mode == stat.st_mode
            && self.uid == stat.st_uid
            && self.gid == stat.st_gid
            && self.size == stat.st_size
            && self.modification_seconds == stat.st_mtime
            && self.modification_nanoseconds == stat.st_mtime_nsec
            && self.change_seconds == stat.st_ctime
            && self.change_nanoseconds == stat.st_ctime_nsec
    }

    fn same_admission(&self, other: &Self) -> bool {
        self.kind == other.kind
            && self.key == other.key
            && self.links == other.links
            && self.mode == other.mode
            && self.uid == other.uid
            && self.gid == other.gid
            && self.size == other.size
            && self.modification_seconds == other.modification_seconds
            && self.modification_nanoseconds == other.modification_nanoseconds
            && self.change_seconds == other.change_seconds
            && self.change_nanoseconds == other.change_nanoseconds
            && self.digest == other.digest
            && self.link_target == other.link_target
    }

    fn equivalent_destination(&self, stat: &rustix::fs::Stat) -> bool {
        EntryKind::from_mode(stat.st_mode).ok() == Some(self.kind)
            && self.mode & 0o7777 == stat.st_mode & 0o7777
            && self.uid == stat.st_uid
            && self.gid == stat.st_gid
            && (self.kind == EntryKind::Directory || self.size == stat.st_size)
            && self.modification_seconds == stat.st_mtime
            && self.modification_nanoseconds == stat.st_mtime_nsec
    }

    fn timestamps(&self) -> Timestamps {
        Timestamps {
            last_access: Timespec {
                tv_sec: self.access_seconds,
                tv_nsec: i64::try_from(self.access_nanoseconds).unwrap_or(999_999_999),
            },
            last_modification: Timespec {
                tv_sec: self.modification_seconds,
                tv_nsec: i64::try_from(self.modification_nanoseconds).unwrap_or(999_999_999),
            },
        }
    }
}

fn invalid_manifest() -> io::Error {
    io::Error::from_raw_os_error(libc::EIO)
}

fn same_admission(left: Option<&AdmissionRecord>, right: Option<&AdmissionRecord>) -> bool {
    match (left, right) {
        (Some(left), Some(right)) => left.same_admission(right),
        (None, None) => true,
        _ => false,
    }
}

fn same_admissions(
    left: &BTreeMap<PathBuf, Option<AdmissionRecord>>,
    right: &BTreeMap<PathBuf, Option<AdmissionRecord>>,
) -> bool {
    left.len() == right.len()
        && left.iter().all(|(path, left)| {
            right
                .get(path)
                .is_some_and(|right| same_admission(left.as_ref(), right.as_ref()))
        })
}