use chrono::{DateTime, Utc};
use hmac::{Hmac, Mac, NewMac};
use sha2::Sha256;

const DEVICEID_KEY: &str = "DeviceId";
const HOSTNAME_KEY: &str = "HostName";
const SHAREDACCESSKEY_KEY: &str = "SharedAccessKey";

///
#[derive(Debug)]
pub enum TokenError {
    ///
    ConnectionStringMissingRequiredParameter(&'static str),
    ///
    InvalidKeyFormat,
    ///
    InvalidKeyLength,
}

/// Provider for authentication
pub trait TokenSource {
    /// Get the authentication value from the source
    fn get(&self, expiry: &DateTime<Utc>) -> String;
}

impl std::fmt::Debug for dyn TokenSource {
    fn fmt(&self, _: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        unimplemented!("Only implemented by implementations of TokenSource")
    }
}

/// Provide an existing SAS as authentication source
#[derive(Debug, Clone)]
pub struct SasTokenSource<'a> {
    sas: &'a str,
}

impl<'a> SasTokenSource<'_> {
    /// Make the token source with the given SAS
    pub fn new(sas: &'a str) -> SasTokenSource<'_> {
        SasTokenSource { sas }
    }
}

impl<'a> TokenSource for SasTokenSource<'_> {
    fn get(&self, _: &DateTime<Utc>) -> String {
        self.sas.to_string()
    }
}

/// Authenticate using the devices key
#[derive(Debug, Clone)]
pub struct DeviceKeyTokenSource<'a> {
    resource_uri: String,
    key: &'a str,
}

impl<'a> DeviceKeyTokenSource<'_> {
    /// Make the source from the devices individual details
    pub fn new(
        hub: &str,
        device_id: &str,
        key: &'a str,
    ) -> Result<DeviceKeyTokenSource<'a>, TokenError> {
        // Verify key is base64
        let b64_key = base64::decode(&key).map_err(|_| TokenError::InvalidKeyFormat)?;
        // Verify key is the right length for Hmac
        Hmac::<Sha256>::new_varkey(&b64_key).map_err(|_| TokenError::InvalidKeyLength)?;

        Ok(DeviceKeyTokenSource {
            resource_uri: format!("{}%2Fdevices%2F{}", hub, device_id),
            key,
        })
    }

    /// Make the source from a connection string for the device
    pub fn new_from_connection_string(
        connection_string: &'a str,
    ) -> Result<DeviceKeyTokenSource<'a>, TokenError> {
        let mut key = None;
        let mut device_id = None;
        let mut hub = None;

        let parts: Vec<&str> = connection_string.split(';').collect();
        for p in parts {
            let s: Vec<&str> = p.split('=').collect();
            match s[0] {
                SHAREDACCESSKEY_KEY => key = Some(s[1]),
                DEVICEID_KEY => device_id = Some(s[1]),
                HOSTNAME_KEY => hub = Some(s[1]),
                _ => (), // Ignore extraneous component in the connection string
            }
        }

        let hub = hub.ok_or(TokenError::ConnectionStringMissingRequiredParameter(
            HOSTNAME_KEY,
        ))?;
        let device_id = device_id.ok_or(TokenError::ConnectionStringMissingRequiredParameter(
            DEVICEID_KEY,
        ))?;
        let key = key.ok_or(TokenError::ConnectionStringMissingRequiredParameter(
            SHAREDACCESSKEY_KEY,
        ))?;

        Ok(Self::new(hub, device_id, key)?)
    }
}

impl<'a> TokenSource for DeviceKeyTokenSource<'_> {
    fn get(&self, expiry: &DateTime<Utc>) -> String {
        let expiry_timestamp = expiry.timestamp();

        let to_sign = format!("{}\n{}", &self.resource_uri, expiry_timestamp);

        let token = generate_token(self.key, &to_sign);

        let sas = format!(
            "SharedAccessSignature sr={}&{}&se={}",
            self.resource_uri, token, expiry_timestamp
        );

        sas
    }
}

pub(crate) fn generate_token(key: &str, message: &str) -> String {
    // Checked base64 and hmac in new so should be safe to unwrap here
    let key = base64::decode(&key).unwrap();
    let mut mac = Hmac::<Sha256>::new_varkey(&key).unwrap();
    mac.update(message.as_bytes());
    let mac_result = mac.finalize();
    let signature = base64::encode(mac_result.into_bytes());

    let pairs = &vec![("sig", signature)];
    serde_urlencoded::to_string(pairs).unwrap()
}

#[cfg(test)]
mod tests {
    // Note this useful idiom: importing names from outer (for mod tests) scope.
    use super::*;
    use chrono::prelude::*;

    #[test]
    fn test_sas_source() {
        // Should pass value through unchanged
        let sas_token_source = SasTokenSource::new("MySasString");
        assert_eq!(sas_token_source.get(&Utc::now()), "MySasString".to_string());
    }

    #[test]
    fn test_key_source() {
        let key_source = DeviceKeyTokenSource::new(
            "pontifex.azure-devices.net",
            "FirstDevice",
            "O+H9VTcdJP0TQkl7bh4nVG0OJNrEataMpuWB54D0VEc=",
        )
        .unwrap();
        let expiry = Utc.ymd(2020, 6, 28).and_hms(14, 08, 25);
        assert_eq!(key_source.get(&expiry), "SharedAccessSignature sr=pontifex.azure-devices.net%2Fdevices%2FFirstDevice&sig=CKYVArtLm72J2UNWLb4V3XqPc679Ig3LX83G3nPExUc%3D&se=1593353305");
    }
}
