use crate::client::{deserialize_key_package, MlsClient, CIPHERSUITE};
use crate::error::{DaveError, DaveResult};
use pyo3::prelude::*;
use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;

use openmls::prelude::*;
use openmls_traits::OpenMlsProvider;
use tls_codec::{Deserialize, Serialize};

/// DAVE-specific constants
pub const DAVE_PROTOCOL_VERSION: u16 = 1;
pub const SFRAME_LABEL: &str = "Discord Secure Frames v0";
pub const SENDER_KEY_LABEL: &str = "Discord Sender Key v0";
pub const MEDIA_KEY_LENGTH: usize = 32;

#[pyclass]
pub struct MlsGroupWrapper {
    group: Arc<RwLock<MlsGroup>>,
    client: MlsClient,
    #[allow(dead_code)]
    member_indices: Arc<RwLock<HashMap<Vec<u8>, LeafNodeIndex>>>,
}

#[pyclass]
#[derive(Clone)]
pub struct CommitResult {
    #[pyo3(get)]
    pub commit_message: Vec<u8>,
    #[pyo3(get)]
    pub welcome_message: Option<Vec<u8>>,
    #[pyo3(get)]
    pub group_info: Option<Vec<u8>>,
}

#[pyclass]
#[derive(Clone)]
pub struct ProcessedMessage {
    #[pyo3(get)]
    pub message_type: String,
    #[pyo3(get)]
    pub sender_identity: Option<Vec<u8>>,
    #[pyo3(get)]
    pub epoch: u64,
}

#[pymethods]
impl MlsGroupWrapper {
    #[new]
    #[pyo3(signature = (client, group_id))]
    pub fn new(client: &MlsClient, group_id: Vec<u8>) -> PyResult<Self> {
        Self::create(client.clone(), group_id).map_err(Into::into)
    }

    #[staticmethod]
    #[pyo3(signature = (client, welcome_bytes))]
    pub fn from_welcome(client: &MlsClient, welcome_bytes: Vec<u8>) -> PyResult<Self> {
        Self::join_from_welcome(client.clone(), &welcome_bytes).map_err(Into::into)
    }

    #[getter]
    pub fn group_id(&self) -> Vec<u8> {
        self.group.read().group_id().as_slice().to_vec()
    }

    #[getter]
    pub fn epoch(&self) -> u64 {
        self.group.read().epoch().as_u64()
    }

    #[getter]
    pub fn member_count(&self) -> usize {
        self.group.read().members().count()
    }

    pub fn members(&self) -> Vec<Vec<u8>> {
        self.group
            .read()
            .members()
            .map(|m| get_credential_identity(&m.credential))
            .collect()
    }

    pub fn add_member(&mut self, key_package_bytes: Vec<u8>) -> PyResult<CommitResult> {
        self.add_member_internal(&key_package_bytes)
            .map_err(Into::into)
    }

    pub fn add_members(&mut self, key_packages: Vec<Vec<u8>>) -> PyResult<CommitResult> {
        self.add_members_internal(&key_packages).map_err(Into::into)
    }

    pub fn remove_member(&mut self, identity: Vec<u8>) -> PyResult<CommitResult> {
        self.remove_member_internal(&identity).map_err(Into::into)
    }

    pub fn self_update(&mut self) -> PyResult<CommitResult> {
        self.self_update_internal().map_err(Into::into)
    }

    pub fn process_message(&mut self, message_bytes: Vec<u8>) -> PyResult<ProcessedMessage> {
        self.process_message_internal(&message_bytes)
            .map_err(Into::into)
    }

    pub fn export_secret(&self, label: &str, context: Vec<u8>, length: usize) -> PyResult<Vec<u8>> {
        self.group
            .read()
            .export_secret(self.client.get_crypto().crypto(), label, &context, length)
            .map_err(|e| DaveError::Mls(format!("{:?}", e)).into())
    }

