use std::collections::BTreeMap;
use std::ffi::OsString;
use std::fs::{self, File};
use std::io::{self, BufRead, BufReader, Read, Write};
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use lazaret_core::configuration::{
    ConfigurationLoadRequest, SandboxLayoutV1, load_trusted_configuration,
};
use lazaret_core::identity::{CoordinatorId, DomainId, GenerationId};
use lazaret_core::records::{
    ArtifactState, HostProcessOwnershipSlot, PreparedArtifact, ProcessOwnershipCommitAck,
};
use lazaret_linux::process::{
    ProcessOwnershipContext, TargetCommand, child_broker, current_process_identity,
    dispatch_private_process_role,
};
use lazaret_state_anchor::{
    FuseOverlayfsMount, ProtectedRuntimeFiles, RuntimeStateAnchorAuthority,
    RuntimeStateAnchorInherited, StateAnchorAcquisition, StateAnchorLaunchAuthority,
    StateAnchorLayout, StateAnchorSandboxView, dispatch_private_state_anchor_recovery,
    recover_private_state_anchor_mount, unmount_after_coordinator_loss,
};
use rustix::process::Signal;

static SEQUENCE: AtomicU64 = AtomicU64::new(0);

fn main() -> ExitCode {
    let arguments = std::env::args_os().skip(1).collect::<Vec<_>>();
    if arguments
        .first()
        .is_some_and(|value| value == "--private-recovery-holder")
    {
        return exit(run_private_recovery_holder(&arguments));
    }
    if let Some(result) = dispatch_private_state_anchor_recovery(&arguments) {
        return exit(result);
    }
    if let Some(result) = dispatch_private_process_role(&arguments) {
        return exit(result);
    }
    exit(run())
}

fn exit(result: io::Result<()>) -> ExitCode {
    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("State Anchor lifecycle fixture failed: {error}");
            ExitCode::FAILURE
        }
    }
}

