//! SFTP backend, over libssh2.
//!
//! URL form: `sftp://user@host:port/path` (an absolute remote path;
//! `ssh://` is accepted as an alias). A URL without a user takes it
//! from a matching lftp bookmark (see [`crate::lftp`]), then `$USER`.
//! Authentication is tried in order: the key file named by
//! `$BEEPING_SSH_KEY` (if set), the ssh-agent, the standard key files
//! (`~/.ssh/id_ed25519`, `~/.ssh/id_rsa`), `$BEEPING_SSH_PASSWORD`
//! (if set), and finally a password from a matching lftp bookmark.
//!
//! Host keys are checked against `~/.ssh/known_hosts` (or the file
//! named by `$BEEPING_SSH_KNOWN_HOSTS`): a mismatch is a hard error
//! (possible man-in-the-middle), and an unknown host is trusted on
//! first use and appended to the file.
//!
//! At most `?connections=N` (default 4) SSH sessions are held at once,
//! staying comfortably under a typical sshd's session limits.

use std::{
    io::Read as _,
    io::Write as _,
    net::TcpStream,
    path::PathBuf,
    sync::atomic::{AtomicBool, Ordering},
};

use repository::{Backend, ChunkId, Error, ObjectKey, ObjectKind};
use ssh2::{CheckResult, ErrorCode, KnownHostFileKind, RenameFlags, Session, Sftp};
use url::Url;

use crate::{
    lftp, percent_decode,
    pool::{Pool, temp_suffix},
};

/// SFTP status code for a missing file, per the protocol.
const SFTP_NO_SUCH_FILE: i32 = 2;

const DEFAULT_CONNECTIONS: usize = 4;

pub struct SftpBackend {
    host: String,
    port: u16,
    user: String,
    /// Password found in an lftp bookmark, if any; the last resort of
    /// the authentication order.
    bookmark_password: Option<String>,
    root: PathBuf,
    pool: Pool<SftpConnection>,
    /// Set once the server has declined to report free space, so a
    /// backup watching its floor stops asking. The extension is
    /// optional, and a refusal costs a pooled retry and a fresh dial
    /// every time it is tried.
    space_unknown: AtomicBool,
}

/// One authenticated session plus its SFTP channel. The session must
/// outlive the channel, so they travel together.
struct SftpConnection {
    _session: Session,
    sftp: Sftp,
}

impl SftpBackend {
    pub fn from_url(url: &Url) -> Result<SftpBackend, Error> {
        let host = url
            .host_str()
            .ok_or_else(|| Error::Backend(format!("{url} has no host")))?
            .to_string();

        let port = url.port().unwrap_or(22);

        let mut user = match url.username() {
            "" => None,
            user => Some(percent_decode(user)),
        };

        let mut bookmark_password = None;

        if let Some(found) = lftp::find_credentials("sftp", &host, port, user.as_deref()) {
            user.get_or_insert(found.user);
            bookmark_password = found.password;
        }

        let user = match user.or_else(|| std::env::var("USER").ok()) {
            Some(user) => user,
            None => return Err(Error::Backend(format!("{url} has no user name"))),
        };

        Ok(SftpBackend {
            host,
            port,
            user,
            bookmark_password,
            root: PathBuf::from(url.path()),
            pool: Pool::new(crate::connection_limit(url, DEFAULT_CONNECTIONS)?),
            space_unknown: AtomicBool::new(false),
        })
    }

    fn with_connection<T>(
        &self,
        op: impl FnMut(&mut SftpConnection) -> Result<T, Error>,
    ) -> Result<T, Error> {
        // SSH has no clean "server full" signal to classify; the
        // conservative cap does the work here
        self.pool.with(|| self.connect(), |_| false, op)
    }

    fn connect(&self) -> Result<SftpConnection, Error> {
        let tcp = TcpStream::connect((self.host.as_str(), self.port))
            .map_err(|err| Error::Backend(format!("sftp: connecting: {err}")))?;

        let mut session = Session::new().map_err(sftp_error)?;
        session.set_tcp_stream(tcp);
        session.handshake().map_err(sftp_error)?;

        self.check_host_key(&session)?;
        self.authenticate(&session)?;

        let sftp = session.sftp().map_err(sftp_error)?;

        Ok(SftpConnection {
            _session: session,
            sftp,
        })
    }