    pub fn derive_media_key(&self) -> PyResult<Vec<u8>> {
        self.export_secret(SFRAME_LABEL, vec![], MEDIA_KEY_LENGTH)
    }

    pub fn derive_sender_key(&self, sender_id: u64) -> PyResult<Vec<u8>> {
        let context = sender_id.to_be_bytes().to_vec();
        self.export_secret(SENDER_KEY_LABEL, context, MEDIA_KEY_LENGTH)
    }

    pub fn get_epoch_authenticator(&self) -> Vec<u8> {
        self.group
            .read()
            .epoch_authenticator()
            .as_slice()
            .to_vec()
    }

    pub fn merge_pending_commit(&mut self) -> PyResult<()> {
        self.group
            .write()
            .merge_pending_commit(self.client.get_crypto())
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;
        Ok(())
    }

    pub fn export_ratchet_tree(&self) -> PyResult<Vec<u8>> {
        self.group
            .read()
            .export_ratchet_tree()
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)).into())
    }

    pub fn __repr__(&self) -> String {
        let group = self.group.read();
        format!(
            "MlsGroup(id={}, epoch={}, members={})",
            hex::encode(group.group_id().as_slice()),
            group.epoch().as_u64(),
            group.members().count()
        )
    }
}

/// Extract identity bytes from a credential
fn get_credential_identity(credential: &Credential) -> Vec<u8> {
    match credential.credential_type() {
        CredentialType::Basic => {
            credential.serialized_content().to_vec()
        }
        _ => {
            credential.serialized_content().to_vec()
        }
    }
}

