use std::collections::HashMap;
use std::path::Path;

use mls_rs::crypto::{SignaturePublicKey, SignatureSecretKey};

use mls_rs::group::ReceivedMessage;
use mls_rs::identity::basic::{BasicCredential, BasicIdentityProvider};
use mls_rs::identity::SigningIdentity;
use mls_rs::{CipherSuite, CipherSuiteProvider, Client, CryptoProvider, ExtensionList, MlsMessage};
use mls_rs_crypto_rustcrypto::RustCryptoProvider;
use mls_rs_provider_sqlite::connection_strategy::FileConnectionStrategy;
use mls_rs_provider_sqlite::storage::{
    SqLiteGroupStateStorage, SqLiteKeyPackageStorage, SqLitePreSharedKeyStorage,
};
use mls_rs_provider_sqlite::SqLiteDataStorageEngine;

const MLS_CIPHERSUITE: CipherSuite = CipherSuite::CURVE25519_AES128;

/// Result of processing an incoming MLS message.
#[derive(Debug)]
pub enum ProcessedMessage {
    /// Decrypted application data.
    Application(Vec<u8>),
    /// Commit processed, epoch advanced.
    Commit,
    /// Other MLS message type (proposal, etc.)
    Other,
}

/// Errors from MLS operations.
#[derive(Debug, thiserror::Error)]
pub enum MlsEngineError {
    #[error("MLS client error: {0}")]
    Client(String),

    #[error("MLS group error: {0}")]
    Group(String),

    #[error("group not found for id: {0}")]
    GroupNotFound(String),

    #[error("message serialization error: {0}")]
    Serialization(String),

    #[error("storage error: {0}")]
    Storage(String),

    #[error("unexpected message type")]
    UnexpectedMessageType,
}

/// Concrete mls-rs config type after building with SQLite + RustCrypto + BasicIdentity.
type SdkMlsConfig = mls_rs::client_builder::WithCryptoProvider<
    RustCryptoProvider,
    mls_rs::client_builder::WithIdentityProvider<
        BasicIdentityProvider,
        mls_rs::client_builder::WithPskStore<
            SqLitePreSharedKeyStorage,
            mls_rs::client_builder::WithKeyPackageRepo<
                SqLiteKeyPackageStorage,
                mls_rs::client_builder::WithGroupStateStorage<
                    SqLiteGroupStateStorage,
                    mls_rs::client_builder::BaseConfig,
                >,
            >,
        >,
    >,
>;

type MlsClient = Client<SdkMlsConfig>;
type MlsGroup = mls_rs::Group<SdkMlsConfig>;

/// Wraps mls-rs with persistent SQLite storage for MLS group operations.
///
/// # Example
///
/// ```no_run
/// use skytale_sdk::mls_engine::MlsEngine;
/// let engine = MlsEngine::new(
///     std::path::Path::new("/tmp/mls-data"),
///     b"agent-1".to_vec(),
/// ).unwrap();
/// ```
pub struct MlsEngine {
    /// In-memory cache of active groups keyed by hex-encoded group_id.
    groups: HashMap<String, MlsGroup>,
    /// The mls-rs client used to create/join/load groups.
    client: MlsClient,
}

impl MlsEngine {
    /// Create a new MLS engine with state persisted to the given directory.
    pub fn new(data_dir: &Path, identity: Vec<u8>) -> Result<Self, MlsEngineError> {
        // Check for symlink before create_dir_all to prevent a TOCTOU attack
        // where an attacker replaces the directory with a symlink.
        if data_dir.exists() {
            #[cfg(unix)]
            {
                let metadata = std::fs::symlink_metadata(data_dir).map_err(|e| {
                    MlsEngineError::Client(format!("failed to check data dir: {e}"))
                })?;
                if metadata.file_type().is_symlink() {
                    return Err(MlsEngineError::Client(
                        "data directory is a symlink, refusing to use".into(),
                    ));
                }
            }
        }

        std::fs::create_dir_all(data_dir).map_err(|e| MlsEngineError::Storage(e.to_string()))?;

        // Restrict data directory to owner-only access (rwx------).
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(data_dir, std::fs::Permissions::from_mode(0o700))
                .map_err(|e| MlsEngineError::Storage(e.to_string()))?;
        }

        let db_path = data_dir.join("mls.db");
        let conn_strategy = FileConnectionStrategy::new(&db_path);
        // NOTE: mls-rs key types do not implement ZeroizeOnDrop — key material
        // in the SQLite database and in-memory mls-rs structures is not
        // automatically zeroized when dropped.
        let storage_engine = SqLiteDataStorageEngine::new(conn_strategy)
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;

