use std::ffi::{OsStr, OsString};
use std::fs::File;
use std::io;
use std::os::fd::{AsFd, AsRawFd, OwnedFd};
use std::os::unix::fs::FileExt;
use std::time::{Duration, Instant};

use lazaret_core::codec::{decode, encode};
use lazaret_core::identity::{CoordinatorId, DomainId, GenerationId};
use lazaret_core::records::{
    HostProcessOwnershipSlot, ProcessOwnershipCommitAck, StagedProcessOwnership,
};
use lazaret_linux::capability;
use lazaret_linux::process::{
    ChildStandardIo, ParentEdge, ProcessOwnershipContext, RetainedProcess, TargetCommand,
    child_broker, current_process_identity,
};
use rustix::event::{PollFd, PollFlags, Timespec, poll};
use rustix::fs::{MemfdFlags, memfd_create};
use rustix::io::{FdFlags, fcntl_setfd, read, write};
use rustix::process::Signal;
use rustix::thread::CapabilitySet;
use serde::{Deserialize, Serialize};

use crate::sandbox::RunscCommand;

const PRIVATE_ROLE_ARGUMENT: &str = "--lazaret-runtime-supervisor-role";
const MAX_CONTROL_MESSAGE: usize = 16 * 1024;
const READINESS_TIMEOUT: Duration = Duration::from_secs(30);
const TERMINATION_GRACE: Duration = Duration::from_secs(2);

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum AttachedRunscSignal {
    Stop,
    Continue,
    Kill,
}

