chaincraft 0.3.2

//! Shared objects and messages for distributed state management

use uuid::Uuid;
use std::any::Any;
use std::sync::Arc;
use std::fmt::{self, Debug};
use std::collections::HashMap;
use tokio::sync::RwLock;
use serde::{Serialize, Deserialize};
use crate::error::{ChainCraftError, CryptoError, SerializationError, Result};
use sha2::{Sha256, Digest};
use async_trait::async_trait;
use bincode;
use hex;

/// Unique identifier for shared objects
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SharedObjectId(Uuid);

impl SharedObjectId {
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }

    pub fn as_uuid(&self) -> &Uuid {
        &self.0
    }

    pub fn into_uuid(self) -> Uuid {
        self.0
    }
    
    pub fn from_uuid(uuid: Uuid) -> Self {
        Self(uuid)
    }
}

impl fmt::Display for SharedObjectId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Message types for inter-node communication
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum MessageType {
    /// Peer discovery message
    PeerDiscovery,
    /// Request for local peers
    RequestLocalPeers,
    /// Response with local peers
    LocalPeers,
    /// Request for shared object update
    RequestSharedObjectUpdate,
    /// Response with shared object data
    SharedObjectUpdate,
    /// Request to get an object
    Get,
    /// Request to set an object
    Set,
    /// Request to delete an object
    Delete,
    /// Response containing requested data
    Response,
    /// Notification of changes
    Notification,
    /// Heartbeat/ping message
    Heartbeat,
    /// Error response
    Error,
    /// Custom application message
    Custom(String),
}

impl fmt::Display for MessageType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            MessageType::PeerDiscovery => write!(f, "PEER_DISCOVERY"),
            MessageType::RequestLocalPeers => write!(f, "REQUEST_LOCAL_PEERS"),
            MessageType::LocalPeers => write!(f, "LOCAL_PEERS"),
            MessageType::RequestSharedObjectUpdate => write!(f, "REQUEST_SHARED_OBJECT_UPDATE"),
            MessageType::SharedObjectUpdate => write!(f, "SHARED_OBJECT_UPDATE"),
            MessageType::Get => write!(f, "GET"),
            MessageType::Set => write!(f, "SET"),
            MessageType::Delete => write!(f, "DELETE"),
            MessageType::Response => write!(f, "RESPONSE"),
            MessageType::Notification => write!(f, "NOTIFICATION"),
            MessageType::Heartbeat => write!(f, "HEARTBEAT"),
            MessageType::Error => write!(f, "ERROR"),
            MessageType::Custom(name) => write!(f, "{}", name),
        }
    }
}

/// A message that can be shared between nodes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SharedMessage {
    /// Unique message identifier
    pub id: SharedObjectId,
    /// Message type
    pub message_type: MessageType,
    /// Target object ID (if applicable)
    pub target_id: Option<SharedObjectId>,
    /// Message payload
    pub data: serde_json::Value,
    /// Timestamp when message was created
    pub timestamp: chrono::DateTime<chrono::Utc>,
    /// Optional signature for authenticated messages
    #[serde(with = "serde_bytes", default)]
    pub signature: Option<Vec<u8>>,
    /// Hash of the message content
    pub hash: String,
}

impl SharedMessage {
    /// Create a new shared message
    pub fn new(message_type: MessageType, data: serde_json::Value) -> Self {
        let mut message = Self {
            id: SharedObjectId::new(),
            message_type,
            target_id: None,
            data,
            timestamp: chrono::Utc::now(),
            signature: None,
            hash: String::new(),
        };
        message.hash = message.calculate_hash();
        message
    }

    /// Create a new message with a target
    pub fn new_with_target(
        message_type: MessageType,
        target_id: SharedObjectId,
        data: serde_json::Value,
    ) -> Self {
        let mut message = Self {
            id: SharedObjectId::new(),
            message_type,
            target_id: Some(target_id),
            data,
            timestamp: chrono::Utc::now(),
            signature: None,
            hash: String::new(),
        };
        message.hash = message.calculate_hash();
        message
    }

    /// Create a custom message
    pub fn custom<T: Serialize>(message_type: impl Into<String>, data: T) -> Result<Self> {
        let data = serde_json::to_value(data)
            .map_err(|e| ChainCraftError::Serialization(SerializationError::Json(e)))?;
        Ok(Self::new(MessageType::Custom(message_type.into()), data))
    }

    /// Sign this message with the given private key
    pub fn sign(&mut self, private_key: &crate::crypto::PrivateKey) -> Result<()> {
        let message_bytes = self.to_bytes()?;
        let signature = private_key.sign(&message_bytes)?;
        self.signature = Some(signature.to_bytes());
        Ok(())
    }