    /// Verifies the server against `~/.ssh/known_hosts`, trusting and
    /// recording hosts seen for the first time.
    fn check_host_key(&self, session: &Session) -> Result<(), Error> {
        let (key, key_type) = session
            .host_key()
            .ok_or_else(|| Error::Backend("sftp: server offered no host key".to_string()))?;

        let mut known_hosts = session.known_hosts().map_err(sftp_error)?;

        let file = std::env::var_os("BEEPING_SSH_KNOWN_HOSTS")
            .map(PathBuf::from)
            .or_else(|| dirs::home_dir().map(|home| home.join(".ssh/known_hosts")));

        if let Some(file) = &file {
            // A missing file just means no hosts are known yet
            let _ = known_hosts.read_file(file, KnownHostFileKind::OpenSSH);
        }

        match known_hosts.check_port(&self.host, self.port, key) {
            CheckResult::Match => Ok(()),
            CheckResult::Mismatch => Err(Error::Backend(format!(
                "sftp: HOST KEY MISMATCH for {}: the server's key does not \
                 match ~/.ssh/known_hosts; someone may be intercepting the \
                 connection",
                self.host
            ))),
            CheckResult::NotFound | CheckResult::Failure => {
                // Trust on first use, and remember the decision
                let name = if self.port == 22 {
                    self.host.clone()
                } else {
                    format!("[{}]:{}", self.host, self.port)
                };

                known_hosts
                    .add(&name, key, "added by beeping", key_type.into())
                    .map_err(sftp_error)?;

                if let Some(file) = &file {
                    known_hosts
                        .write_file(file, KnownHostFileKind::OpenSSH)
                        .map_err(sftp_error)?;
                }

                Ok(())
            }
        }
    }

    fn authenticate(&self, session: &Session) -> Result<(), Error> {
        if let Some(key) = std::env::var_os("BEEPING_SSH_KEY")
            && session
                .userauth_pubkey_file(&self.user, None, std::path::Path::new(&key), None)
                .is_ok()
            && session.authenticated()
        {
            return Ok(());
        }

        if session.userauth_agent(&self.user).is_ok() && session.authenticated() {
            return Ok(());
        }

        if let Some(home) = dirs::home_dir() {
            for name in ["id_ed25519", "id_rsa"] {
                let key = home.join(".ssh").join(name);

                if key.exists()
                    && session
                        .userauth_pubkey_file(&self.user, None, &key, None)
                        .is_ok()
                    && session.authenticated()
                {
                    return Ok(());
                }
            }
        }

        let passwords = std::env::var("BEEPING_SSH_PASSWORD")
            .ok()
            .into_iter()
            .chain(self.bookmark_password.clone());

        for password in passwords {
            if session.userauth_password(&self.user, &password).is_ok() && session.authenticated() {
                return Ok(());
            }
        }

        Err(Error::Backend(format!(
            "sftp: could not authenticate as {}@{} (tried agent, key \
             files, $BEEPING_SSH_PASSWORD, and lftp bookmarks)",
            self.user, self.host
        )))
    }

    fn object_path(&self, key: &ObjectKey) -> PathBuf {
        let mut path = self.root.clone();

        for segment in key.segments() {
            path.push(segment);
        }

        path
    }

    fn dir_names(&self, dir: PathBuf) -> Result<Vec<String>, Error> {
        self.with_connection(|connection| match connection.sftp.readdir(&dir) {
            Ok(entries) => Ok(entries
                .into_iter()
                .filter_map(|(path, _)| {
                    path.file_name()
                        .and_then(|name| name.to_str())
                        .map(str::to_string)
                })
                .collect()),
            Err(err) if is_not_found(&err) => Ok(Vec::new()),
            Err(err) => Err(sftp_error(err)),
        })
    }
}

impl Backend for SftpBackend {
    fn put(&self, key: &ObjectKey, data: &[u8]) -> Result<(), Error> {
        let path = self.object_path(key);

        self.with_connection(|connection| {
            let sftp = &connection.sftp;

            if sftp.stat(&path).is_ok() {
                // Write-once: the object is already this one
                return Ok(());
            }

            // Create missing parents; errors mean "exists" more
            // often than not, and real trouble fails the write below
            let mut parent = self.root.clone();
            let segments = key.segments();

            let _ = sftp.mkdir(&parent, 0o755);

            for segment in &segments[..segments.len() - 1] {
                parent.push(segment);
                let _ = sftp.mkdir(&parent, 0o755);
            }

            // Write beside the final name, then rename into place
            let temp = {
                let mut temp = path.clone().into_os_string();
                temp.push(temp_suffix());
                PathBuf::from(temp)
            };

            {
                let mut file = sftp.create(&temp).map_err(sftp_error)?;
                file.write_all(data)
                    .map_err(|err| Error::Backend(format!("sftp: writing: {err}")))?;
                // Not every server implements the fsync extension;
                // decline is not failure
                let _ = file.fsync();
            }

            let flags = RenameFlags::ATOMIC | RenameFlags::OVERWRITE | RenameFlags::NATIVE;

            match sftp.rename(&temp, &path, Some(flags)) {
                Ok(()) => Ok(()),
                Err(_) if sftp.stat(&path).is_ok() => {
                    // Lost a benign race to a concurrent writer
                    let _ = sftp.unlink(&temp);

                    Ok(())
                }
                Err(err) => Err(sftp_error(err)),
            }
        })
    }

