use crate::{get_session_as_id_from_uuid, ApAddress, ApContext};
use anyhow::{anyhow, Result};
use base64::engine::{general_purpose, Engine as _};
use jdt_maybe_multiple::MaybeMultiple;
use openmls::{
    group::GroupId,
    prelude::{
        tls_codec::Serialize, BasicCredential, CredentialWithKey, DeserializeBytes, KeyPackage,
        KeyPackageBundle, KeyPackageIn, MlsMessageBodyOut, MlsMessageOut, OpenMlsProvider, Welcome,
    },
};
use openmls_basic_credential::SignatureKeyPair;
use openmls_rust_crypto::{MemoryStorage, OpenMlsRustCrypto};
use serde::Deserialize;
use serde_with::serde_as;
use std::collections::HashMap;
use std::fmt::{Debug, Display, Formatter};
use uuid::Uuid;
use vodozemac::olm::Session;
use wasm_bindgen::{
    convert::{FromWasmAbi, IntoWasmAbi},
    describe::WasmDescribe,
    prelude::*,
};

#[derive(serde::Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct ApSession {
    #[serde(rename = "@context")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub context: Option<ApContext>,
    #[serde(rename = "type")]
    pub kind: ApSessionType,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    pub to: ApAddress,
    pub attributed_to: ApAddress,
    pub instrument: MaybeMultiple<ApInstrument>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reference: Option<String>,
    pub uuid: Option<String>,
}

impl Debug for ApSession {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Session[from: {}, to: {}{}{}]",
            self.attributed_to,
            self.to,
            self.uuid
                .as_ref()
                .map(|u| format!(", uuid: {}", u))
                .unwrap_or_default(),
            self.reference
                .as_ref()
                .map(|r| format!(", ref: {}", r))
                .unwrap_or_default()
        )
    }
}

impl Display for ApSession {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        Debug::fmt(self, f)
    }
}

impl Default for ApSession {
    fn default() -> ApSession {
        let uuid = Uuid::new_v4().to_string();
        ApSession {
            context: Some(ApContext::default()),
            kind: ApSessionType::default(),
            id: Some(get_session_as_id_from_uuid(uuid.clone())),
            to: ApAddress::None,
            attributed_to: ApAddress::None,
            instrument: MaybeMultiple::None,
            reference: Option::None,
            uuid: Some(uuid),
        }
    }
}

#[derive(serde::Serialize, Deserialize, Clone, Debug, Default, PartialEq)]
pub enum ApSessionType {
    #[default]
    #[serde(alias = "encrypted_session")]
    EncryptedSession,
}

#[derive(serde::Serialize, Deserialize, Clone, Debug, Default, PartialEq, Eq, Ord, PartialOrd)]
pub enum ApInstrumentType {
    #[default]
    #[serde(alias = "olm_identity_key")]
    OlmIdentityKey,
    #[serde(alias = "olm_one_time_key")]
    OlmOneTimeKey,
    #[serde(alias = "olm_session")]
    OlmSession,
    #[serde(alias = "olm_account")]
    OlmAccount,
    #[serde(alias = "vault_item")]
    VaultItem,
    #[serde(alias = "mls_signer")]
    MlsCredentials,
    #[serde(alias = "mls_key_package")]
    MlsKeyPackage,
    #[serde(alias = "mls_storage")]
    MlsStorage,
    #[serde(alias = "mls_welcome")]
    MlsWelcome,
    #[serde(alias = "mls_group_id")]
    MlsGroupId,
}

impl Display for ApInstrumentType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        Debug::fmt(self, f)
    }
}

#[derive(serde::Serialize, Deserialize, Clone, Default, PartialEq, Eq, Ord, PartialOrd)]
pub struct ApInstrument {
    #[serde(rename = "type")]
    pub kind: ApInstrumentType,
    pub id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub hash: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub uuid: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub url: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mutation_of: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub conversation: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub activity: Option<String>,
}

impl Debug for ApInstrument {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Instrument[{}[{}{}{}]",
            self.kind,
            self.id
                .as_ref()
                .or(self.uuid.as_ref())
                .map(|i| i.as_str())
                .unwrap_or("unnamed"),
            self.name
                .as_ref()
                .map(|n| format!(", {}", n))
                .unwrap_or_default(),
            self.hash
                .as_ref()
                .map(|h| format!(", hash:{}", h))
                .unwrap_or_default()
        )
    }
}

impl Display for ApInstrument {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        Debug::fmt(self, f)
    }
}

impl WasmDescribe for ApInstrument {
    fn describe() {
        JsValue::describe()
    }
}

impl FromWasmAbi for ApInstrument {
    type Abi = <JsValue as FromWasmAbi>::Abi;

    unsafe fn from_abi(js: Self::Abi) -> Self {
        let value = JsValue::from_abi(js);
        serde_wasm_bindgen::from_value(value).unwrap_or_default()
    }
}

impl IntoWasmAbi for ApInstrument {
    type Abi = <JsValue as IntoWasmAbi>::Abi;

    fn into_abi(self) -> Self::Abi {
        serde_wasm_bindgen::to_value(&self)
            .unwrap_or_default()
            .into_abi()
    }
}

