//! Client transport: TCP, rustls, TOFU and the HELLO handshake.
//!
//! The handshake records the server's SPKI fingerprint but does not act on it;
//! the trust decision happens after the handshake and before any credentials
//! are sent, so a frontend can ask the player. A changed key is a hard error.

use std::fmt;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::mpsc;
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

use parking_lot::Mutex;
use rustls::client::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::DigitallySignedStruct;
use rusqlite::Connection;

use crate::client::database;
use crate::protocol::LineBuffer;
use crate::tofu::{decide, spki_fingerprint, TrustDecision};
use crate::version::{parse_hello, HelloReply, VERSION};

/// How long to wait for TCP + TLS before giving up.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
/// Read timeout once the stream is usable.
const READ_TIMEOUT: Duration = Duration::from_secs(10);

/// What a frontend shows when a server key is unknown.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TrustPrompt {
    pub addr: String,
    pub fingerprint: String,
}

/// A usable connection after trust, version and optional auto-login.
pub struct Connected {
    pub stream: TlsStream,
    pub user: Option<String>,
    /// Set when stored credentials were rejected; the session is still usable.
    pub auth_error: Option<String>,
    pub server_version: String,
    pub latest: String,
    pub url: String,
    /// Round trip of the HELLO exchange.
    pub elapsed: Duration,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TransportError {
    Timeout,
    TrustDeclined,
    KeyChanged {
        expected: String,
        observed: String,
    },
    VersionMismatch(String),
    Closed,
    Unexpected(String),
    Io(String),
}

impl fmt::Display for TransportError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TransportError::Timeout => write!(f, "Connection timed out"),
            TransportError::TrustDeclined => write!(f, "Certificate was not trusted"),
            TransportError::KeyChanged { expected, observed } => write!(
                f,
                "HOST KEY CHANGED for this server!\n\
                 expected fingerprint: {expected}\n\
                 observed fingerprint: {observed}\n\
                 If this is intentional (server reinstall), forget this host to reset trust."
            ),
            TransportError::VersionMismatch(message) => write!(f, "{message}"),
            TransportError::Closed => write!(f, "Server closed the connection during HELLO"),
            TransportError::Unexpected(message) => write!(f, "{message}"),
            TransportError::Io(message) => write!(f, "Connection error: {message}"),
        }
    }
}

impl std::error::Error for TransportError {}

fn io_failure(error: std::io::Error) -> TransportError {
    TransportError::Io(error.to_string())
}

/// Connect, verify the key, check versions and try the stored login.
///
/// `trust` is called only for a server we have never seen; returning `false`
/// aborts before any application data is sent.
pub fn connect(
    addr: &str,
    conn: &Connection,
    unsafe_mode: bool,
    trust: &mut dyn FnMut(&TrustPrompt) -> bool,
) -> Result<Connected, TransportError> {
    let stored = database::get_fingerprint(conn, addr);
    let (tls, observed) = tls_connect(addr)?;
    let observed = observed.ok_or_else(|| {
        TransportError::Unexpected("the server presented no certificate fingerprint".to_string())
    })?;

    if !unsafe_mode {
        match decide(stored.as_deref(), &observed) {
            TrustDecision::Trusted => {}
            TrustDecision::TrustNew => {
                let prompt = TrustPrompt {
                    addr: addr.to_string(),
                    fingerprint: observed.clone(),
                };
                if !trust(&prompt) {
                    return Err(TransportError::TrustDeclined);
                }
                database::save_fingerprint(conn, addr, &observed)
                    .map_err(|error| TransportError::Io(error.to_string()))?;
            }
            TrustDecision::Mismatch { expected, actual } => {
                return Err(TransportError::KeyChanged {
                    expected,
                    observed: actual,
                })
            }
        }
    }

    let mut tls = tls;
    tls.set_read_timeout(Some(READ_TIMEOUT))
        .map_err(io_failure)?;

    let start = Instant::now();
    tls.write_line(&format!("HELLO {VERSION}"))
        .map_err(io_failure)?;
    let mut response = String::new();
    if tls.read_line(&mut response).map_err(io_failure)? == 0 {
        return Err(TransportError::Closed);
    }
    let elapsed = start.elapsed();

    let response = response.trim();
    if let Some(message) = response.strip_prefix("ERROR") {
        return Err(TransportError::VersionMismatch(message.trim().to_string()));
    }
    let hello: HelloReply =
        parse_hello(response).ok_or_else(|| TransportError::Unexpected(response.to_string()))?;

    // Credentials are keyed by the pinned fingerprint, so any address that
    // presents this key can log in with the same stored account.
    let mut user = None;
    let mut auth_error = None;
    if let Some((username, password)) = database::get_credentials(conn, &observed) {
        tls.write_line(&format!("AUTH {username} {password}"))
            .map_err(io_failure)?;
        let mut auth = String::new();
        tls.read_line(&mut auth).map_err(io_failure)?;
        if auth.trim().starts_with("OK") {
            user = Some(username);
        } else {
            auth_error = Some(auth.trim().to_string());
        }
    }

    Ok(Connected {
        stream: tls,
        user,
        auth_error,
        server_version: hello.version,
        latest: hello.latest,
        url: hello.url,
        elapsed,
    })
}

