//! rustls-based client TLS with Trust-On-First-Use (TOFU) host verification.
//!
//! Model (SSH known_hosts style):
//! - First contact: accept the certificate and record the SHA-256 of its
//!   SubjectPublicKeyInfo (SPKI). Hashing the SPKI rather than the whole
//!   certificate means routine renewal with the same key does not break
//!   clients; key rotation does — by design.
//! - Later contacts: the presented key must hash to the stored fingerprint,
//!   otherwise the handshake fails with a loud error.
//!
//! Handshake signatures are accepted without webpki validation; identity is
//! established purely by the pinned fingerprint. An impostor cannot present a
//! different certificate without triggering the mismatch rejection.

use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::Arc;

use parking_lot::Mutex;
use rustls::client::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::DigitallySignedStruct;
use zagan_bet::tofu::{decide, spki_fingerprint, TrustDecision};

pub struct TlsConnector {
    config: Arc<rustls::ClientConfig>,
    server_name: rustls::ServerName,
    observed: Arc<Mutex<Option<String>>>,
}

impl TlsConnector {
    /// `stored` is the previously trusted fingerprint for this host, if any.
    /// With `bypass_tofu` (unsafe mode) any certificate is accepted.
    pub fn new(
        server_name: &str,
        stored: Option<String>,
        bypass_tofu: bool,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let observed = Arc::new(Mutex::new(None));
        let verifier = Arc::new(TofuVerifier {
            stored,
            bypass_tofu,
            observed: Arc::clone(&observed),
        });
        let config = rustls::ClientConfig::builder()
            .with_safe_defaults()
            .with_custom_certificate_verifier(verifier)
            .with_no_client_auth();
        Ok(Self {
            config: Arc::new(config),
            server_name: server_name.try_into()?,
            observed,
        })
    }

    pub fn connect(&self, stream: TcpStream) -> Result<TlsStream, Box<dyn std::error::Error + Send + Sync>> {
        let mut conn =
            rustls::ClientConnection::new(Arc::clone(&self.config), self.server_name.clone())?;
        let mut sock = stream;
        conn.complete_io(&mut sock)?;

        Ok(TlsStream {
            session: Arc::new(Mutex::new(Session { sock, tls: conn })),
        })
    }

    /// Fingerprint of the certificate seen during the last `connect`.
    pub fn observed_fingerprint(&self) -> Option<String> {
        self.observed.lock().clone()
    }
}

struct TofuVerifier {
    stored: Option<String>,
    bypass_tofu: bool,
    observed: Arc<Mutex<Option<String>>>,
}

impl ServerCertVerifier for TofuVerifier {
    fn verify_server_cert(
        &self,
        end_entity: &rustls::Certificate,
        _intermediates: &[rustls::Certificate],
        _server_name: &rustls::ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp_response: &[u8],
        _now: std::time::SystemTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        let fp = spki_fingerprint(&end_entity.0)
            .map_err(|e| rustls::Error::General(format!("certificate parse failed: {e}")))?;
        *self.observed.lock() = Some(fp.clone());

        if self.bypass_tofu {
            return Ok(ServerCertVerified::assertion());
        }

        match decide(self.stored.as_deref(), &fp) {
            TrustDecision::TrustNew | TrustDecision::Trusted => Ok(ServerCertVerified::assertion()),
            TrustDecision::Mismatch { expected, actual } => Err(rustls::Error::General(format!(
                "HOST KEY CHANGED for this server!\n\
                 expected fingerprint: {expected}\n\
                 observed fingerprint: {actual}\n\
                 If this is intentional (server reinstall), use the client's\n\
                 forget-host command to reset trust for this address."
            ))),
        }
    }

    // Identity is pinned via fingerprint comparison above; handshake
    // signatures are not additionally verified against a trust store.
    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::Certificate,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::Certificate,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        vec![
            rustls::SignatureScheme::RSA_PKCS1_SHA256,
            rustls::SignatureScheme::RSA_PKCS1_SHA384,
            rustls::SignatureScheme::RSA_PKCS1_SHA512,
            rustls::SignatureScheme::RSA_PSS_SHA256,
            rustls::SignatureScheme::RSA_PSS_SHA384,
            rustls::SignatureScheme::RSA_PSS_SHA512,
            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
            rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
            rustls::SignatureScheme::ECDSA_NISTP521_SHA512,
        ]
    }
}

/// TCP socket + TLS session behind one mutex so clones share the stream.
struct Session {
    sock: TcpStream,
    tls: rustls::ClientConnection,
}

impl Session {
    fn drive(&mut self) -> std::io::Result<()> {
        Ok(self.tls.complete_io(&mut self.sock).map(|_| ())?)
    }
}

#[derive(Clone)]
pub struct TlsStream {
    session: Arc<Mutex<Session>>,
}

impl TlsStream {
    pub fn read_line(&mut self, buf: &mut String) -> std::io::Result<usize> {
        let mut byte_buf = [0u8; 1024];
        let mut s = self.session.lock();
        loop {
            match s.tls.reader().read(&mut byte_buf) {
                Ok(0) => return Ok(0),
                Ok(n) => {
                    buf.push_str(std::str::from_utf8(&byte_buf[..n]).unwrap_or(""));
                    if buf.contains('\n') {
                        return Ok(n);
                    }
                }
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                Err(e) => return Err(e),
            }
        }
    }

    pub fn write_line(&mut self, msg: &str) -> std::io::Result<()> {
        let mut s = self.session.lock();
        s.tls.writer().write_all(msg.as_bytes())?;
        s.tls.writer().write_all(b"\n")?;
        s.drive()
    }

    /// Clones share the same underlying TLS session.
    pub fn try_clone(&self) -> std::io::Result<TlsStream> {
        Ok(self.clone())
    }
}

impl Read for TlsStream {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let mut s = self.session.lock();
        loop {
            match s.tls.reader().read(buf) {
                Ok(n) => return Ok(n),
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                Err(e) => return Err(e),
            }
        }
    }
}

impl Write for TlsStream {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let mut s = self.session.lock();
        s.tls.writer().write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        let mut s = self.session.lock();
        s.tls.writer().flush()?;
        s.drive()
    }
}