        // Restrict db file to owner-only read/write (rw-------).
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            if db_path.exists() {
                std::fs::set_permissions(&db_path, std::fs::Permissions::from_mode(0o600))
                    .map_err(|e| MlsEngineError::Storage(e.to_string()))?;
            }
        }

        let group_state = storage_engine
            .group_state_storage()
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;
        let key_package = storage_engine
            .key_package_storage()
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;
        let psk = storage_engine
            .pre_shared_key_storage()
            .map_err(|e| MlsEngineError::Storage(format!("{e:?}")))?;

        let crypto = RustCryptoProvider::default();
        let cs = crypto
            .cipher_suite_provider(MLS_CIPHERSUITE)
            .ok_or_else(|| MlsEngineError::Client("unsupported cipher suite".into()))?;

        // Persist the signing key so the agent keeps the same MLS identity
        // across restarts instead of generating a new one each time.
        let key_path = data_dir.join("signing_key.bin");
        let (secret_key, public_key) = if key_path.exists() {
            let key_bytes = std::fs::read(&key_path)
                .map_err(|e| MlsEngineError::Client(format!("failed to read signing key: {e}")))?;
            deserialize_signing_key(&key_bytes)?
        } else {
            let keys = cs
                .signature_key_generate()
                .map_err(|e| MlsEngineError::Client(format!("{e:?}")))?;
            let key_bytes = serialize_signing_key(&keys);
            std::fs::write(&key_path, &key_bytes)
                .map_err(|e| MlsEngineError::Client(format!("failed to write signing key: {e}")))?;
            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                std::fs::set_permissions(&key_path, std::fs::Permissions::from_mode(0o600))
                    .map_err(|e| {
                        MlsEngineError::Client(format!("failed to set key permissions: {e}"))
                    })?;
            }
            keys
        };

        let basic_cred = BasicCredential::new(identity);
        let signing_id = SigningIdentity::new(basic_cred.into_credential(), public_key);

        let client = Client::builder()
            .group_state_storage(group_state)
            .key_package_repo(key_package)
            .psk_store(psk)
            .identity_provider(BasicIdentityProvider)
            .crypto_provider(RustCryptoProvider::default())
            .signing_identity(signing_id, secret_key, MLS_CIPHERSUITE)
            .build();

        Ok(Self {
            groups: HashMap::new(),
            client,
        })
    }

    /// Create a new MLS group. Returns the group_id bytes.
    pub fn create_group(&mut self, group_id: &[u8]) -> Result<Vec<u8>, MlsEngineError> {
        let group = self
            .client
            .create_group_with_id(
                group_id.to_vec(),
                ExtensionList::default(),
                Default::default(),
                None,
            )
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        let id = hex::encode(group.group_id());
        self.groups.insert(id, group);
        Ok(group_id.to_vec())
    }

    /// Generate a KeyPackage for others to add us to their group.
    pub fn generate_key_package(&self) -> Result<Vec<u8>, MlsEngineError> {
        let kp_msg = self
            .client
            .generate_key_package_message(Default::default(), Default::default(), None)
            .map_err(|e| MlsEngineError::Client(format!("{e:?}")))?;
        kp_msg
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))
    }

    /// Add a member to a group we own. Returns `(commit_bytes, welcome_bytes)`.
    pub fn add_member(
        &mut self,
        group_id: &[u8],
        key_package_bytes: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>), MlsEngineError> {
        let group = self.get_group_mut(group_id)?;

        let kp_msg = MlsMessage::from_bytes(key_package_bytes)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        let commit_output = group
            .commit_builder()
            .add_member(kp_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
            .build()
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        group
            .apply_pending_commit()
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        let commit_bytes = commit_output
            .commit_message
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        let welcome_bytes = commit_output
            .welcome_messages
            .first()
            .ok_or_else(|| MlsEngineError::Group("no welcome message produced".into()))?
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        Ok((commit_bytes, welcome_bytes))
    }

    /// Join a group from a Welcome message. Returns the group_id.
    pub fn join_from_welcome(&mut self, welcome_bytes: &[u8]) -> Result<Vec<u8>, MlsEngineError> {
        let welcome_msg = MlsMessage::from_bytes(welcome_bytes)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        let (group, _info) = self
            .client
            .join_group(None, &welcome_msg, None)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;

        let group_id = group.group_id().to_vec();
        let id = hex::encode(&group_id);
        self.groups.insert(id, group);
        Ok(group_id)
    }

    /// Encrypt a message for a group.
    #[tracing::instrument(skip(self, group_id, plaintext), fields(payload.size = plaintext.len()))]
    pub fn encrypt(
        &mut self,
        group_id: &[u8],
        plaintext: &[u8],
    ) -> Result<Vec<u8>, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = group
            .encrypt_application_message(plaintext, vec![])
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;
        mls_msg
            .to_bytes()
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))
    }

    /// Decrypt a message from a group.
    #[tracing::instrument(skip(self, group_id, ciphertext), fields(payload.size = ciphertext.len()))]
    pub fn decrypt(
        &mut self,
        group_id: &[u8],
        ciphertext: &[u8],
    ) -> Result<Vec<u8>, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = MlsMessage::from_bytes(ciphertext)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        match group
            .process_incoming_message(mls_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
        {
            ReceivedMessage::ApplicationMessage(app_msg) => Ok(app_msg.data().to_vec()),
            _ => Err(MlsEngineError::UnexpectedMessageType),
        }
    }

    /// Process an incoming commit (epoch advancement).
    pub fn process_commit(
        &mut self,
        group_id: &[u8],
        commit_bytes: &[u8],
    ) -> Result<(), MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = MlsMessage::from_bytes(commit_bytes)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        match group
            .process_incoming_message(mls_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
        {
            ReceivedMessage::Commit(_) => Ok(()),
            _ => Err(MlsEngineError::UnexpectedMessageType),
        }
    }

    /// Process any incoming MLS message, letting mls-rs determine the type.
    ///
    /// Use this instead of `decrypt`/`process_commit` when the message type is
    /// unknown (e.g., the relay labels everything as `Msg`).
    pub fn process_incoming(
        &mut self,
        group_id: &[u8],
        ciphertext: &[u8],
    ) -> Result<ProcessedMessage, MlsEngineError> {
        let group = self.get_group_mut(group_id)?;
        let mls_msg = MlsMessage::from_bytes(ciphertext)
            .map_err(|e| MlsEngineError::Serialization(format!("{e:?}")))?;

        match group
            .process_incoming_message(mls_msg)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?
        {
            ReceivedMessage::ApplicationMessage(app_msg) => {
                Ok(ProcessedMessage::Application(app_msg.data().to_vec()))
            }
            ReceivedMessage::Commit(_) => Ok(ProcessedMessage::Commit),
            _ => Ok(ProcessedMessage::Other),
        }
    }

    /// Load a group from persistent storage by its ID.
    pub fn load_group(&mut self, group_id: &[u8]) -> Result<(), MlsEngineError> {
        let group = self
            .client
            .load_group(group_id)
            .map_err(|e| MlsEngineError::Group(format!("{e:?}")))?;
        let id = hex::encode(group_id);
        self.groups.insert(id, group);
        Ok(())
    }

    /// Remove a group from the in-memory cache.
    ///
    /// Returns `true` if the group was present, `false` otherwise.
    pub fn remove_group(&mut self, channel_id: &[u8]) -> bool {
        let id = hex::encode(channel_id);
        self.groups.remove(&id).is_some()
    }

    fn get_group_mut(&mut self, group_id: &[u8]) -> Result<&mut MlsGroup, MlsEngineError> {
        let id = hex::encode(group_id);
        if !self.groups.contains_key(&id) {
            if self.groups.len() >= MAX_ACTIVE_GROUPS {
                return Err(MlsEngineError::Client("too many active groups".into()));
            }
            // Try to load from persistent storage
            if let Ok(group) = self.client.load_group(group_id) {
                self.groups.insert(id.clone(), group);
            }
        }
        self.groups
            .get_mut(&id)
            .ok_or(MlsEngineError::GroupNotFound(id))
    }
}