impl ApInstrument {
    pub fn is_olm_identity_key(&self) -> bool {
        matches!(self.kind, ApInstrumentType::OlmIdentityKey)
    }

    pub fn is_olm_one_time_key(&self) -> bool {
        matches!(self.kind, ApInstrumentType::OlmOneTimeKey)
    }

    pub fn is_olm_session(&self) -> bool {
        matches!(self.kind, ApInstrumentType::OlmSession)
    }

    pub fn is_olm_account(&self) -> bool {
        matches!(self.kind, ApInstrumentType::OlmAccount)
    }

    pub fn is_vault_item(&self) -> bool {
        matches!(self.kind, ApInstrumentType::VaultItem)
    }

    pub fn is_mls_credentials(&self) -> bool {
        matches!(self.kind, ApInstrumentType::MlsCredentials)
    }

    pub fn is_mls_key_package(&self) -> bool {
        matches!(self.kind, ApInstrumentType::MlsKeyPackage)
    }

    pub fn is_mls_storage(&self) -> bool {
        matches!(self.kind, ApInstrumentType::MlsStorage)
    }

    pub fn is_mls_welcome(&self) -> bool {
        matches!(self.kind, ApInstrumentType::MlsWelcome)
    }

    pub fn is_mls_group_id(&self) -> bool {
        matches!(self.kind, ApInstrumentType::MlsGroupId)
    }
}

type EncryptableVaultItem<E> = (String, E);
impl<E> TryFrom<EncryptableVaultItem<E>> for ApInstrument
where
    E: Fn(Vec<u8>) -> Vec<u8>,
{
    type Error = anyhow::Error;

    fn try_from((data, encrypt): EncryptableVaultItem<E>) -> Result<Self, Self::Error> {
        let encrypted = encrypt(data.as_bytes().to_vec());
        let content = Some(general_purpose::STANDARD.encode(encrypted));

        Ok(ApInstrument {
            kind: ApInstrumentType::VaultItem,
            id: None,
            content,
            hash: None,
            uuid: None,
            name: None,
            url: None,
            mutation_of: None,
            conversation: None,
            activity: None,
        })
    }
}

type EncryptableOlmSession<H> = (Session, Vec<u8>, H);
impl<H> TryFrom<EncryptableOlmSession<H>> for ApInstrument
where
    H: Fn(Vec<u8>) -> String,
{
    type Error = anyhow::Error;

    fn try_from((session, key, hash): EncryptableOlmSession<H>) -> Result<Self, Self::Error> {
        let pickle = session.pickle();
        let hash = Some(hash(serde_json::to_string(&pickle)?.into_bytes()));
        let content = Some(pickle.encrypt(key.as_slice().try_into()?));
        Ok(ApInstrument {
            kind: ApInstrumentType::OlmSession,
            id: None,
            content,
            hash,
            uuid: None,
            name: None,
            url: None,
            mutation_of: None,
            conversation: None,
            activity: None,
        })
    }
}

#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct CredentialKeyPair {
    pub credential_with_key: CredentialWithKey,
    pub key_pair: SignatureKeyPair,
}

type AddressAndSignatureKeyPair = (ApAddress, SignatureKeyPair);
impl From<AddressAndSignatureKeyPair> for CredentialKeyPair {
    fn from((address, keys): AddressAndSignatureKeyPair) -> Self {
        let credential = BasicCredential::new(address.to_string().as_bytes().to_vec());

        CredentialKeyPair {
            credential_with_key: CredentialWithKey {
                credential: credential.into(),
                signature_key: keys.public().into(),
            },
            key_pair: keys,
        }
    }
}

type EncryptableCredentialKeyPair<F> = (CredentialKeyPair, F);
impl<F> From<EncryptableCredentialKeyPair<F>> for ApInstrument
where
    F: Fn(Vec<u8>) -> Vec<u8>,
{
    fn from((keys, encrypt): EncryptableCredentialKeyPair<F>) -> Self {
        let v = serde_json::to_vec(&keys).unwrap();
        let e = encrypt(v);
        let content = Some(general_purpose::STANDARD.encode(&e));

        ApInstrument {
            kind: ApInstrumentType::MlsCredentials,
            content,
            ..Default::default()
        }
    }
}

impl From<GroupId> for ApInstrument {
    fn from(group_id: GroupId) -> Self {
        let v = serde_json::to_vec(&group_id).unwrap();
        let content = Some(general_purpose::STANDARD.encode(&v));

        ApInstrument {
            kind: ApInstrumentType::MlsGroupId,
            content,
            ..Default::default()
        }
    }
}

impl TryFrom<ApInstrument> for GroupId {
    type Error = anyhow::Error;

    fn try_from(instrument: ApInstrument) -> Result<Self> {
        if !instrument.is_mls_group_id() {
            return Err(anyhow!("Instrument is not a GroupId"));
        }

        let content = instrument.content.ok_or(anyhow!("content must be Some"))?;
        let v = general_purpose::STANDARD.decode(content)?;
        serde_json::from_slice(v.as_slice()).map_err(anyhow::Error::msg)
    }
}