    /// Verify the signature of this message
    pub fn verify_signature(&self, public_key: &crate::crypto::PublicKey) -> Result<bool> {
        if let Some(sig_bytes) = &self.signature {
            // Create a copy without signature for verification
            let mut message_copy = self.clone();
            message_copy.signature = None;
            let message_bytes = message_copy.to_bytes()?;
            
            // Reconstruct signature
            let signature = if sig_bytes.len() == 64 {
                if public_key.algorithm() == "Ed25519" {
                    let sig_bytes_array: [u8; 64] = sig_bytes.clone().try_into().unwrap();
                    let sig_result = ed25519_dalek::Signature::from_bytes(&sig_bytes_array);
                    if let Ok(sig) = sig_result {
                        crate::crypto::Signature::Ed25519(sig)
                    } else {
                        return Err(ChainCraftError::Crypto(CryptoError::InvalidSignature));
                    }
                } else {
                    let sig_result = k256::ecdsa::Signature::from_bytes(sig_bytes.as_slice().into());
                    if let Ok(sig) = sig_result {
                        crate::crypto::Signature::Secp256k1(sig)
                    } else {
                        return Err(ChainCraftError::Crypto(CryptoError::InvalidSignature));
                    }
                }
            } else {
                return Err(ChainCraftError::Crypto(CryptoError::InvalidSignature));
            };
            
            public_key.verify(&message_bytes, &signature)
        } else {
            Ok(false)
        }
    }

    /// Calculate the hash of this message
    pub fn calculate_hash(&self) -> String {
        let mut hasher = Sha256::new();
        if let Ok(id_bytes) = bincode::serialize(self.id.as_uuid()) {
            hasher.update(&id_bytes);
        }
        hasher.update(self.message_type.to_string().as_bytes());
        hasher.update(self.data.to_string().as_bytes());
        hasher.update(self.timestamp.to_rfc3339().as_bytes());
        hex::encode(hasher.finalize())
    }

    /// Verify the message hash
    pub fn verify_hash(&self) -> bool {
        self.hash == self.calculate_hash()
    }

    /// Get the message size in bytes
    pub fn size(&self) -> usize {
        bincode::serialized_size(self).unwrap_or(0) as usize
    }

    /// Convert message to bytes for signing/verification
    pub fn to_bytes(&self) -> Result<Vec<u8>> {
        bincode::serialize(self).map_err(|e| ChainCraftError::Serialization(SerializationError::Binary(e)))
    }

    /// Create message from bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        bincode::deserialize(bytes).map_err(|e| ChainCraftError::Serialization(SerializationError::Binary(e)))
    }

    /// Serialize to JSON
    pub fn to_json(&self) -> Result<String> {
        serde_json::to_string(self).map_err(|e| ChainCraftError::Serialization(SerializationError::Json(e)))
    }

    /// Deserialize from JSON
    pub fn from_json(json: &str) -> Result<Self> {
        serde_json::from_str(json).map_err(|e| ChainCraftError::Serialization(SerializationError::Json(e)))
    }
}

impl PartialEq for SharedMessage {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id && self.hash == other.hash
    }
}

impl Eq for SharedMessage {}

/// Trait for objects that can be shared and synchronized across nodes
#[async_trait]
pub trait SharedObject: Send + Sync + Debug {
    /// Get the unique identifier for this shared object
    fn id(&self) -> SharedObjectId;

    /// Get the type name of this shared object
    fn type_name(&self) -> &'static str;

    /// Validate if a message is valid for this shared object
    async fn is_valid(&self, message: &SharedMessage) -> Result<bool>;

    /// Process a validated message and update the object state
    async fn add_message(&mut self, message: SharedMessage) -> Result<()>;

    /// Check if this object supports merkleized synchronization
    fn is_merkleized(&self) -> bool;

    /// Get the latest state digest for synchronization
    async fn get_latest_digest(&self) -> Result<String>;

    /// Check if the object has a specific digest
    async fn has_digest(&self, digest: &str) -> Result<bool>;

    /// Validate if a digest is valid
    async fn is_valid_digest(&self, digest: &str) -> Result<bool>;

    /// Add a digest to the object
    async fn add_digest(&mut self, digest: String) -> Result<bool>;

    /// Get messages for gossip protocol
    async fn gossip_messages(&self, digest: Option<&str>) -> Result<Vec<SharedMessage>>;