impl AttachedRunscSignal {
    const fn signal(self) -> Signal {
        match self {
            Self::Stop => Signal::STOP,
            Self::Continue => Signal::CONT,
            Self::Kill => Signal::KILL,
        }
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum SupervisorEvent {
    Staged(StagedProcessOwnership),
    Ready {
        staged: StagedProcessOwnership,
        sentry_pid_inventory: u32,
    },
    SignalApplied(AttachedRunscSignal),
    AttachedRunscExited,
    Stopped,
    Refused(String),
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
enum CoordinatorMessage {
    StagedAck(ProcessOwnershipCommitAck),
    SandboxCommitted,
    SignalAttachedRunsc(AttachedRunscSignal),
    Stop,
}

pub struct CoordinatorChannel {
    descriptor: OwnedFd,
}

impl CoordinatorChannel {
    #[must_use]
    pub const fn new(descriptor: OwnedFd) -> Self {
        Self { descriptor }
    }

    /// # Errors
    /// Returns an error for a closed, malformed, oversized, or unknown message.
    pub fn receive(&self) -> io::Result<SupervisorEvent> {
        receive_message(&self.descriptor)
    }

    /// # Errors
    /// Returns an error when the exact typed acknowledgment cannot be delivered.
    pub fn acknowledge_staged(&self, acknowledgment: ProcessOwnershipCommitAck) -> io::Result<()> {
        send_message(
            &self.descriptor,
            &CoordinatorMessage::StagedAck(acknowledgment),
        )
    }

    /// # Errors
    /// Returns an error when the durable sandbox commit cannot be acknowledged.
    pub fn acknowledge_sandbox_committed(&self) -> io::Result<()> {
        send_message(&self.descriptor, &CoordinatorMessage::SandboxCommitted)
    }

    /// # Errors
    /// Returns an error when the exact retained-runsc action cannot be requested.
    pub fn signal_attached_runsc(&self, signal: AttachedRunscSignal) -> io::Result<()> {
        send_message(
            &self.descriptor,
            &CoordinatorMessage::SignalAttachedRunsc(signal),
        )
    }

    /// # Errors
    /// Returns an error when ordered sandbox stop cannot be requested.
    pub fn stop(&self) -> io::Result<()> {
        send_message(&self.descriptor, &CoordinatorMessage::Stop)
    }
}

pub struct RuntimeSupervisorLaunch {
    target: TargetCommand,
    ownership: ProcessOwnershipContext,
}

impl RuntimeSupervisorLaunch {
    /// # Errors
    /// Refuses any command other than the deterministic foreground attached-runsc form.
    pub fn from_runsc_command(
        ownership: ProcessOwnershipContext,
        command: &RunscCommand,
    ) -> io::Result<Self> {
        if ownership.slot != HostProcessOwnershipSlot::SandboxRuntimeSupervisor
            || command.program() != OsStr::new("runsc")
            || !command
                .arguments()
                .iter()
                .any(|argument| argument == OsStr::new("run"))
            || command
                .arguments()
                .iter()
                .any(|argument| matches!(argument.to_str(), Some("-detach" | "--detach" | "-d")))
        {
            return Err(io::Error::other(
                "Runtime Supervisor requires foreground attached runsc",
            ));
        }
        let mut target = TargetCommand::new("/proc/self/exe")
            .argument(PRIVATE_ROLE_ARGUMENT)
            .argument("--domain")
            .argument(hex(ownership.domain_id.as_bytes()))
            .argument("--generation")
            .argument(ownership.generation_id.get().to_string())
            .argument("--coordinator")
            .argument(hex(ownership.coordinator_id.as_bytes()))
            .argument("--runsc")
            .argument(command.program());
        for argument in command.arguments() {
            target = target.argument(argument);
        }
        Ok(Self {
            target: target.with_lifetime_pipe_parent_edge(),
            ownership,
        })
    }

    #[must_use]
    pub const fn target(&self) -> &TargetCommand {
        &self.target
    }

    #[must_use]
    pub const fn ownership(&self) -> ProcessOwnershipContext {
        self.ownership
    }

    /// # Errors
    /// Refuses a Runtime Supervisor target without the lifetime-pipe parent edge.
    pub fn child_standard_io(
        &self,
        private_control: OwnedFd,
        lifetime_pipe: OwnedFd,
    ) -> io::Result<ChildStandardIo> {
        if self.target.parent_edge() != ParentEdge::LifetimePipe {
            return Err(io::Error::other(
                "Runtime Supervisor launch requires the lifetime-pipe parent edge",
            ));
        }
        Ok(ChildStandardIo::private_control_with_lifetime_pipe(
            private_control,
            lifetime_pipe,
        ))
    }
}

/// Runs the Runtime Supervisor's private early role before ordinary binary dispatch.
///
/// # Errors
/// Returns a fail-closed parse, acquisition, capability, channel, or process error.
#[must_use]
pub fn dispatch_private_runtime_role(arguments: &[OsString]) -> Option<io::Result<()>> {
    if arguments
        .first()
        .is_none_or(|argument| argument != PRIVATE_ROLE_ARGUMENT)
    {
        return None;
    }
    let result =
        parse_runtime_arguments(&arguments[1..]).and_then(|parsed| run_supervisor(&parsed));
    if let Err(error) = &result {
        let _ = send_message(
            std::io::stdout(),
            &SupervisorEvent::Refused(error.to_string()),
        );
    }
    Some(result)
}

struct RuntimeArguments {
    ownership: ProcessOwnershipContext,
    runsc_program: OsString,
    runsc_arguments: Vec<OsString>,
}

fn parse_runtime_arguments(arguments: &[OsString]) -> io::Result<RuntimeArguments> {
    let mut arguments = arguments.iter();
    let domain = parse_fixed_hex(named(&mut arguments, "--domain")?)?;
    let generation = named(&mut arguments, "--generation")?
        .to_str()
        .and_then(|value| value.parse().ok())
        .ok_or_else(|| io::Error::other("Runtime Supervisor generation is invalid"))?;
    let coordinator = parse_fixed_hex(named(&mut arguments, "--coordinator")?)?;
    if arguments.next().is_none_or(|value| value != "--runsc") {
        return Err(io::Error::other(
            "Runtime Supervisor runsc separator is missing",
        ));
    }
    let program = arguments
        .next()
        .ok_or_else(|| io::Error::other("Runtime Supervisor runsc program is missing"))?;
    if program != "runsc" {
        return Err(io::Error::other("Runtime Supervisor target is not runsc"));
    }
    Ok(RuntimeArguments {
        ownership: ProcessOwnershipContext {
            domain_id: DomainId::new(domain),
            generation_id: GenerationId::new(generation),
            coordinator_id: CoordinatorId::new(coordinator),
            slot: HostProcessOwnershipSlot::SandboxAttachedRunsc,
        },
        runsc_program: (*program).clone(),
        runsc_arguments: arguments.cloned().collect(),
    })
}

fn named<'a>(
    arguments: &mut impl Iterator<Item = &'a OsString>,
    name: &str,
) -> io::Result<&'a OsString> {
    if arguments.next().is_none_or(|value| value != name) {
        return Err(io::Error::other(format!(
            "Runtime Supervisor argument {name} is missing"
        )));
    }
    arguments
        .next()
        .ok_or_else(|| io::Error::other(format!("Runtime Supervisor argument {name} has no value")))
}

fn parse_fixed_hex(value: &OsStr) -> io::Result<[u8; 32]> {
    let value = value
        .to_str()
        .ok_or_else(|| io::Error::other("Runtime Supervisor identity is not UTF-8"))?;
    if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
        return Err(io::Error::other("Runtime Supervisor identity is invalid"));
    }
    let mut bytes = [0_u8; 32];
    for (index, byte) in bytes.iter_mut().enumerate() {
        *byte = u8::from_str_radix(&value[index * 2..index * 2 + 2], 16)
            .map_err(|_| io::Error::other("Runtime Supervisor identity is invalid"))?;
    }
    Ok(bytes)
}

