use anyhow::{Context, Result};
use nix::unistd::Uid;
use quinn::{SendStream, RecvStream};
use zeroize::Zeroizing;

use crate::{read_msg, write_msg, ClientConfig, Message, ser::ByteSlice};
use super::Question;

#[derive(Debug, Serialize, Deserialize)]
pub struct AuthRequest<'a> {
    #[serde(borrow)]
    pub nonce: ByteSlice<'a>,
}

impl<'a> AuthRequest<'a> {
    pub const NAMESPACE: &str = "QUINOA AUTHENTICATION";

    pub fn new(nonce: &'a [u8]) -> AuthRequest<'a> {
        Self {
            nonce: nonce.into(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct AuthResponse<'a> {
    #[serde(borrow)]
    signature: &'a str,
}

#[cfg(feature = "server")]
async fn server_authorize_key(user_id: Uid, public_key: &ssh_key::public::PublicKey) -> Result<()> {
    const PATH: &str = "/.ssh/authorized_keys";

    let mut user_passwd_buf = Vec::new();
    let user_passwd = crate::passwd::Passwd::from_uid(user_id, &mut user_passwd_buf)?
        .context("No passwd entry for user")?;

    let home = user_passwd.dir.to_str()?;

    let mut authorized_keys = String::with_capacity(home.len() + PATH.len());
    authorized_keys.push_str(home);
    authorized_keys.push_str(PATH);
    match tokio::fs::read_to_string(&authorized_keys).await {
        Ok(authorized_keys) => {
            let mut authorized_keys = ssh_key::authorized_keys::AuthorizedKeys::new(&authorized_keys);

            while let Some(r) = authorized_keys.next() {
                let entry = r?;
                if entry.public_key().key_data() == public_key.key_data() {
                    return Ok(());
                }
            }
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
        Err(e) => {
            error!("unable to read authorized keys file at {:?}: {}", authorized_keys, e);
        }
    };
    Err(anyhow!("Key provided by the client is not authorized to connect"))
}

#[cfg(feature = "server")]
pub async fn server_authenticate(send: &mut SendStream, recv: &mut RecvStream, user_id: Uid, public_key: &str) -> Result<()> {
    use rand::RngCore;

    let public_key = ssh_key::public::PublicKey::from_openssh(public_key)?;

    server_authorize_key(user_id, &public_key).await?;

    let mut nonce = vec![0; 256];
    rand::thread_rng().try_fill_bytes(&mut nonce)?;
    let request = AuthRequest::new(&nonce);
    let request = Message::from_value(&request)?;
    request.write_ref(send).await?;

    let mut buf = Vec::new();
    let response = read_msg::<AuthResponse>(recv, &mut buf).await??;

    let sig = ssh_key::SshSig::from_pem(&response.signature)?;
    public_key.verify(AuthRequest::NAMESPACE, &request.0, &sig)?;
    
    Ok(())
}

pub async fn client_authenticate(cfg: &ClientConfig, send: &mut SendStream, recv: &mut RecvStream, username: &str) -> Result<()> {
    let private_key = Zeroizing::new(tokio::fs::read_to_string(&cfg.ssh_key_file).await?);
    let private_key = ssh_key::private::PrivateKey::from_openssh(private_key)
        .context("Loading private key")?;
    let public_key = private_key.public_key();
    let public_key = public_key.to_openssh()
        .context("Encoding public key")?;
    info!("loaded private key: {:?}", private_key.algorithm());

    write_msg(send, &super::Hello {
        username,
        auth_method: super::Method::SshKey {
            public_key: &public_key,
        },
    }).await?;
    info!("sent hello");

    let mut buf = Vec::new();
    _ = read_msg::<AuthRequest>(recv, &mut buf).await??;
    info!("read auth request");

    let private_key = if private_key.is_encrypted() {
        let password = Zeroizing::new(rpassword::prompt_password("Enter the key's passphrase: ")?);
        private_key.decrypt(&password).context("Incorrect passphrase")?
    } else {
        private_key
    };

    let sig = private_key.sign(AuthRequest::NAMESPACE, ssh_key::HashAlg::Sha512, &buf)?;
    let sig = sig.to_pem(ssh_key::LineEnding::LF)?;
    info!("signed auth request");

    write_msg(send, &AuthResponse { signature: &sig }).await?;
    info!("wrote auth request");

    let Question::LoggedIn = read_msg(recv, &mut buf).await?? else {
        bail!("Received an unexpected question")
    };

    Ok(())
}