#[allow(clippy::too_many_lines)]
fn run() -> io::Result<()> {
    match fs::metadata("/dev/fuse") {
        Ok(_) => {}
        Err(error) if error.kind() == io::ErrorKind::NotFound => {
            eprintln!("skipped: /dev/fuse is unavailable");
            return Ok(());
        }
        Err(error) => return Err(error),
    }
    let tree = TestTree::new().map_err(|error| io::Error::other(format!("test tree: {error}")))?;
    let configuration = configuration(&tree)
        .map_err(|error| io::Error::other(format!("configuration: {error}")))?;
    let lower = tree.path.join("prepared-lower");
    fs::create_dir_all(lower.join("usr/bin"))?;
    fs::create_dir_all(lower.join("etc"))?;
    fs::write(lower.join("usr/bin/base-tool"), b"base")?;
    for name in ["resolv.conf", "hosts", "hostname"] {
        fs::write(lower.join("etc").join(name), b"image")?;
    }
    let metadata = lower.metadata()?;
    let artifact = PreparedArtifact {
        prepared_artifact_id: lazaret_core::identity::PreparedArtifactId::new([8; 32]),
        source_identity: lazaret_core::identity::SourceIdentity::new(b"fixture".to_vec()),
        selected_manifest_digest: lazaret_core::identity::SelectedManifestDigest::new([9; 32]),
        policy_id: lazaret_core::identity::PolicyId::new([10; 32]),
        id_map_digest: lazaret_core::identity::IdMapDigest::new([11; 32]),
        path: lazaret_core::identity::HostPath::new(lower.as_os_str().as_encoded_bytes().to_vec()),
        device: lazaret_core::identity::DeviceId::new(metadata.dev()),
        inode: lazaret_core::identity::InodeId::new(metadata.ino()),
        state: ArtifactState::Available,
    };
    let layout = StateAnchorLayout::create(&configuration)
        .map_err(|error| io::Error::other(format!("State Anchor layout: {error}")))?;
    exercise_private_recovery(&layout)
        .map_err(|error| io::Error::other(format!("private recovery exercise: {error}")))?;

    let launch =
        StateAnchorLaunchAuthority::capture(&layout, &artifact).map_err(io::Error::other)?;
    let runtime_authority =
        RuntimeStateAnchorAuthority::from_inherited(RuntimeStateAnchorInherited {
            lower_fd: launch.lower_fd(),
            lower_label: launch.lower_label(),
            persistent_root_fd: launch.persistent_root_fd(),
            persistent_anchor_fd: launch.persistent_anchor_fd(),
            persistent_anchor_label: launch.persistent_anchor_label(),
            runtime_root_fd: launch.runtime_root_fd(),
            runtime_anchor_fd: launch.runtime_anchor_fd(),
            runtime_anchor_label: launch.runtime_anchor_label(),
            expected_lower: launch.expected_lower(),
        })
        .map_err(io::Error::other)?;
    let expected = std::cell::Cell::new(None);
    let runtime_mount = runtime_authority
        .acquire(
            current_process_identity()?,
            ProcessOwnershipContext {
                domain_id: DomainId::new([1; 32]),
                generation_id: GenerationId::new(1),
                coordinator_id: CoordinatorId::new([3; 32]),
                slot: HostProcessOwnershipSlot::StateAnchorDaemon,
            },
            |staged| {
                expected.set(Some(staged));
                Ok(())
            },
            |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
            |readiness| {
                if expected.get() != Some(readiness.staged())
                    || lazaret_linux::filesystem::current_mount_identity(layout.merged())?
                        != readiness.mount()
                {
                    return Err(io::Error::other(
                        "inherited State Anchor readiness changed authority",
                    ));
                }
                Ok(())
            },
        )
        .map_err(io::Error::other)?;
    let runtime_targets = [
        PathBuf::from("etc/resolv.conf"),
        PathBuf::from("etc/hosts"),
        PathBuf::from("etc/hostname"),
    ];
    runtime_mount
        .require_persistent_destinations(&[PathBuf::from("usr"), PathBuf::from("etc")])
        .map_err(io::Error::other)?;
    runtime_mount
        .require_transient_override_targets(&runtime_targets)
        .map_err(io::Error::other)?;
    fs::write(layout.merged().join("etc/hosts"), b"stale")?;
    if runtime_mount
        .require_transient_override_targets(&runtime_targets)
        .is_ok()
    {
        return Err(io::Error::other(
            "upper runtime-file entry survived transient override preflight",
        ));
    }
    let held_directory = layout.merged().join("root/open-lock");
    fs::create_dir_all(&held_directory)?;
    let held = File::open(&held_directory)?;
    fs::remove_dir(&held_directory)?;
    drop(held);
    if layout.work().join("work").read_dir()?.next().is_none() {
        return Err(io::Error::other(
            "open-directory removal did not create fuse-overlayfs work residue",
        ));
    }
    runtime_mount.unmount().map_err(io::Error::other)?;
    layout.inspect_state().map_err(io::Error::other)?;
    fs::remove_file(layout.upper().join("etc/hosts"))?;

    let first = acquire(&layout, &artifact)?;
    if layout.merged().join("root/open-lock").exists() {
        return Err(io::Error::other(
            "removed open directory returned after normalized remount",
        ));
    }
    fs::write(
        layout.merged().join("usr/bin/installed-tool"),
        b"persistent",
    )?;
    fs::write(layout.merged().join("etc/developer.conf"), b"persistent")?;
    if lower.join("usr/bin/installed-tool").exists()
        || lower.join("etc/developer.conf").exists()
        || !layout.upper().join("usr/bin/installed-tool").exists()
        || !layout.upper().join("etc/developer.conf").exists()
    {
        return Err(io::Error::other(
            "writes did not remain upper-only over the read-only lower",
        ));
    }
    let view = StateAnchorSandboxView::from_mounted(&layout, &artifact, first.mount_identity())
        .map_err(io::Error::other)?;
    view.revalidate().map_err(io::Error::other)?;
    layout
        .replace_protected_runtime_files(&ProtectedRuntimeFiles {
            resolv_conf: b"nameserver 192.0.2.53\n",
            hosts: b"127.0.0.1 localhost\n",
            hostname: b"lazaret\n",
        })
        .map_err(io::Error::other)?;
    let runtime_sources = view.protected_runtime_mounts().collect::<Vec<_>>();
    if runtime_sources.len() != 3
        || runtime_sources
            .iter()
            .any(|(source, _)| source.starts_with(&lower) || source.starts_with(layout.upper()))
    {
        return Err(io::Error::other(
            "protected runtime files are not host-private runtime material",
        ));
    }
    if lower.join("etc/resolv.conf").read_link().is_ok()
        || fs::read(lower.join("etc/resolv.conf"))? != b"image"
        || layout.upper().join("etc/resolv.conf").exists()
        || fs::read(&runtime_sources[0].0)? != b"nameserver 192.0.2.53\n"
    {
        return Err(io::Error::other(
            "runtime-file replacement escaped transient storage",
        ));
    }
    drop(view);
    let wrong = lazaret_linux::filesystem::current_mount_identity(layout.upper())?;
    if unmount_after_coordinator_loss(&layout, wrong).is_ok() {
        return Err(io::Error::other(
            "ownership mismatch unmounted foreign identity",
        ));
    }
    first.unmount().map_err(io::Error::other)?;
    layout.inspect_state().map_err(io::Error::other)?;

    let second = acquire(&layout, &artifact)?;
    if fs::read(layout.merged().join("usr/bin/installed-tool"))? != b"persistent"
        || fs::read(layout.merged().join("etc/developer.conf"))? != b"persistent"
    {
        return Err(io::Error::other(
            "upper state did not persist across remount",
        ));
    }
    let identity = second.mount_identity();
    second.process().send_signal(Signal::KILL)?;
    while !second.process().has_exited()? {
        std::thread::yield_now();
    }
    unmount_after_coordinator_loss(&layout, identity).map_err(io::Error::other)?;
    second
        .finish_after_external_unmount()
        .map_err(io::Error::other)?;
    drop(tree);
    Ok(())
}

