#![doc = "Dependency-neutral client side of the Lazaret control protocol."]

use std::fmt;
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::fd::{AsFd, AsRawFd};
use std::os::unix::fs::{FileTypeExt, MetadataExt};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use lazaret_core::codec::{CodecError, decode, encode};
use lazaret_core::configuration::ConfinedPathPlan;
use lazaret_core::identity::{CertificateDigest, ClientId};
use lazaret_core::protocol::{
    ProtocolValidationError, client_enrollment_binding, verify_tls_credential,
};
use lazaret_core::records::{
    ClientEnrollmentRequest, GenerationIdentity, ProtocolPublicationIdentity, ProtocolVersion,
    TlsCredential, TlsRole,
};
use rcgen::{
    CertificateParams, DistinguishedName, DnType, ExtendedKeyUsagePurpose, KeyPair,
    KeyUsagePurpose, PKCS_ED25519,
};
use rustix::fs::{Mode, OFlags, openat};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName};
use rustls::{ClientConfig, ClientConnection, RootCertStore};
use sha2::{Digest, Sha256};

pub const CONTROL_ALPN: &[u8] = b"lazaret-control/1";
pub const CONTROL_SERVER_NAME: &str = "coordinator.lazaret.invalid";
pub const CONTROL_SOCKET_NAME: &str = "c.sock";
pub const PUBLICATION_IDENTITY_NAME: &str = "protocol-identity.postcard";
pub const MAX_CONTROL_BYTES: usize = 65_536;

#[derive(Debug)]
pub enum ControlClientError {
    Io(io::Error),
    Codec(CodecError),
    Certificate(rcgen::Error),
    Tls(rustls::Error),
    Protocol(ProtocolValidationError),
    InvalidPath(&'static str),
    InvalidPublication(&'static str),
    InvalidCertificateProfile,
    OversizeFrame { length: usize, maximum: usize },
}

impl fmt::Display for ControlClientError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Io(error) => write!(formatter, "control I/O failed: {error}"),
            Self::Codec(error) => write!(formatter, "control record failed validation: {error}"),
            Self::Certificate(error) => write!(formatter, "certificate operation failed: {error}"),
            Self::Tls(error) => write!(formatter, "TLS failed: {error}"),
            Self::Protocol(error) => write!(formatter, "protocol validation failed: {error}"),
            Self::InvalidPath(reason) => write!(formatter, "invalid control path: {reason}"),
            Self::InvalidPublication(reason) => {
                write!(formatter, "invalid control publication: {reason}")
            }
            Self::InvalidCertificateProfile => {
                formatter.write_str("TLS certificate profile is invalid")
            }
            Self::OversizeFrame { length, maximum } => {
                write!(formatter, "control frame length {length} exceeds {maximum}")
            }
        }
    }
}

impl std::error::Error for ControlClientError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Io(error) => Some(error),
            Self::Codec(error) => Some(error),
            Self::Certificate(error) => Some(error),
            Self::Tls(error) => Some(error),
            Self::Protocol(error) => Some(error),
            Self::InvalidPath(_)
            | Self::InvalidPublication(_)
            | Self::InvalidCertificateProfile
            | Self::OversizeFrame { .. } => None,
        }
    }
}

impl From<io::Error> for ControlClientError {
    fn from(error: io::Error) -> Self {
        Self::Io(error)
    }
}

impl From<CodecError> for ControlClientError {
    fn from(error: CodecError) -> Self {
        Self::Codec(error)
    }
}

impl From<rcgen::Error> for ControlClientError {
    fn from(error: rcgen::Error) -> Self {
        Self::Certificate(error)
    }
}

impl From<rustls::Error> for ControlClientError {
    fn from(error: rustls::Error) -> Self {
        Self::Tls(error)
    }
}

impl From<ProtocolValidationError> for ControlClientError {
    fn from(error: ProtocolValidationError) -> Self {
        Self::Protocol(error)
    }
}

pub struct PendingClientEnrollment {
    request: ClientEnrollmentRequest,
    private_key: PrivatePkcs8KeyDer<'static>,
}

impl PendingClientEnrollment {
    /// Creates a CSR bound to the exact generation and Client CLI identity.
    ///
    /// # Errors
    /// Returns an error when encoding, key generation, or CSR signing fails.
    pub fn create(
        generation_identity: GenerationIdentity,
        cli_id: ClientId,
    ) -> Result<Self, ControlClientError> {
        let binding = client_enrollment_binding(&generation_identity, cli_id)?;
        let mut params = CertificateParams::default();
        params.distinguished_name = DistinguishedName::new();
        params.distinguished_name.push(DnType::CommonName, binding);
        params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
        params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ClientAuth];
        let key = KeyPair::generate_for(&PKCS_ED25519)?;
        let csr = params.serialize_request(&key)?;
        Ok(Self {
            request: ClientEnrollmentRequest {
                protocol_version: ProtocolVersion::V1,
                generation_identity,
                cli_id,
                certificate_signing_request_der: csr.der().to_vec(),
            },
            private_key: PrivatePkcs8KeyDer::from(key.serialize_der()),
        })
    }

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

    #[must_use]
    pub fn private_key(&self) -> PrivateKeyDer<'static> {
        PrivateKeyDer::Pkcs8(self.private_key.clone_key())
    }
}