fn run_supervisor(arguments: &RuntimeArguments) -> io::Result<()> {
    require_lifetime_held()?;
    let control = std::io::stdout();
    let parent = current_process_identity()?;
    let pid_file = File::from(memfd_create("lazaret-runsc-pid", MemfdFlags::empty())?);
    fcntl_setfd(pid_file.as_fd(), FdFlags::empty())?;
    let runsc = runsc_target(arguments, pid_file.as_raw_fd())?;
    let staged = std::cell::Cell::new(None);
    let process = child_broker().acquire_with_stdio(
        parent,
        arguments.ownership,
        &runsc,
        attached_runsc_standard_io(&runsc)?,
        |ownership| {
            staged.set(Some(ownership));
            send_message(&control, &SupervisorEvent::Staged(ownership))
        },
        |ownership| match receive_message(&control)? {
            CoordinatorMessage::StagedAck(acknowledgment)
                if acknowledgment.acknowledges(&ownership) =>
            {
                Ok(acknowledgment)
            }
            _ => Err(io::Error::other(
                "Runtime Supervisor received a mismatched staged acknowledgment",
            )),
        },
    )?;
    if let Err(error) = fcntl_setfd(pid_file.as_fd(), FdFlags::CLOEXEC) {
        return fail_with_exact_cleanup(process, error.into());
    }
    let Some(staged) = staged.get() else {
        return fail_with_exact_cleanup(
            process,
            io::Error::other("attached runsc ownership was not staged"),
        );
    };
    let sentry_pid_inventory = match wait_for_runsc_ready(&pid_file, &process) {
        Ok(pid) => pid,
        Err(error) => return fail_with_exact_cleanup(process, error),
    };
    if let Err(error) = capability::irreversibly_drop(CapabilitySet::NET_ADMIN) {
        return fail_with_exact_cleanup(process, error);
    }
    if let Err(error) = capability::require_regain_refused(CapabilitySet::NET_ADMIN) {
        return fail_with_exact_cleanup(process, error);
    }
    if let Err(error) = send_message(
        &control,
        &SupervisorEvent::Ready {
            staged,
            sentry_pid_inventory,
        },
    ) {
        return fail_with_exact_cleanup(process, error);
    }
    let committed: CoordinatorMessage = match receive_message(&control) {
        Ok(committed) => committed,
        Err(error) => return fail_with_exact_cleanup(process, error),
    };
    if committed != CoordinatorMessage::SandboxCommitted {
        return fail_with_exact_cleanup(
            process,
            io::Error::other("sandbox commit acknowledgment was not received"),
        );
    }
    supervise(&control, process)
}