#[allow(deprecated)]
fn run_private_recovery_holder(arguments: &[OsString]) -> io::Result<()> {
    if arguments.len() != 2 {
        return Err(io::Error::other("private holder arguments are malformed"));
    }
    let merged = PathBuf::from(&arguments[1]);
    let uid = rustix::process::getuid().as_raw();
    let gid = rustix::process::getgid().as_raw();
    rustix::thread::unshare(rustix::thread::UnshareFlags::NEWUSER)
        .map_err(|error| io::Error::other(format!("holder user namespace: {error}")))?;
    fs::write("/proc/self/setgroups", "deny")
        .map_err(|error| io::Error::other(format!("holder setgroups: {error}")))?;
    fs::write("/proc/self/uid_map", format!("0 {uid} 1\n"))
        .map_err(|error| io::Error::other(format!("holder uid map: {error}")))?;
    fs::write("/proc/self/gid_map", format!("0 {gid} 1\n"))
        .map_err(|error| io::Error::other(format!("holder gid map: {error}")))?;
    rustix::thread::unshare(rustix::thread::UnshareFlags::NEWNS)
        .map_err(|error| io::Error::other(format!("holder mount namespace: {error}")))?;
    rustix::mount::mount_change(
        "/",
        rustix::mount::MountPropagationFlags::PRIVATE | rustix::mount::MountPropagationFlags::REC,
    )
    .map_err(|error| io::Error::other(format!("holder private propagation: {error}")))?;
    rustix::mount::mount(
        "tmpfs",
        &merged,
        "tmpfs",
        rustix::mount::MountFlags::empty(),
        c"size=1m",
    )
    .map_err(|error| io::Error::other(format!("holder synthetic mount: {error}")))?;
    let user = lazaret_linux::filesystem::namespace_identity(Path::new("/proc/self/ns/user"))?;
    let mount_namespace =
        lazaret_linux::filesystem::namespace_identity(Path::new("/proc/self/ns/mnt"))?;
    let mount = lazaret_linux::filesystem::current_mount_identity(&merged)?;
    writeln!(
        io::stdout(),
        "{} {} {} {} {} {} {}",
        user.device.get(),
        user.inode.get(),
        mount_namespace.device.get(),
        mount_namespace.inode.get(),
        mount.mount_id.get(),
        mount.device.get(),
        mount.inode.get(),
    )?;
    io::stdout().flush()?;
    loop {
        std::thread::park();
    }
}