    fn get(&self, key: &ObjectKey) -> Result<Option<Vec<u8>>, Error> {
        let path = self.object_path(key);

        self.with_connection(|connection| match connection.sftp.open(&path) {
            Ok(mut file) => {
                let mut data = Vec::new();

                file.read_to_end(&mut data)
                    .map_err(|err| Error::Backend(format!("sftp: reading: {err}")))?;

                Ok(Some(data))
            }
            Err(err) if is_not_found(&err) => Ok(None),
            Err(err) => Err(sftp_error(err)),
        })
    }

    fn contains(&self, key: &ObjectKey) -> Result<bool, Error> {
        let path = self.object_path(key);

        self.with_connection(|connection| match connection.sftp.stat(&path) {
            Ok(_) => Ok(true),
            Err(err) if is_not_found(&err) => Ok(false),
            Err(err) => Err(sftp_error(err)),
        })
    }

    fn list(
        &self,
        kind: ObjectKind,
        visit: &mut dyn FnMut(ObjectKey) -> Result<(), Error>,
    ) -> Result<(), Error> {
        match kind {
            ObjectKind::Header => {
                if self.contains(&ObjectKey::Header)? {
                    visit(ObjectKey::Header)?;
                }
            }
            // The flat kinds: one directory of hex-named objects each
            ObjectKind::Snapshot | ObjectKind::Prune | ObjectKind::Lock => {
                for name in self.dir_names(self.root.join(kind.directory()))? {
                    if let Some(key) = kind.key_for(&name) {
                        visit(key)?;
                    }
                }
            }
            ObjectKind::Chunk => {
                for prefix in self.dir_names(self.root.join("chunks"))? {
                    if prefix.len() != 2 || !prefix.bytes().all(|b| b.is_ascii_hexdigit()) {
                        continue;
                    }

                    for name in self.dir_names(self.root.join("chunks").join(&prefix))? {
                        if let Ok(id) = ChunkId::from_hex(name) {
                            visit(ObjectKey::Chunk(id))?;
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Asks the server for the filesystem statistics of the repository
    /// directory, over the OpenSSH `fstatvfs` SFTP extension.
    ///
    /// Servers that do not implement it simply cannot say, which is not
    /// an error — the floor is not applied rather than the backup
    /// failing — and the refusal is remembered so the question is asked
    /// only once.
    fn free_space(&self) -> Result<Option<u64>, Error> {
        if self.space_unknown.load(Ordering::Acquire) {
            return Ok(None);
        }

        let stats = self.with_connection(|connection| {
            let mut dir = connection.sftp.opendir(&self.root).map_err(sftp_error)?;

            dir.statvfs().map_err(sftp_error)
        });

        match stats {
            Ok(stats) => Ok(Some(stats.f_bavail.saturating_mul(stats.f_frsize))),
            Err(_) => {
                self.space_unknown.store(true, Ordering::Release);

                Ok(None)
            }
        }
    }

    /// Adds up everything under the repository directory, one readdir
    /// per directory — a few hundred round trips for the fanned-out
    /// chunk store, which is why this is asked once a run and not more.
    fn used_space(&self) -> Result<Option<u64>, Error> {
        self.with_connection(|connection| tree_bytes(&connection.sftp, &self.root))
            .map(Some)
    }

    fn delete(&self, key: &ObjectKey) -> Result<(), Error> {
        let path = self.object_path(key);

        self.with_connection(|connection| match connection.sftp.unlink(&path) {
            Ok(()) => Ok(()),
            Err(err) if is_not_found(&err) => Ok(()),
            Err(err) => Err(sftp_error(err)),
        })
    }
}

/// How many bytes the tree under `path` holds, by the sizes the server
/// reports. A directory that is not there holds nothing.
fn tree_bytes(sftp: &Sftp, path: &std::path::Path) -> Result<u64, Error> {
    let entries = match sftp.readdir(path) {
        Ok(entries) => entries,
        Err(err) if is_not_found(&err) => return Ok(0),
        Err(err) => return Err(sftp_error(err)),
    };

    let mut total = 0;

    for (child, stat) in entries {
        total += if stat.is_dir() {
            tree_bytes(sftp, &child)?
        } else {
            stat.size.unwrap_or(0)
        };
    }

    Ok(total)
}

fn is_not_found(err: &ssh2::Error) -> bool {
    matches!(err.code(), ErrorCode::SFTP(SFTP_NO_SUCH_FILE))
}

fn sftp_error(err: ssh2::Error) -> Error {
    Error::Backend(format!("sftp: {err}"))
}