fn runsc_target(arguments: &RuntimeArguments, pid_file: i32) -> io::Result<TargetCommand> {
    let run = arguments
        .runsc_arguments
        .iter()
        .position(|argument| argument == "run")
        .ok_or_else(|| io::Error::other("attached runsc command has no run operation"))?;
    let mut target = TargetCommand::new(arguments.runsc_program.clone());
    for argument in &arguments.runsc_arguments[..=run] {
        target = target.argument(argument);
    }
    target = target.argument(format!("-pid-file=/proc/self/fd/{pid_file}"));
    for argument in &arguments.runsc_arguments[run + 1..] {
        target = target.argument(argument);
    }
    Ok(target.with_parent_death_kill_edge())
}

fn attached_runsc_standard_io(target: &TargetCommand) -> io::Result<ChildStandardIo> {
    if target.parent_edge() != ParentEdge::ParentDeathKill {
        return Err(io::Error::other(
            "attached runsc requires the parent-death KILL edge",
        ));
    }
    Ok(ChildStandardIo::silenced())
}

fn wait_for_runsc_ready(pid_file: &File, process: &RetainedProcess) -> io::Result<u32> {
    let deadline = Instant::now() + READINESS_TIMEOUT;
    let mut bytes = [0_u8; 32];
    loop {
        require_lifetime_held()?;
        if process.has_exited()? {
            return Err(io::Error::other("attached runsc exited before readiness"));
        }
        let length = pid_file.read_at(&mut bytes, 0)?;
        if length > 0 {
            return std::str::from_utf8(&bytes[..length])
                .ok()
                .map(str::trim)
                .and_then(|value| value.parse().ok())
                .filter(|pid| *pid > 0)
                .ok_or_else(|| io::Error::other("runsc pid-file readiness is malformed"));
        }
        if Instant::now() >= deadline {
            return Err(io::Error::other("attached runsc readiness timed out"));
        }
        std::thread::sleep(Duration::from_millis(5));
    }
}

fn supervise(control: &std::io::Stdout, process: RetainedProcess) -> io::Result<()> {
    loop {
        let exited = match process.has_exited() {
            Ok(exited) => exited,
            Err(error) => return fail_with_exact_cleanup(process, error),
        };
        if exited {
            child_broker().reap(process)?;
            send_message(control, &SupervisorEvent::AttachedRunscExited)?;
            return Ok(());
        }
        let lease_lost = match lifetime_lost() {
            Ok(lost) => lost,
            Err(error) => return fail_with_exact_cleanup(process, error),
        };
        if lease_lost {
            terminate_and_reap(process)?;
            return Ok(());
        }
        let mut descriptors = [PollFd::new(control, PollFlags::IN)];
        if let Err(error) = poll(
            &mut descriptors,
            Some(&Timespec {
                tv_sec: 0,
                tv_nsec: 10_000_000,
            }),
        ) {
            return fail_with_exact_cleanup(process, error.into());
        }
        let events = descriptors[0].revents();
        if events.intersects(PollFlags::HUP | PollFlags::ERR) {
            return fail_with_exact_cleanup(
                process,
                io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "runtime control channel closed",
                ),
            );
        }
        if events.intersects(PollFlags::IN) {
            let message = match receive_message(control) {
                Ok(message) => message,
                Err(error) => return fail_with_exact_cleanup(process, error),
            };
            match message {
                CoordinatorMessage::SignalAttachedRunsc(signal) => {
                    if let Err(error) = process.send_signal(signal.signal()) {
                        return fail_with_exact_cleanup(process, error);
                    }
                    if let Err(error) =
                        send_message(control, &SupervisorEvent::SignalApplied(signal))
                    {
                        return fail_with_exact_cleanup(process, error);
                    }
                }
                CoordinatorMessage::Stop => {
                    terminate_and_reap(process)?;
                    send_message(control, &SupervisorEvent::Stopped)?;
                    return Ok(());
                }
                CoordinatorMessage::StagedAck(_) | CoordinatorMessage::SandboxCommitted => {
                    return fail_with_exact_cleanup(
                        process,
                        io::Error::other("unexpected duplicate coordinator message"),
                    );
                }
            }
        }
    }
}