pub struct ControlClient;

impl ControlClient {
    /// Discovers a publication through retained Runtime Domain authority.
    ///
    /// # Errors
    /// Refuses changed authority and malformed or non-private publication objects.
    pub fn discover_confined(
        runtime_domain_root: &ConfinedPathPlan,
        uid: u32,
    ) -> Result<(File, PathBuf, ProtocolPublicationIdentity), ControlClientError> {
        let root =
            File::from(runtime_domain_root.open_authority().map_err(|_| {
                ControlClientError::InvalidPath("runtime Domain authority changed")
            })?);
        let root_metadata = root.metadata()?;
        if !root_metadata.is_dir()
            || root_metadata.uid() != uid
            || root_metadata.mode() & 0o7777 != 0o700
        {
            return Err(ControlClientError::InvalidPath(
                "runtime Domain root is not invoking-user-owned mode 0700",
            ));
        }
        let mut identity = File::from(
            openat(
                root.as_fd(),
                PUBLICATION_IDENTITY_NAME,
                OFlags::RDONLY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
            )
            .map_err(io::Error::from)?,
        );
        let metadata = identity.metadata()?;
        if !metadata.file_type().is_file()
            || metadata.uid() != uid
            || metadata.mode() & 0o7777 != 0o600
            || metadata.nlink() != 1
            || metadata.len() > MAX_CONTROL_BYTES as u64
        {
            return Err(ControlClientError::InvalidPublication(
                "publication identity is not an exact private regular file",
            ));
        }
        let socket = File::from(
            openat(
                root.as_fd(),
                CONTROL_SOCKET_NAME,
                OFlags::PATH | OFlags::CLOEXEC | OFlags::NOFOLLOW,
                Mode::empty(),
            )
            .map_err(io::Error::from)?,
        );
        let metadata = socket.metadata()?;
        if !metadata.file_type().is_socket()
            || metadata.uid() != uid
            || metadata.mode() & 0o7777 != 0o600
        {
            return Err(ControlClientError::InvalidPublication(
                "control socket is not an invoking-user-private socket",
            ));
        }
        let mut bytes = Vec::new();
        identity.read_to_end(&mut bytes)?;
        let socket_path = PathBuf::from(format!(
            "/proc/self/fd/{}/{}",
            root.as_raw_fd(),
            CONTROL_SOCKET_NAME
        ));
        Ok((root, socket_path, decode(&bytes)?))
    }