impl TryFrom<MlsMessageOut> for ApInstrument {
    type Error = anyhow::Error;

    fn try_from(message: MlsMessageOut) -> Result<Self> {
        let body = message.body();

        match body {
            MlsMessageBodyOut::Welcome(welcome) => Ok(Self {
                kind: ApInstrumentType::MlsWelcome,
                content: Some(general_purpose::STANDARD.encode(welcome.tls_serialize_detached()?)),
                ..Default::default()
            }),
            MlsMessageBodyOut::KeyPackage(key) => Ok(Self {
                kind: ApInstrumentType::MlsKeyPackage,
                content: Some(general_purpose::STANDARD.encode(key.tls_serialize_detached()?)),
                ..Default::default()
            }),
            _ => Err(anyhow!("unimplemented MlsMessageOut")),
        }
    }
}

impl TryFrom<ApInstrument> for Welcome {
    type Error = anyhow::Error;

    fn try_from(instrument: ApInstrument) -> Result<Self> {
        let content = instrument.content.ok_or(anyhow!("Content must be Some"))?;
        match instrument.kind {
            ApInstrumentType::MlsWelcome => {
                Ok(Welcome::tls_deserialize_bytes(&general_purpose::STANDARD.decode(content)?)?.0)
            }
            _ => Err(anyhow!("ApInstrument not a Welcome")),
        }
    }
}

impl TryFrom<ApInstrument> for KeyPackage {
    type Error = anyhow::Error;

    fn try_from(instrument: ApInstrument) -> Result<Self> {
        let provider = OpenMlsRustCrypto::default();
        let content = instrument.content.ok_or(anyhow!("Content must be Some"))?;
        match instrument.kind {
            ApInstrumentType::MlsKeyPackage => Ok(KeyPackageIn::tls_deserialize_bytes(
                &general_purpose::STANDARD.decode(content)?,
            )?
            .0
            .validate(provider.crypto(), openmls::prelude::ProtocolVersion::Mls10)?),
            _ => Err(anyhow!("ApInstrument not a KeyPackage")),
        }
    }
}

impl From<KeyPackageBundle> for ApInstrument {
    fn from(keys: KeyPackageBundle) -> Self {
        Self {
            kind: ApInstrumentType::MlsKeyPackage,
            content: Some(
                general_purpose::STANDARD
                    .encode(keys.key_package().tls_serialize_detached().unwrap()),
            ),
            ..Default::default()
        }
    }
}

#[serde_as]
#[derive(serde::Serialize, serde::Deserialize, Debug)]
struct MlsMemoryStorageSerializable {
    #[serde_as(as = "Vec<(_, _)>")]
    pub values: HashMap<Vec<u8>, Vec<u8>>,
}

impl From<HashMap<Vec<u8>, Vec<u8>>> for MlsMemoryStorageSerializable {
    fn from(map: HashMap<Vec<u8>, Vec<u8>>) -> Self {
        MlsMemoryStorageSerializable { values: map }
    }
}

type EncryptableMemoryStorage<'a, E, H> = (&'a MemoryStorage, Option<String>, E, H);

impl<'a, E, H> From<EncryptableMemoryStorage<'a, E, H>> for ApInstrument
where
    E: Fn(Vec<u8>) -> Vec<u8>,
    H: Fn(Vec<u8>) -> String,
{
    fn from((storage, mutation_of, encrypt, hash): EncryptableMemoryStorage<'a, E, H>) -> Self {
        let values = &storage.values;
        let s: MlsMemoryStorageSerializable = values.read().unwrap().clone().into();
        let v = serde_json::to_vec(&s).unwrap();
        let e = encrypt(v.clone());
        let h = hash(v);
        let b = general_purpose::STANDARD.encode(&e);

        ApInstrument {
            kind: ApInstrumentType::MlsStorage,
            content: Some(b),
            hash: Some(h),
            mutation_of,
            ..Default::default()
        }
    }
}

impl ApInstrument {
    pub fn to_provider<F>(&self, decrypt: F) -> Result<OpenMlsRustCrypto>
    where
        F: Fn(String) -> Vec<u8>,
    {
        if !self.is_mls_storage() {
            return Err(anyhow!("Instrument is not MlsStorage"));
        }

        let p = OpenMlsRustCrypto::default();
        let c = self
            .content
            .clone()
            .ok_or(anyhow!("content must be Some"))?;
        let d = decrypt(c);
        let v: MlsMemoryStorageSerializable = serde_json::from_slice(&d)?;
        let mut s = p.storage().values.write().unwrap();
        s.clear();
        s.extend(v.values);
        drop(s);

        Ok(p)
    }

    pub fn to_credentials<F>(&self, decrypt: F) -> Result<CredentialKeyPair>
    where
        F: Fn(String) -> Vec<u8>,
    {
        if !self.is_mls_credentials() {
            return Err(anyhow!("Instrument not MlsCredentials"));
        }

        let c = self
            .content
            .clone()
            .ok_or(anyhow!("content must be Some"))?;
        let d = decrypt(c);
        let v: CredentialKeyPair = serde_json::from_slice(&d)?;

        Ok(v)
    }
}