fn exercise_private_recovery(layout: &StateAnchorLayout<'_>) -> io::Result<()> {
    let target = TargetCommand::new("/proc/self/exe")
        .argument("--private-recovery-holder")
        .argument(layout.merged_label());
    let mut captured = child_broker()
        .run_transient_captured(&target)
        .map_err(|error| io::Error::other(format!("starting private holder: {error}")))?;
    let values = {
        let (stdout, stderr) = captured.streams();
        let mut reader = BufReader::new(stdout);
        let mut line = String::new();
        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            match reader.read_line(&mut line) {
                Ok(0) if Instant::now() < deadline => std::thread::yield_now(),
                Ok(0) => {
                    let mut detail = String::new();
                    match stderr.read_to_string(&mut detail) {
                        Ok(_) => {}
                        Err(error) if error.kind() == io::ErrorKind::WouldBlock => {}
                        Err(error) => return Err(error),
                    }
                    return Err(io::Error::other(format!(
                        "private holder closed without evidence: {}",
                        detail.trim()
                    )));
                }
                Ok(_) => break,
                Err(error)
                    if error.kind() == io::ErrorKind::WouldBlock && Instant::now() < deadline =>
                {
                    std::thread::yield_now();
                }
                Err(error) => return Err(error),
            }
        }
        line.split_whitespace()
            .map(|value| {
                value
                    .parse::<u64>()
                    .map_err(|_| io::Error::other("private holder output is malformed"))
            })
            .collect::<io::Result<Vec<_>>>()?
    };
    if values.len() != 7 {
        return Err(io::Error::other("private holder evidence is incomplete"));
    }
    let holder = captured.into_process();
    let user = lazaret_core::identity::NamespaceIdentity {
        device: lazaret_core::identity::DeviceId::new(values[0]),
        inode: lazaret_core::identity::InodeId::new(values[1]),
    };
    let mount_namespace = lazaret_core::identity::NamespaceIdentity {
        device: lazaret_core::identity::DeviceId::new(values[2]),
        inode: lazaret_core::identity::InodeId::new(values[3]),
    };
    let mount = lazaret_core::identity::MountIdentity {
        mount_id: lazaret_core::identity::MountId::new(values[4]),
        device: lazaret_core::identity::DeviceId::new(values[5]),
        inode: lazaret_core::identity::InodeId::new(values[6]),
    };
    let wrong = lazaret_core::identity::MountIdentity {
        mount_id: lazaret_core::identity::MountId::new(mount.mount_id.get() + 1),
        ..mount
    };
    if recover_private_state_anchor_mount(layout, &holder, user, mount_namespace, wrong).is_ok() {
        return Err(io::Error::other("private recovery accepted wrong mount"));
    }
    let observed =
        recover_private_state_anchor_mount(layout, &holder, user, mount_namespace, mount)
            .map_err(io::Error::other)?;
    if observed == mount {
        return Err(io::Error::other("private recovery did not prove absence"));
    }
    holder.send_signal(Signal::KILL)?;
    while !holder.has_exited()? {
        std::thread::yield_now();
    }
    child_broker().reap(holder)
}