    /// Builds a TLS 1.3 client configuration pinned to the publication CA.
    ///
    /// # Errors
    /// Refuses malformed CA or client credential material.
    pub fn config(
        publication: &ProtocolPublicationIdentity,
        client: Option<(&TlsCredential, PrivateKeyDer<'static>)>,
    ) -> Result<ClientConfig, ControlClientError> {
        let mut roots = RootCertStore::empty();
        roots.add(CertificateDer::from(
            publication.tls_ca_certificate_der.clone(),
        ))?;
        let builder = ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
            .with_root_certificates(roots);
        let mut config = if let Some((credential, private_key)) = client {
            verify_tls_credential(
                credential,
                &publication.generation_identity,
                TlsRole::ClientCli,
                credential.subject_id(),
            )?;
            builder.with_client_auth_cert(
                vec![CertificateDer::from(credential.certificate_der.clone())],
                private_key,
            )?
        } else {
            builder.with_no_client_auth()
        };
        config.alpn_protocols = vec![CONTROL_ALPN.to_vec()];
        config.enable_early_data = false;
        config.resumption = rustls::client::Resumption::disabled();
        Ok(config)
    }

    /// Connects a pinned TLS stream to the control socket.
    ///
    /// # Errors
    /// Returns socket, TLS, or certificate-profile failures.
    pub fn connect(
        socket_path: &Path,
        config: ClientConfig,
        publication: &ProtocolPublicationIdentity,
    ) -> Result<
        rustls::StreamOwned<ClientConnection, std::os::unix::net::UnixStream>,
        ControlClientError,
    > {
        let socket = std::os::unix::net::UnixStream::connect(socket_path)?;
        let server_name = ServerName::try_from(CONTROL_SERVER_NAME)
            .map_err(|_| ControlClientError::InvalidCertificateProfile)?;
        let mut connection = ClientConnection::new(Arc::new(config), server_name)?;
        let mut socket = socket;
        while connection.is_handshaking() {
            connection.complete_io(&mut socket)?;
        }
        let [certificate] = connection.peer_certificates().unwrap_or_default() else {
            return Err(ControlClientError::InvalidCertificateProfile);
        };
        if connection.protocol_version() != Some(rustls::ProtocolVersion::TLSv1_3)
            || connection.alpn_protocol() != Some(CONTROL_ALPN)
            || certificate_digest(certificate) != publication.coordinator_tls_certificate_sha256
        {
            return Err(ControlClientError::InvalidCertificateProfile);
        }
        Ok(rustls::StreamOwned::new(connection, socket))
    }

    /// Writes one bounded length-delimited control record.
    ///
    /// # Errors
    /// Returns encoding, size, or stream failures.
    pub fn send<T: serde::Serialize>(
        stream: &mut impl Write,
        record: &T,
    ) -> Result<(), ControlClientError> {
        let bytes = encode(record)?;
        if bytes.len() > MAX_CONTROL_BYTES {
            return Err(ControlClientError::OversizeFrame {
                length: bytes.len(),
                maximum: MAX_CONTROL_BYTES,
            });
        }
        let length = u32::try_from(bytes.len()).map_err(|_| ControlClientError::OversizeFrame {
            length: bytes.len(),
            maximum: MAX_CONTROL_BYTES,
        })?;
        stream.write_all(&length.to_be_bytes())?;
        stream.write_all(&bytes)?;
        stream.flush()?;
        Ok(())
    }

    /// Reads one bounded length-delimited control record.
    ///
    /// # Errors
    /// Returns framing, size, decoding, or stream failures.
    pub fn receive<T>(stream: &mut impl Read) -> Result<T, ControlClientError>
    where
        T: serde::de::DeserializeOwned + serde::Serialize,
    {
        let mut length = [0_u8; 4];
        stream.read_exact(&mut length)?;
        let length = u32::from_be_bytes(length) as usize;
        if length > MAX_CONTROL_BYTES {
            return Err(ControlClientError::OversizeFrame {
                length,
                maximum: MAX_CONTROL_BYTES,
            });
        }
        let mut bytes = vec![0_u8; length];
        stream.read_exact(&mut bytes)?;
        Ok(decode(&bytes)?)
    }
}

fn certificate_digest(certificate: impl AsRef<[u8]>) -> CertificateDigest {
    CertificateDigest::new(sha256(certificate.as_ref()))
}

fn sha256(bytes: &[u8]) -> [u8; 32] {
    Sha256::digest(bytes).into()
}

#[cfg(test)]
mod tests {
    use lazaret_core::identity::{
        ClientId, ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId, SandboxId,
    };
    use lazaret_core::records::{ControlRequest, GenerationIdentity};

    use super::*;

    #[test]
    fn enrollment_request_retains_exact_generation_and_client_identity() {
        let generation = GenerationIdentity {
            agent_profile: "profile".to_owned(),
            domain_id: DomainId::new([1; 32]),
            generation_id: GenerationId::new(2),
            configuration_fingerprint: ConfigurationFingerprint::new([3; 32]),
            coordinator_id: CoordinatorId::new([4; 32]),
            sandbox_id: SandboxId::new([5; 32]),
        };
        let cli_id = ClientId::new([6; 32]);
        let pending = PendingClientEnrollment::create(generation.clone(), cli_id).unwrap();

        assert_eq!(pending.request().generation_identity, generation);
        assert_eq!(pending.request().cli_id, cli_id);
        assert!(!pending.request().certificate_signing_request_der.is_empty());
        assert!(matches!(pending.private_key(), PrivateKeyDer::Pkcs8(_)));
    }

    #[test]
    fn control_framing_round_trips_and_rejects_oversize_length() {
        let request = ControlRequest::Enroll(
            PendingClientEnrollment::create(
                GenerationIdentity {
                    agent_profile: "profile".to_owned(),
                    domain_id: DomainId::new([7; 32]),
                    generation_id: GenerationId::new(8),
                    configuration_fingerprint: ConfigurationFingerprint::new([9; 32]),
                    coordinator_id: CoordinatorId::new([10; 32]),
                    sandbox_id: SandboxId::new([11; 32]),
                },
                ClientId::new([12; 32]),
            )
            .unwrap()
            .request()
            .clone(),
        );
        let mut bytes = Vec::new();
        ControlClient::send(&mut bytes, &request).unwrap();
        assert_eq!(
            ControlClient::receive::<ControlRequest>(&mut bytes.as_slice()).unwrap(),
            request
        );

        let oversized = u32::try_from(MAX_CONTROL_BYTES + 1).unwrap().to_be_bytes();
        assert!(ControlClient::receive::<ControlRequest>(&mut oversized.as_slice()).is_err());
    }
}