use std::ffi::OsString;
use std::io::{self, Read, Write};
use std::os::unix::ffi::OsStringExt as _;
use std::path::PathBuf;
use std::process::ExitCode;
use std::time::Duration;

use ed25519_dalek::SigningKey;
use lazaret_core::codec::decode;
use lazaret_core::protocol::VerifiedSession;
use lazaret_core::records::{
    EscalationOutcome, ProtocolPublicationIdentity, SandboxCredentialBundle, SessionBinding,
    SessionRegistrationResponse,
};
use lazaret_escalation_client::{ClientEvent, EscalationClient};
use lazaret_linux::filesystem::RegularFileMailbox;
use lazaret_sandbox_contract::{
    CREDENTIALS_DESTINATION, MAILBOX_REQUESTS_DESTINATION, MAILBOX_RESULTS_DESTINATION,
};
use nix::sys::signal::{SigSet, SigmaskHow, Signal};
use nix::sys::signalfd::{SfdFlags, SignalFd};

struct Arguments {
    inbox: PathBuf,
    outbox: PathBuf,
    publication: PathBuf,
    sandbox_credentials: PathBuf,
    registration: PathBuf,
    session_binding: PathBuf,
    justification: String,
    command_argv: Vec<Vec<u8>>,
}

fn main() -> ExitCode {
    match run() {
        Ok(code) => ExitCode::from(code),
        Err(error) => {
            eprintln!("lazaret escalation refused: {error}");
            ExitCode::from(2)
        }
    }
}

fn run() -> Result<u8, Box<dyn std::error::Error>> {
    let arguments = parse_arguments_with_session(
        std::env::args_os().skip(1).collect(),
        std::env::var_os("LAZARET_SESSION_ID"),
    )?;
    let publication: ProtocolPublicationIdentity =
        read_record(&arguments.publication).map_err(|error| context("publication", error))?;
    let bundle: SandboxCredentialBundle = read_record(&arguments.sandbox_credentials)
        .map_err(|error| context("sandbox credentials", error))?;
    let registration: SessionRegistrationResponse = read_record(&arguments.registration)
        .map_err(|error| context("session registration", error))?;
    let binding: SessionBinding = read_record(&arguments.session_binding)
        .map_err(|error| context("session binding", error))?;
    let sandbox_key_bytes: [u8; 32] = bundle
        .sandbox_private_key
        .as_slice()
        .try_into()
        .map_err(|error| context("sandbox signing key", error))?;
    let sandbox_key = SigningKey::from_bytes(&sandbox_key_bytes);
    let session = VerifiedSession::new(
        publication.trust_manifest_sha256,
        bundle.trust_manifest.tls_ca_sha256,
        &bundle.trust_manifest,
        binding,
        bundle.coordinator_protocol_credential(),
        bundle.sandbox_protocol_credential(),
        registration.session_capability,
    )
    .map_err(|error| context("session authority", error))?;
    let client = EscalationClient::new(
        RegularFileMailbox::open(&arguments.inbox)
            .map_err(|error| context("request mailbox", error))?,
        RegularFileMailbox::open(&arguments.outbox)
            .map_err(|error| context("result mailbox", error))?,
        session,
        sandbox_key,
    );
    let mut cancellation =
        CancellationSignals::register().map_err(|error| context("cancellation signals", error))?;
    let mut request = client
        .request(arguments.command_argv, arguments.justification)
        .map_err(|error| context("request publication", error))?;
    loop {
        if cancellation.requested()? {
            request.cancel()?;
        }
        match request
            .poll()
            .map_err(|error| context("result collection", error))?
        {
            ClientEvent::Output { stream, bytes, .. } => match stream {
                lazaret_core::records::ResultStream::Stdout => {
                    io::stdout().write_all(&bytes)?;
                    io::stdout().flush()?;
                }
                lazaret_core::records::ResultStream::Stderr => {
                    io::stderr().write_all(&bytes)?;
                    io::stderr().flush()?;
                }
            },
            ClientEvent::Terminal(result) => {
                return Ok(match result.outcome {
                    EscalationOutcome::Completed => result
                        .exit_code
                        .and_then(|code| u8::try_from(code).ok())
                        .unwrap_or(1),
                    EscalationOutcome::Denied => 3,
                    EscalationOutcome::TimedOut => 4,
                    EscalationOutcome::Cancelled => 5,
                    EscalationOutcome::Failed => 6,
                });
            }
            ClientEvent::Pending => std::thread::sleep(Duration::from_millis(10)),
        }
    }
}