    /// Get messages since a specific digest
    async fn get_messages_since_digest(&self, digest: &str) -> Result<Vec<SharedMessage>>;

    /// Get the current state as a serializable value
    async fn get_state(&self) -> Result<serde_json::Value>;

    /// Reset the object to a clean state
    async fn reset(&mut self) -> Result<()>;

    /// Serialize the object to JSON
    fn to_json(&self) -> Result<serde_json::Value>;

    /// Update the object from JSON data
    async fn from_json(&mut self, data: serde_json::Value) -> Result<()>;

    /// Clone the object as a trait object
    fn clone_box(&self) -> Box<dyn SharedObject>;

    /// Get a reference to self as Any for downcasting
    fn as_any(&self) -> &dyn Any;

    /// Get a mutable reference to self as Any for downcasting
    fn as_any_mut(&mut self) -> &mut dyn Any;

    /// Called when the object is accessed
    async fn on_access(&mut self) -> Result<()>;

    /// Called when the object is modified
    async fn on_modify(&mut self) -> Result<()>;

    /// Called when the object is deleted
    async fn on_delete(&mut self) -> Result<()>;

    /// Validate the object's current state
    async fn validate(&self) -> Result<bool>;

    /// Get object metadata
    fn metadata(&self) -> HashMap<String, String>;
}

impl Clone for Box<dyn SharedObject> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

/// Digest for tracking shared object state
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StateDigest {
    /// The digest hash
    pub hash: String,
    /// Timestamp when digest was created
    pub timestamp: chrono::DateTime<chrono::Utc>,
    /// Number of messages included in this digest
    pub message_count: u64,
}

impl StateDigest {
    /// Create a new state digest
    pub fn new(hash: String, message_count: u64) -> Self {
        Self {
            hash,
            timestamp: chrono::Utc::now(),
            message_count,
        }
    }

    /// Calculate a digest from messages
    pub fn from_messages(messages: &[SharedMessage]) -> Self {
        let mut hasher = Sha256::new();
        for message in messages {
            hasher.update(message.hash.as_bytes());
        }
        let hash = hex::encode(hasher.finalize());
        Self::new(hash, messages.len() as u64)
    }
}

/// Registry for managing shared objects
pub struct SharedObjectRegistry {
    objects: HashMap<SharedObjectId, Box<dyn SharedObject>>,
    #[allow(dead_code)]
    lock: Arc<RwLock<()>>,
}

impl Default for SharedObjectRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl SharedObjectRegistry {
    /// Create a new empty registry
    pub fn new() -> Self {
        Self {
            objects: HashMap::new(),
            lock: Arc::new(RwLock::new(())),
        }
    }

    /// Register a shared object
    pub fn register(&mut self, object: Box<dyn SharedObject>) -> SharedObjectId {
        let id = object.id();
        self.objects.insert(id.clone(), object);
        id
    }

    /// Get a reference to an object by ID
    pub fn get(&self, id: &SharedObjectId) -> Option<&dyn SharedObject> {
        self.objects.get(id).map(|obj| obj.as_ref())
    }

    /// Get a mutable reference to a shared object by ID
    pub fn get_mut(&mut self, id: &SharedObjectId) -> Option<&mut (dyn SharedObject + '_)> {
        self.objects.get_mut(id).map(|obj| obj.as_mut())
    }

    /// Remove an object from the registry
    pub fn remove(&mut self, id: &SharedObjectId) -> Option<Box<dyn SharedObject>> {
        self.objects.remove(id)
    }

    /// Get all object IDs
    pub fn ids(&self) -> Vec<SharedObjectId> {
        self.objects.keys().cloned().collect()
    }

    /// Get the number of registered objects
    pub fn len(&self) -> usize {
        self.objects.len()
    }

    /// Check if the registry is empty
    pub fn is_empty(&self) -> bool {
        self.objects.is_empty()
    }

    /// Clear all objects from the registry
    pub fn clear(&mut self) {
        self.objects.clear();
    }

    /// Get objects by type
    pub fn get_by_type(&self, type_name: &str) -> Vec<&dyn SharedObject> {
        self.objects
            .values()
            .filter(|obj| obj.type_name() == type_name)
            .map(|obj| obj.as_ref())
            .collect()
    }

    /// Check if an object exists
    pub fn contains(&self, id: &SharedObjectId) -> bool {
        self.objects.contains_key(id)
    }
}

impl Debug for SharedObjectRegistry {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SharedObjectRegistry")
            .field("object_count", &self.objects.len())
            .field("object_ids", &self.ids())
            .finish()
    }
}