use std::{fmt::Debug, sync::Arc, time::Duration};

use anyhow::{Context as _, Result as AnyhowResult};
use russh::{
    client::{AuthResult, Handle, Handler},
    keys::PrivateKeyWithHashAlg,
};
use tokio::net::ToSocketAddrs;

pub(crate) struct ClientHandler {}

impl Handler for ClientHandler {
    type Error = russh::Error;

    async fn check_server_key(
        &mut self,
        server_public_key: &russh::keys::ssh_key::PublicKey,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }
}

pub(crate) async fn connect_ssh<A: ToSocketAddrs + Debug, K: AsRef<[u8]>, P: Into<String>>(
    connect_timeout: Duration,
    auth_timeout: Duration,
    addr: A,
    username: &str,
    key: Option<K>,
    password: Option<P>,
) -> AnyhowResult<Handle<ClientHandler>> {
    let ssh_config = Arc::new(russh::client::Config {
        inactivity_timeout: None,
        ..Default::default()
    });
    log::info!(
        "Connecting to SSH {:?} {} key & {} password",
        &addr,
        if key.is_some() { "with" } else { "without" },
        if password.is_some() {
            "with"
        } else {
            "without"
        }
    );
    let mut connection = tokio::time::timeout(
        connect_timeout,
        russh::client::connect(ssh_config, addr, crate::ssh::ClientHandler {}),
    )
    .await
    .context("Connection to host timed out")?
    .context("Cannot connect to host")?;

    let key = match key {
        Some(key) => {
            let privkey = russh::keys::PrivateKey::from_openssh(key).context("Invalid SSH key")?;
            Some(PrivateKeyWithHashAlg::new(
                Arc::new(privkey),
                connection.best_supported_rsa_hash().await?.flatten(),
            ))
        }
        None => None,
    };

    let mut auth_success = false;
    if !auth_success {
        if let Some(key) = key {
            log::info!("Trying to authenticate using key");
            match tokio::time::timeout(
                auth_timeout,
                connection.authenticate_publickey(username, key),
            )
            .await
            {
                Ok(Ok(AuthResult::Success)) => {
                    auth_success = true;
                }
                Ok(Ok(AuthResult::Failure {
                    remaining_methods,
                    partial_success,
                })) => {
                    if partial_success || !remaining_methods.contains(&russh::MethodKind::Password)
                    {
                        if partial_success {
                            anyhow::bail!("SSH auth failed: multi-factor auth is not supported");
                        } else {
                            anyhow::bail!("SSH auth failed: no auth methods left");
                        }
                    }
                    log::warn!("SSH public key rejected");
                }
                Ok(Err(e)) => {
                    log::warn!("SSH public key auth failed: {e}");
                }
                Err(_) => {
                    log::warn!("SSH public key auth timed out");
                }
            }
        } else {
            log::info!("No key provided");
        }
    }
    if !auth_success {
        if let Some(password) = password {
            log::info!("Trying to authenticate using password");
            match tokio::time::timeout(
                auth_timeout,
                connection.authenticate_password(username, password),
            )
            .await
            {
                Ok(Ok(AuthResult::Success)) => {
                    auth_success = true;
                }
                Ok(Ok(AuthResult::Failure {
                    remaining_methods: _,
                    partial_success,
                })) => {
                    if partial_success {
                        anyhow::bail!("SSH auth failed: multi-factor auth is not supported");
                    } else {
                        anyhow::bail!("SSH auth failed: no auth methods left");
                    }
                }
                Ok(Err(e)) => {
                    log::warn!("SSH password auth failed: {e}");
                }
                Err(_) => {
                    log::warn!("SSH password auth timed out");
                }
            }
        } else {
            log::info!("No password provided");
        }
    }
    if !auth_success {
        anyhow::bail!("SSH auth failed: no auth methods left");
    }
    Ok(connection)
}