use crate::error::{DaveError, DaveResult};
use openmls::prelude::*;
use openmls_basic_credential::SignatureKeyPair;
use openmls_rust_crypto::OpenMlsRustCrypto;
use openmls_traits::OpenMlsProvider;
use pyo3::prelude::*;
use std::sync::Arc;
use tls_codec::{Deserialize, Serialize};

pub const CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

/// Internal struct to hold signature key data that can be shared
pub struct SignatureKeyData {
    pub public_key: Vec<u8>,
    pub signature_scheme: SignatureScheme,
}

#[pyclass]
pub struct MlsClient {
    pub(crate) identity: Vec<u8>,
    pub(crate) credential_with_key: CredentialWithKey,
    pub(crate) signature_key_data: SignatureKeyData,
    pub(crate) crypto: Arc<OpenMlsRustCrypto>,
}

impl Clone for MlsClient {
    fn clone(&self) -> Self {
        // Share the crypto provider Arc - this is crucial for key package storage
        MlsClient {
            identity: self.identity.clone(),
            credential_with_key: self.credential_with_key.clone(),
            signature_key_data: SignatureKeyData {
                public_key: self.signature_key_data.public_key.clone(),
                signature_scheme: self.signature_key_data.signature_scheme,
            },
            crypto: self.crypto.clone(), // Share the Arc, don't create new provider
        }
    }
}

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

    #[staticmethod]
    #[pyo3(signature = (identity_str))]
    pub fn from_string(identity_str: &str) -> PyResult<Self> {
        Self::create(identity_str.as_bytes().to_vec()).map_err(Into::into)
    }

    #[getter]
    pub fn identity(&self) -> Vec<u8> {
        self.identity.clone()
    }

    #[getter]
    pub fn identity_str(&self) -> PyResult<String> {
        String::from_utf8(self.identity.clone())
            .map_err(|e| DaveError::Serialization(e.to_string()).into())
    }

    #[getter]
    pub fn public_key(&self) -> Vec<u8> {
        self.signature_key_data.public_key.clone()
    }

    pub fn generate_key_package(&self) -> PyResult<Vec<u8>> {
        self.create_key_package().map_err(Into::into)
    }

    pub fn generate_key_packages(&self, count: usize) -> PyResult<Vec<Vec<u8>>> {
        (0..count)
            .map(|_| self.create_key_package())
            .collect::<DaveResult<Vec<_>>>()
            .map_err(Into::into)
    }

    pub fn __repr__(&self) -> String {
        format!(
            "MlsClient(identity={:?})",
            String::from_utf8_lossy(&self.identity)
        )
    }
}

impl MlsClient {
    pub(crate) fn create(identity: Vec<u8>) -> DaveResult<Self> {
        let crypto = Arc::new(OpenMlsRustCrypto::default());

        let signature_keys = SignatureKeyPair::new(CIPHERSUITE.signature_algorithm())
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        signature_keys
            .store(crypto.storage())
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        let public_key = signature_keys.public().to_vec();
        let signature_scheme = CIPHERSUITE.signature_algorithm();

        let credential = Credential::new(CredentialType::Basic, identity.clone());

        let credential_with_key = CredentialWithKey {
            credential,
            signature_key: signature_keys.public().into(),
        };

        Ok(MlsClient {
            identity,
            credential_with_key,
            signature_key_data: SignatureKeyData {
                public_key,
                signature_scheme,
            },
            crypto,
        })
    }

    pub(crate) fn create_key_package(&self) -> DaveResult<Vec<u8>> {
        let signer = self.get_signer()?;

        let key_package_bundle = KeyPackage::builder()
            .build(
                CIPHERSUITE,
                self.crypto.as_ref(),
                &signer,
                self.credential_with_key.clone(),
            )
            .map_err(|e| DaveError::Mls(format!("{:?}", e)))?;

        key_package_bundle
            .key_package()
            .tls_serialize_detached()
            .map_err(|e| DaveError::Serialization(format!("{:?}", e)))
    }

    pub(crate) fn get_crypto(&self) -> &OpenMlsRustCrypto {
        self.crypto.as_ref()
    }

    pub(crate) fn get_signer(&self) -> DaveResult<SignatureKeyPair> {
        SignatureKeyPair::read(
            self.crypto.storage(),
            &self.signature_key_data.public_key,
            self.signature_key_data.signature_scheme,
        )
        .ok_or_else(|| DaveError::Crypto("Failed to read signature keys from storage".into()))
    }
}

/// Deserialize a KeyPackage from bytes
pub fn deserialize_key_package(data: &[u8]) -> DaveResult<KeyPackageIn> {
    KeyPackageIn::tls_deserialize_exact(data)
        .map_err(|e| DaveError::Serialization(format!("{:?}", e)))
}