impl MlsGroupWrapper {
    fn create(client: MlsClient, group_id: Vec<u8>) -> DaveResult<Self> {
        let signer = client.get_signer()?;

        let group_config = MlsGroupCreateConfig::builder()
            .use_ratchet_tree_extension(true)
            .ciphersuite(CIPHERSUITE)
            .build();

        let group = MlsGroup::new_with_group_id(
            client.get_crypto(),
            &signer,
            &group_config,
            GroupId::from_slice(&group_id),
            client.credential_with_key.clone(),
        )
        .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        Ok(MlsGroupWrapper {
            group: Arc::new(RwLock::new(group)),
            client,
            member_indices: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    fn join_from_welcome(client: MlsClient, welcome_bytes: &[u8]) -> DaveResult<Self> {
        let mls_message_in = MlsMessageIn::tls_deserialize_exact(welcome_bytes)
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        let welcome = match mls_message_in.extract() {
            MlsMessageBodyIn::Welcome(welcome) => welcome,
            _ => return Err(DaveError::InvalidState("Not a welcome message".into())),
        };

        let group_config = MlsGroupJoinConfig::builder()
            .use_ratchet_tree_extension(true)
            .build();

        let group = StagedWelcome::new_from_welcome(
            client.get_crypto(),
            &group_config,
            welcome,
            None,
        )
        .map_err(|e| DaveError::Mls(format!("{:?}", e)))?
        .into_group(client.get_crypto())
        .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        Ok(MlsGroupWrapper {
            group: Arc::new(RwLock::new(group)),
            client,
            member_indices: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    fn add_member_internal(&mut self, key_package_bytes: &[u8]) -> DaveResult<CommitResult> {
        let key_package_in = deserialize_key_package(key_package_bytes)?;
        let crypto = self.client.get_crypto();
        let key_package: KeyPackage = key_package_in
            .validate(crypto.crypto(), ProtocolVersion::default())
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        let signer = self.client.get_signer()?;

        let (mls_message_out, welcome_out, _group_info) = self
            .group
            .write()
            .add_members(
                crypto,
                &signer,
                &[key_package],
            )
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        let commit_bytes: Vec<u8> = mls_message_out
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        let welcome_bytes: Vec<u8> = welcome_out
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        Ok(CommitResult {
            commit_message: commit_bytes,
            welcome_message: Some(welcome_bytes),
            group_info: None,
        })
    }

    fn add_members_internal(&mut self, key_packages: &[Vec<u8>]) -> DaveResult<CommitResult> {
        let crypto = self.client.get_crypto();
        let validated: Vec<KeyPackage> = key_packages
            .iter()
            .map(|kp| {
                let kp_in = deserialize_key_package(kp)?;
                kp_in
                    .validate(crypto.crypto(), ProtocolVersion::default())
                    .map_err(|e| DaveError::Mls(format!("{:?}", e)))
            })
            .collect::<DaveResult<Vec<_>>>()?;

        let signer = self.client.get_signer()?;

        let (mls_message_out, welcome_out, _group_info) = self
            .group
            .write()
            .add_members(
                crypto,
                &signer,
                &validated,
            )
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        let commit_bytes: Vec<u8> = mls_message_out
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        let welcome_bytes: Vec<u8> = welcome_out
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        Ok(CommitResult {
            commit_message: commit_bytes,
            welcome_message: Some(welcome_bytes),
            group_info: None,
        })
    }

    fn remove_member_internal(&mut self, identity: &[u8]) -> DaveResult<CommitResult> {
        let group = self.group.read();
        let member = group
            .members()
            .find(|m| get_credential_identity(&m.credential) == identity)
            .ok_or_else(|| DaveError::InvalidState("Member not found".into()))?;

        let leaf_index = member.index;
        drop(group);

        let signer = self.client.get_signer()?;

        let (mls_message_out, _welcome_out, _group_info) = self
            .group
            .write()
            .remove_members(
                self.client.get_crypto(),
                &signer,
                &[leaf_index],
            )
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        let commit_bytes: Vec<u8> = mls_message_out
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        Ok(CommitResult {
            commit_message: commit_bytes,
            welcome_message: None,
            group_info: None,
        })
    }

    fn self_update_internal(&mut self) -> DaveResult<CommitResult> {
        let signer = self.client.get_signer()?;

        let commit_bundle = self
            .group
            .write()
            .self_update(
                self.client.get_crypto(),
                &signer,
                LeafNodeParameters::default(),
            )
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        let commit_bytes: Vec<u8> = commit_bundle
            .commit()
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        Ok(CommitResult {
            commit_message: commit_bytes,
            welcome_message: None,
            group_info: None,
        })
    }

    fn process_message_internal(&mut self, message_bytes: &[u8]) -> DaveResult<ProcessedMessage> {
        let mls_message_in = MlsMessageIn::tls_deserialize_exact(message_bytes)
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))?;

        let protocol_message: ProtocolMessage = mls_message_in
            .into_protocol_message()
            .ok_or_else(|| DaveError::InvalidState("Not a protocol message".into()))?;

        let processed = self
            .group
            .write()
            .process_message(self.client.get_crypto(), protocol_message)
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        let sender_identity: Option<Vec<u8>> = match processed.sender() {
            Sender::Member(leaf_index) => {
                let group = self.group.read();
                let identity = group
                    .members()
                    .find(|m| &m.index == leaf_index)
                    .map(|m| get_credential_identity(&m.credential));
                identity
            }
            _ => None,
        };

        let (message_type, epoch) = match processed.into_content() {
            ProcessedMessageContent::ApplicationMessage(_) => {
                ("application".to_string(), self.group.read().epoch().as_u64())
            }
            ProcessedMessageContent::ProposalMessage(_) => {
                ("proposal".to_string(), self.group.read().epoch().as_u64())
            }
            ProcessedMessageContent::StagedCommitMessage(staged) => {
                let epoch = self.group.read().epoch().as_u64() + 1;
                self.group
                    .write()
                    .merge_staged_commit(self.client.get_crypto(), *staged)
                    .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;
                ("commit".to_string(), epoch)
            }
            ProcessedMessageContent::ExternalJoinProposalMessage(_) => {
                ("external_join".to_string(), self.group.read().epoch().as_u64())
            }
        };

        Ok(ProcessedMessage {
            message_type,
            sender_identity,
            epoch,
        })
    }
}