use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

pub const PROTOCOL_VERSION: u16 = 1;
pub const SESSION_KEY_BYTES: usize = 32;
pub const AES_GCM_NONCE_BYTES: usize = 12;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload", rename_all = "snake_case")]
pub enum WireMessage {
    Hello(HelloMessage),
    HelloAck(HelloAckMessage),
    Encrypted(EncryptedEnvelope),
    Error(ProtocolErrorMessage),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HelloMessage {
    pub protocol_version: u16,
    pub key_id: String,
    pub session_key: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HelloAckMessage {
    pub protocol_version: u16,
    pub key_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedEnvelope {
    pub enc: String,
    pub key_id: String,
    pub nonce: String,
    pub ciphertext: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProtocolErrorMessage {
    pub message: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload", rename_all = "snake_case")]
pub enum ClientMessage {
    Subscribe { topic: String },
    Ping,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload", rename_all = "snake_case")]
pub enum ServerMessage {
    Snapshot { records: Vec<RecordSummary> },
    RecordChanged { record: RecordSummary },
    Pong,
    Error { message: String },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecordSummary {
    pub id: Uuid,
    pub title: String,
    pub updated_at: DateTime<Utc>,
}

pub mod crypto {
    use aes_gcm::{
        aead::{Aead, AeadCore, KeyInit, OsRng},
        Aes256Gcm, Key, Nonce,
    };
    use base64::{engine::general_purpose::STANDARD, Engine};
    use rand_core::RngCore;
    use serde::{de::DeserializeOwned, Serialize};

    use crate::{EncryptedEnvelope, AES_GCM_NONCE_BYTES, PROTOCOL_VERSION, SESSION_KEY_BYTES};

    #[derive(Debug)]
    pub enum CryptoError {
        InvalidBase64(base64::DecodeError),
        InvalidKeyLength(usize),
        InvalidNonceLength(usize),
        Serialize(serde_json::Error),
        Deserialize(serde_json::Error),
        Encrypt,
        Decrypt,
    }

    impl std::fmt::Display for CryptoError {
        fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            match self {
                Self::InvalidBase64(error) => write!(formatter, "invalid base64: {error}"),
                Self::InvalidKeyLength(length) => write!(formatter, "invalid key length: {length}"),
                Self::InvalidNonceLength(length) => {
                    write!(formatter, "invalid nonce length: {length}")
                }
                Self::Serialize(error) => write!(formatter, "serialize failed: {error}"),
                Self::Deserialize(error) => write!(formatter, "deserialize failed: {error}"),
                Self::Encrypt => write!(formatter, "encryption failed"),
                Self::Decrypt => write!(formatter, "decryption failed"),
            }
        }
    }

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

    #[derive(Debug, Clone)]
    pub struct SessionKey {
        key_id: String,
        key: [u8; SESSION_KEY_BYTES],
    }

    impl SessionKey {
        pub fn generate() -> Self {
            let mut key = [0_u8; SESSION_KEY_BYTES];
            OsRng.fill_bytes(&mut key);

            Self {
                key_id: uuid::Uuid::new_v4().to_string(),
                key,
            }
        }

        pub fn from_base64(key_id: String, value: &str) -> Result<Self, CryptoError> {
            let decoded = STANDARD.decode(value).map_err(CryptoError::InvalidBase64)?;
            if decoded.len() != SESSION_KEY_BYTES {
                return Err(CryptoError::InvalidKeyLength(decoded.len()));
            }

            let mut key = [0_u8; SESSION_KEY_BYTES];
            key.copy_from_slice(&decoded);

            Ok(Self { key_id, key })
        }

        pub fn key_id(&self) -> &str {
            &self.key_id
        }

        pub fn to_base64(&self) -> String {
            STANDARD.encode(self.key)
        }

        pub fn encrypt<T: Serialize>(&self, value: &T) -> Result<EncryptedEnvelope, CryptoError> {
            let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&self.key));
            let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
            let plaintext = serde_json::to_vec(value).map_err(CryptoError::Serialize)?;
            let ciphertext = cipher
                .encrypt(&nonce, plaintext.as_ref())
                .map_err(|_| CryptoError::Encrypt)?;

            Ok(EncryptedEnvelope {
                enc: format!("aes-256-gcm-v{PROTOCOL_VERSION}"),
                key_id: self.key_id.clone(),
                nonce: STANDARD.encode(nonce),
                ciphertext: STANDARD.encode(ciphertext),
            })
        }

        pub fn decrypt<T: DeserializeOwned>(
            &self,
            envelope: &EncryptedEnvelope,
        ) -> Result<T, CryptoError> {
            let nonce_bytes = STANDARD
                .decode(&envelope.nonce)
                .map_err(CryptoError::InvalidBase64)?;
            if nonce_bytes.len() != AES_GCM_NONCE_BYTES {
                return Err(CryptoError::InvalidNonceLength(nonce_bytes.len()));
            }

            let ciphertext = STANDARD
                .decode(&envelope.ciphertext)
                .map_err(CryptoError::InvalidBase64)?;
            let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&self.key));
            let plaintext = cipher
                .decrypt(Nonce::from_slice(&nonce_bytes), ciphertext.as_ref())
                .map_err(|_| CryptoError::Decrypt)?;

            serde_json::from_slice(&plaintext).map_err(CryptoError::Deserialize)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{crypto::SessionKey, ClientMessage};

    #[test]
    fn session_key_encrypts_and_decrypts_client_message() {
        let session_key = SessionKey::generate();
        let message = ClientMessage::Subscribe {
            topic: "customers".to_string(),
        };

        let envelope = session_key.encrypt(&message).expect("encrypt message");
        let decrypted = session_key
            .decrypt::<ClientMessage>(&envelope)
            .expect("decrypt message");

        assert_eq!(decrypted, message);
    }
}