fn terminate_and_reap(process: RetainedProcess) -> io::Result<()> {
    let term_error = process.send_signal(Signal::TERM).err();
    let deadline = Instant::now() + TERMINATION_GRACE;
    while !process.has_exited()? && Instant::now() < deadline {
        std::thread::sleep(Duration::from_millis(5));
    }
    if !process.has_exited()? {
        process.send_signal(Signal::KILL)?;
    }
    child_broker().reap(process)?;
    if let Some(error) = term_error {
        Err(error)
    } else {
        Ok(())
    }
}

fn fail_with_exact_cleanup(process: RetainedProcess, error: io::Error) -> io::Result<()> {
    let cleanup = terminate_and_reap(process);
    match cleanup {
        Ok(()) => Err(error),
        Err(cleanup) => Err(io::Error::new(
            error.kind(),
            format!("{error}; exact attached-runsc cleanup failed: {cleanup}"),
        )),
    }
}

fn require_lifetime_held() -> io::Result<()> {
    if lifetime_lost()? {
        Err(io::Error::other("Host Coordinator lifetime lease is lost"))
    } else {
        Ok(())
    }
}

fn lifetime_lost() -> io::Result<bool> {
    let stderr = std::io::stderr();
    let mut descriptors = [PollFd::new(&stderr, PollFlags::IN)];
    poll(&mut descriptors, Some(&Timespec::default()))?;
    Ok(descriptors[0]
        .revents()
        .intersects(PollFlags::IN | PollFlags::HUP | PollFlags::ERR))
}

fn send_message<T: Serialize>(descriptor: impl AsFd, message: &T) -> io::Result<()> {
    let encoded = encode(message).map_err(|error| io::Error::other(error.to_string()))?;
    if encoded.len() > MAX_CONTROL_MESSAGE {
        return Err(io::Error::other("runtime control message is oversized"));
    }
    let length = u32::try_from(encoded.len())
        .map_err(|_| io::Error::other("runtime control message is oversized"))?
        .to_be_bytes();
    send_all(descriptor.as_fd(), &length)?;
    send_all(descriptor.as_fd(), &encoded)
}

fn receive_message<T>(descriptor: impl AsFd) -> io::Result<T>
where
    T: serde::de::DeserializeOwned + Serialize,
{
    let mut length = [0_u8; 4];
    receive_exact(descriptor.as_fd(), &mut length)?;
    let length = usize::try_from(u32::from_be_bytes(length))
        .map_err(|_| io::Error::other("runtime control length is invalid"))?;
    if length > MAX_CONTROL_MESSAGE {
        return Err(io::Error::other("runtime control message is oversized"));
    }
    let mut encoded = vec![0_u8; length];
    receive_exact(descriptor.as_fd(), &mut encoded)?;
    decode(&encoded).map_err(|error| io::Error::other(error.to_string()))
}

fn send_all(descriptor: impl AsFd + Copy, mut bytes: &[u8]) -> io::Result<()> {
    while !bytes.is_empty() {
        let sent = write(descriptor, bytes)?;
        if sent == 0 {
            return Err(io::Error::new(
                io::ErrorKind::WriteZero,
                "runtime control channel closed",
            ));
        }
        bytes = &bytes[sent..];
    }
    Ok(())
}

fn receive_exact(descriptor: impl AsFd + Copy, mut bytes: &mut [u8]) -> io::Result<()> {
    while !bytes.is_empty() {
        let received = read(descriptor, &mut *bytes)?;
        if received == 0 {
            return Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "runtime control channel closed",
            ));
        }
        bytes = &mut bytes[received..];
    }
    Ok(())
}

fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write as _;

    bytes.iter().fold(String::new(), |mut output, byte| {
        write!(output, "{byte:02x}").expect("writing to a String cannot fail");
        output
    })
}

#[cfg(test)]
mod tests {
    use std::ffi::OsString;