const MAX_ACTIVE_GROUPS: usize = 1000;

/// Serialize a signing key pair as `[secret_len as u32 LE][secret_bytes][public_bytes]`.
fn serialize_signing_key(keys: &(SignatureSecretKey, SignaturePublicKey)) -> Vec<u8> {
    let secret_bytes: &[u8] = keys.0.as_ref();
    let public_bytes: &[u8] = keys.1.as_ref();
    let mut buf = Vec::with_capacity(4 + secret_bytes.len() + public_bytes.len());
    buf.extend_from_slice(&(secret_bytes.len() as u32).to_le_bytes());
    buf.extend_from_slice(secret_bytes);
    buf.extend_from_slice(public_bytes);
    buf
}

/// Deserialize a signing key pair from the format produced by `serialize_signing_key`.
fn deserialize_signing_key(
    data: &[u8],
) -> Result<(SignatureSecretKey, SignaturePublicKey), MlsEngineError> {
    if data.len() < 4 {
        return Err(MlsEngineError::Client("signing key file too short".into()));
    }
    let secret_len = u32::from_le_bytes(
        data[..4]
            .try_into()
            .map_err(|_| MlsEngineError::Client("invalid key length prefix".into()))?,
    ) as usize;
    if data.len() < 4 + secret_len {
        return Err(MlsEngineError::Client("signing key file truncated".into()));
    }
    let secret_bytes = data[4..4 + secret_len].to_vec();
    let public_bytes = data[4 + secret_len..].to_vec();
    Ok((
        SignatureSecretKey::from(secret_bytes),
        SignaturePublicKey::from(public_bytes),
    ))
}