/// TCP + TLS on a thread so a blocked handshake can time out.
fn tls_connect(addr: &str) -> Result<(TlsStream, Option<String>), TransportError> {
    let (sender, receiver) = mpsc::channel();
    let addr = addr.to_string();
    let server_name = addr
        .split(':')
        .next()
        .unwrap_or("127.0.0.1")
        .to_string();

    thread::spawn(move || {
        let outcome = (|| -> Result<(TlsStream, Option<String>), TransportError> {
            let tcp = TcpStream::connect(&addr).map_err(io_failure)?;
            let connector = TlsConnector::new(&server_name)
                .map_err(|error| TransportError::Io(error.to_string()))?;
            let stream = connector
                .connect(tcp)
                .map_err(|error| TransportError::Io(error.to_string()))?;
            Ok((stream, connector.observed_fingerprint()))
        })();
        let _ = sender.send(outcome);
    });

    match receiver.recv_timeout(CONNECT_TIMEOUT) {
        Ok(outcome) => outcome,
        Err(mpsc::RecvTimeoutError::Timeout) => Err(TransportError::Timeout),
        Err(mpsc::RecvTimeoutError::Disconnected) => Err(TransportError::Unexpected(
            "connection attempt failed".to_string(),
        )),
    }
}

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

impl TlsConnector {
    fn new(server_name: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let observed = Arc::new(Mutex::new(None));
        let verifier = Arc::new(RecordingVerifier {
            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,
        })
    }

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

        Ok(TlsStream {
            session: Arc::new(Mutex::new(Session {
                sock,
                tls: connection,
                lines: LineBuffer::new(),
                read_timeout: Some(READ_TIMEOUT),
            })),
        })
    }

    fn observed_fingerprint(&self) -> Option<String> {
        self.observed.lock().clone()
    }
}

/// Certificate verification is record-only: the trust decision is made after
/// the handshake, before anything is sent, so a frontend can ask first.
struct RecordingVerifier {
    observed: Arc<Mutex<Option<String>>>,
}

impl ServerCertVerifier for RecordingVerifier {
    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 fingerprint = spki_fingerprint(&end_entity.0)
            .map_err(|error| rustls::Error::General(format!("certificate parse failed: {error}")))?;
        *self.observed.lock() = Some(fingerprint);
        Ok(ServerCertVerified::assertion())
    }

    // Identity is pinned by the fingerprint check above; handshake signatures
    // are not additionally validated 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<Duration>,
}

impl Session {
    fn drive(&mut self) -> std::io::Result<()> {
        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 chunk = [0u8; 4096];
            match s.tls.reader().read(&mut chunk) {
                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(&chunk[..n]),
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                Err(error) => return Err(error),
            }
        }
    }

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

    pub fn set_read_timeout(&self, duration: Option<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(error) if error.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                Err(error) => return Err(error),
            }
        }
    }
}

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