    use lazaret_core::identity::{CoordinatorId, DomainId, GenerationId};
    use lazaret_core::records::HostProcessOwnershipSlot;
    use lazaret_linux::ipc::private_stream_pair;
    use lazaret_linux::process::ChildStandardIoPolicy;
    use rustix::pipe::{PipeFlags, pipe_with};

    use super::*;

    #[test]
    fn private_stream_framing_round_trips_typed_events() {
        let (coordinator, supervisor) = private_stream_pair().unwrap();
        let expected = SupervisorEvent::Refused("refused".into());
        send_message(&supervisor, &expected).unwrap();
        assert_eq!(
            CoordinatorChannel::new(coordinator).receive().unwrap(),
            expected
        );
    }

    #[test]
    fn runtime_supervisor_and_runsc_keep_distinct_edge_and_stdio_policies() {
        let ownership = ProcessOwnershipContext {
            domain_id: DomainId::new([1; 32]),
            generation_id: GenerationId::new(7),
            coordinator_id: CoordinatorId::new([2; 32]),
            slot: HostProcessOwnershipSlot::SandboxRuntimeSupervisor,
        };
        let launch = RuntimeSupervisorLaunch {
            target: TargetCommand::new("/proc/self/exe").with_lifetime_pipe_parent_edge(),
            ownership,
        };
        let (_coordinator_control, supervisor_control) = private_stream_pair().unwrap();
        let (supervisor_lifetime, _coordinator_lifetime) = pipe_with(PipeFlags::CLOEXEC).unwrap();
        assert_eq!(launch.target().parent_edge(), ParentEdge::LifetimePipe);
        assert_eq!(
            launch
                .child_standard_io(supervisor_control, supervisor_lifetime)
                .unwrap()
                .policy(),
            ChildStandardIoPolicy::PrivateControlWithLifetimePipe
        );

        let runsc_arguments = RuntimeArguments {
            ownership: ProcessOwnershipContext {
                slot: HostProcessOwnershipSlot::SandboxAttachedRunsc,
                ..ownership
            },
            runsc_program: "runsc".into(),
            runsc_arguments: ["-root=/runtime", "run", "container"]
                .map(OsString::from)
                .into(),
        };
        let runsc = runsc_target(&runsc_arguments, 9).unwrap();
        assert_eq!(runsc.parent_edge(), ParentEdge::ParentDeathKill);
        assert_eq!(
            attached_runsc_standard_io(&runsc).unwrap().policy(),
            ChildStandardIoPolicy::Silenced
        );
    }

    #[test]
    fn runtime_role_parser_preserves_exact_ownership_and_runsc_arguments() {
        let domain = [1_u8; 32];
        let coordinator = [2_u8; 32];
        let arguments = [
            "--domain".into(),
            hex(&domain).into(),
            "--generation".into(),
            "7".into(),
            "--coordinator".into(),
            hex(&coordinator).into(),
            "--runsc".into(),
            "runsc".into(),
            "-root=/runtime".into(),
            "run".into(),
            "container".into(),
        ];
        let parsed = parse_runtime_arguments(&arguments).unwrap();
        assert_eq!(parsed.ownership.domain_id, DomainId::new(domain));
        assert_eq!(parsed.ownership.generation_id, GenerationId::new(7));
        assert_eq!(
            parsed.ownership.coordinator_id,
            CoordinatorId::new(coordinator)
        );
        assert_eq!(
            parsed.ownership.slot,
            HostProcessOwnershipSlot::SandboxAttachedRunsc
        );
        assert_eq!(parsed.runsc_program, OsString::from("runsc"));
        assert_eq!(
            parsed.runsc_arguments,
            ["-root=/runtime", "run", "container"].map(OsString::from)
        );
    }

    #[test]
    fn runtime_role_parser_refuses_non_runsc_target() {
        let arguments = [
            "--domain".into(),
            hex(&[1_u8; 32]).into(),
            "--generation".into(),
            "7".into(),
            "--coordinator".into(),
            hex(&[2_u8; 32]).into(),
            "--runsc".into(),
            "foreign-runtime".into(),
        ];
        assert!(parse_runtime_arguments(&arguments).is_err());
    }
}