fn acquire<'layout, 'configuration>(
    layout: &'layout StateAnchorLayout<'configuration>,
    artifact: &PreparedArtifact,
) -> io::Result<FuseOverlayfsMount<'layout, 'configuration>> {
    let expected = std::cell::Cell::new(None);
    FuseOverlayfsMount::acquire(
        layout,
        StateAnchorAcquisition {
            lower: Path::new(std::ffi::OsStr::from_bytes(artifact.path.as_bytes())),
            expected_lower: (artifact.device.get(), artifact.inode.get()),
            expected_parent: current_process_identity()?,
            ownership: ProcessOwnershipContext {
                domain_id: DomainId::new([1; 32]),
                generation_id: GenerationId::new(2),
                coordinator_id: CoordinatorId::new([3; 32]),
                slot: HostProcessOwnershipSlot::StateAnchorDaemon,
            },
        },
        |staged| {
            expected.set(Some(staged));
            Ok(())
        },
        |staged| {
            if expected.get() != Some(staged) {
                return Err(io::Error::other("staged process identity changed"));
            }
            Ok(ProcessOwnershipCommitAck::v1(staged))
        },
        |readiness| {
            if expected.get() != Some(readiness.staged()) {
                return Err(io::Error::other("ready process identity changed"));
            }
            if lazaret_linux::filesystem::current_mount_identity(layout.merged())?
                != readiness.mount()
            {
                return Err(io::Error::other(
                    "readiness preceded exact mount verification",
                ));
            }
            Ok(())
        },
    )
    .map_err(io::Error::other)
}

struct TestTree {
    path: PathBuf,
}

impl TestTree {
    fn new() -> io::Result<Self> {
        let parent = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("../../target/test-artifacts/state-anchor-lifecycle");
        fs::create_dir_all(&parent)?;
        let parent = parent.canonicalize()?;
        let path = parent.join(format!(
            "lazaret state-anchor,:{}-{}",
            std::process::id(),
            SEQUENCE.fetch_add(1, Ordering::Relaxed)
        ));
        fs::create_dir(&path)?;
        Ok(Self { path })
    }
}

impl Drop for TestTree {
    fn drop(&mut self) {
        let _ = fs::remove_dir_all(&self.path);
    }
}

fn configuration(
    tree: &TestTree,
) -> io::Result<lazaret_core::configuration::ResolvedDomainConfiguration> {
    let configuration_directory = tree.path.join("configuration");
    let workspace = tree.path.join("workspace");
    let runtime = tree.path.join("runtime");
    fs::create_dir(&configuration_directory)?;
    fs::create_dir(&workspace)?;
    fs::create_dir(&runtime)?;
    let configuration_path = configuration_directory.join("lazaret.toml");
    fs::write(
        &configuration_path,
        format!(
            "schema = \"lazaret-trusted-host-config-v1\"\n[storage]\ndata_root = \"{}\"\n[network]\ndefault_allow = [\"192.0.2.10/32\"]\n[profiles.default]\nimage = \"example.invalid/fixture:v1\"\npersistent_dirs = [\"/usr\", \"/etc\"]
",
            tree.path.join("data").display()
        ),
    )?;
    let layout = SandboxLayoutV1::from_protected_destinations([
        "/workspace".to_owned(),
        "/proc".to_owned(),
        "/sys".to_owned(),
        "/dev".to_owned(),
        "/run/lazaret".to_owned(),
    ])
    .map_err(io::Error::other)?;
    load_trusted_configuration(ConfigurationLoadRequest {
        explicit_configuration_path: Some(configuration_path),
        selected_profile: "default".to_owned(),
        workspace,
        home: tree.path.join("home"),
        xdg_config_home: None,
        xdg_data_home: None,
        xdg_runtime_dir: runtime,
        environment: BTreeMap::<String, OsString>::new(),
        sandbox_layout: layout,
    })
    .map_err(io::Error::other)
}