//! 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::DigitallySignedStruct;
use rustls::client::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use zagan_bet::protocol::LineBuffer;
use zagan_bet::tofu::{TrustDecision, decide, spki_fingerprint};

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,
                lines: LineBuffer::new(),
                read_timeout: Some(std::time::Duration::from_secs(10)),
            })),
        })
    }

    /// 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,
    lines: LineBuffer,
    read_timeout: Option<std::time::Duration>,
}

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 s = self.session.lock();
        loop {
            if let Some(line) = s.lines.next_line() {
                buf.push_str(&line);
                buf.push('\n');
                return Ok(line.len() + 1);
            }

            let mut byte_buf = [0u8; 4096];
            match s.tls.reader().read(&mut byte_buf) {
                Ok(0) => {
                    if let Some(rest) = s.lines.take_rest() {
                        let read = rest.len();
                        buf.push_str(&rest);
                        return Ok(read);
                    }
                    return Ok(0);
                }
                Ok(n) => s.lines.push(&byte_buf[..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()
    }

    /// Apply a read timeout to the underlying TCP socket.
    pub fn set_read_timeout(&self, duration: Option<std::time::Duration>) -> std::io::Result<()> {
        let mut s = self.session.lock();
        s.read_timeout = duration;
        s.sock.set_read_timeout(duration)
    }

    /// Return one complete line if the socket already has one, without
    /// blocking. Used to pick up server pushes between requests.
    pub fn read_available_line(&mut self) -> std::io::Result<Option<String>> {
        let mut s = self.session.lock();
        if let Some(line) = s.lines.next_line() {
            return Ok(Some(line));
        }

        // `set_read_timeout(Some(0))` is rejected by std, so poll with the
        // shortest timeout the socket accepts instead of a nonblocking mode.
        s.sock
            .set_read_timeout(Some(std::time::Duration::from_millis(1)))?;
        let drained = (|| -> std::io::Result<()> {
            loop {
                let mut buf = [0u8; 4096];
                match s.tls.reader().read(&mut buf) {
                    Ok(0) => return Ok(()),
                    Ok(n) => s.lines.push(&buf[..n]),
                    // The reader only drains plaintext rustls has already
                    // decrypted; `drive` moves bytes off the socket.
                    Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                        match s.drive() {
                            Ok(()) => continue,
                            Err(error)
                                if matches!(
                                    error.kind(),
                                    std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
                                ) =>
                            {
                                return Ok(());
                            }
                            Err(error) => return Err(error),
                        }
                    }
                    Err(error)
                        if matches!(
                            error.kind(),
                            std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
                        ) =>
                    {
                        return Ok(());
                    }
                    Err(error) => return Err(error),
                }
            }
        })();
        s.sock.set_read_timeout(s.read_timeout)?;
        drained?;
        Ok(s.lines.next_line())
    }
}

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()
    }
}