fn context(label: &str, error: impl std::fmt::Display) -> io::Error {
    io::Error::other(format!("{label}: {error}"))
}

struct CancellationSignals {
    descriptor: SignalFd,
    previous: SigSet,
    published: bool,
}

impl CancellationSignals {
    fn register() -> io::Result<Self> {
        let mut signals = SigSet::empty();
        signals.add(Signal::SIGINT);
        signals.add(Signal::SIGTERM);
        let previous = signals
            .thread_swap_mask(SigmaskHow::SIG_BLOCK)
            .map_err(io::Error::from)?;
        let descriptor =
            SignalFd::with_flags(&signals, SfdFlags::SFD_CLOEXEC | SfdFlags::SFD_NONBLOCK)
                .map_err(io::Error::from)?;
        Ok(Self {
            descriptor,
            previous,
            published: false,
        })
    }

    fn requested(&mut self) -> io::Result<bool> {
        if self.published {
            return Ok(false);
        }
        if self
            .descriptor
            .read_signal()
            .map_err(io::Error::from)?
            .is_some()
        {
            self.published = true;
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

impl Drop for CancellationSignals {
    fn drop(&mut self) {
        let _ = self.previous.thread_set_mask();
    }
}

fn read_record<T>(path: &std::path::Path) -> Result<T, Box<dyn std::error::Error>>
where
    T: lazaret_core::records::SchemaType + serde::de::DeserializeOwned + serde::Serialize,
{
    let mut file = std::fs::File::from(rustix::fs::openat(
        rustix::fs::CWD,
        path,
        rustix::fs::OFlags::RDONLY
            | rustix::fs::OFlags::NONBLOCK
            | rustix::fs::OFlags::NOFOLLOW
            | rustix::fs::OFlags::CLOEXEC,
        rustix::fs::Mode::empty(),
    )?);
    let metadata = file.metadata()?;
    let length = usize::try_from(metadata.len())?;
    if !metadata.file_type().is_file() || length > lazaret_core::codec::MAX_SCHEMA_RECORD_BYTES {
        return Err(io::Error::other("client authority record type or size is invalid").into());
    }
    let mut bytes = Vec::with_capacity(length.saturating_add(1));
    Read::by_ref(&mut file)
        .take((lazaret_core::codec::MAX_SCHEMA_RECORD_BYTES as u64) + 1)
        .read_to_end(&mut bytes)?;
    if bytes.len() != length || file.metadata()?.len() != metadata.len() {
        return Err(io::Error::other("client authority record changed while read").into());
    }
    Ok(decode(&bytes)?)
}

fn parse_arguments_with_session(
    mut arguments: Vec<OsString>,
    session: Option<OsString>,
) -> Result<Arguments, io::Error> {
    let separator = arguments
        .iter()
        .position(|argument| argument == "--")
        .ok_or_else(|| io::Error::other("command separator is missing"))?;
    let command = arguments.split_off(separator + 1);
    arguments.pop();
    if command.is_empty() {
        return Err(io::Error::other("command argv is empty"));
    }
    let mut inbox = None;
    let mut outbox = None;
    let mut publication = None;
    let mut sandbox_credentials = None;
    let mut registration = None;
    let mut session_binding = None;
    let mut justification = None;
    let mut options = arguments.into_iter();
    while let Some(option) = options.next() {
        let value = options
            .next()
            .ok_or_else(|| io::Error::other("client option has no value"))?;
        match option.to_str() {
            Some("--inbox") => inbox = Some(PathBuf::from(value)),
            Some("--outbox") => outbox = Some(PathBuf::from(value)),
            Some("--publication") => publication = Some(PathBuf::from(value)),
            Some("--sandbox-credentials") => sandbox_credentials = Some(PathBuf::from(value)),
            Some("--registration") => registration = Some(PathBuf::from(value)),
            Some("--session-binding") => session_binding = Some(PathBuf::from(value)),
            Some("--justification") => {
                justification = Some(
                    value
                        .into_string()
                        .map_err(|_| io::Error::other("justification is not UTF-8"))?,
                );
            }
            _ => return Err(io::Error::other("unknown client option")),
        }
    }
    if let Some(session) = session {
        let session = session
            .into_string()
            .map_err(|_| io::Error::other("LAZARET_SESSION_ID is not UTF-8"))?;
        if session.len() != 64
            || !session
                .bytes()
                .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
        {
            return Err(io::Error::other("LAZARET_SESSION_ID is invalid"));
        }
        let credentials = PathBuf::from(CREDENTIALS_DESTINATION);
        inbox.get_or_insert_with(|| PathBuf::from(MAILBOX_REQUESTS_DESTINATION));
        outbox.get_or_insert_with(|| PathBuf::from(MAILBOX_RESULTS_DESTINATION));
        publication.get_or_insert_with(|| credentials.join("publication.postcard"));
        sandbox_credentials.get_or_insert_with(|| credentials.join("sandbox-credentials.postcard"));
        registration.get_or_insert_with(|| {
            credentials.join(format!("session-{session}-registration.postcard"))
        });
        session_binding
            .get_or_insert_with(|| credentials.join(format!("session-{session}-binding.postcard")));
    }
    Ok(Arguments {
        inbox: inbox.ok_or_else(|| io::Error::other("inbox path is missing"))?,
        outbox: outbox.ok_or_else(|| io::Error::other("outbox path is missing"))?,
        publication: publication.ok_or_else(|| io::Error::other("publication path is missing"))?,
        sandbox_credentials: sandbox_credentials
            .ok_or_else(|| io::Error::other("sandbox credentials path is missing"))?,
        registration: registration
            .ok_or_else(|| io::Error::other("registration path is missing"))?,
        session_binding: session_binding
            .ok_or_else(|| io::Error::other("session binding path is missing"))?,
        justification: justification.ok_or_else(|| io::Error::other("justification is missing"))?,
        command_argv: command.into_iter().map(OsString::into_vec).collect(),
    })
}

#[cfg(test)]
mod tests {
    use super::parse_arguments_with_session;
    use lazaret_sandbox_contract::{
        CREDENTIALS_DESTINATION, MAILBOX_REQUESTS_DESTINATION, MAILBOX_RESULTS_DESTINATION,
    };
    use std::ffi::OsString;

    fn valid_arguments() -> Vec<OsString> {
        [
            "--inbox",
            "/inbox",
            "--outbox",
            "/outbox",
            "--publication",
            "/publication",
            "--sandbox-credentials",
            "/credentials",
            "--registration",
            "/registration",
            "--session-binding",
            "/binding",
            "--justification",
            "needed",
            "--",
            "/usr/bin/tool",
            "$HOME",
            "a;b",
        ]
        .into_iter()
        .map(OsString::from)
        .collect()
    }

    #[test]
    fn parser_preserves_exact_argv_without_shell_interpretation() {
        let parsed = parse_arguments_with_session(valid_arguments(), None).unwrap();
        assert_eq!(
            parsed.command_argv,
            [
                b"/usr/bin/tool".to_vec(),
                b"$HOME".to_vec(),
                b"a;b".to_vec()
            ]
        );
    }

    #[test]
    fn parser_refuses_missing_separator_unknown_options_and_empty_argv() {
        let mut missing_separator = valid_arguments();
        missing_separator.retain(|argument| argument != "--");
        assert!(parse_arguments_with_session(missing_separator, None).is_err());

        let mut unknown = valid_arguments();
        unknown[0] = OsString::from("--unknown");
        assert!(parse_arguments_with_session(unknown, None).is_err());

        let mut empty = valid_arguments();
        empty.truncate(empty.iter().position(|argument| argument == "--").unwrap() + 1);
        assert!(parse_arguments_with_session(empty, None).is_err());
    }

    #[test]
    fn session_environment_supplies_only_fixed_generation_mount_paths() {
        let arguments = ["--justification", "needed", "--", "/usr/bin/tool"]
            .into_iter()
            .map(OsString::from)
            .collect();
        let session = "11".repeat(32);
        let parsed =
            parse_arguments_with_session(arguments, Some(OsString::from(&session))).unwrap();
        assert_eq!(
            parsed.inbox,
            std::path::Path::new(MAILBOX_REQUESTS_DESTINATION)
        );
        assert_eq!(
            parsed.outbox,
            std::path::Path::new(MAILBOX_RESULTS_DESTINATION)
        );
        assert_eq!(
            parsed.registration,
            std::path::Path::new(
                &std::path::Path::new(CREDENTIALS_DESTINATION)
                    .join(format!("session-{session}-registration.postcard"))
            )
        );

        let arguments = ["--justification", "needed", "--", "/usr/bin/tool"]
            .into_iter()
            .map(OsString::from)
            .collect();
        assert!(
            parse_arguments_with_session(arguments, Some(OsString::from("not-an-identity")))
                .is_err()
        );
    }
}