herodb/src/rpc.rs

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use jsonrpsee::{core::RpcResult, proc_macros::rpc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::server::Server;
use crate::options::DBOption;

/// Database backend types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BackendType {
    Redb,
    Sled,
    // Future: InMemory, Custom(String)
}

/// Database configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabaseConfig {
    pub name: Option<String>,
    pub storage_path: Option<String>,
    pub max_size: Option<u64>,
    pub redis_version: Option<String>,
}

/// Database information returned by metadata queries
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabaseInfo {
    pub id: u64,
    pub name: Option<String>,
    pub backend: BackendType,
    pub encrypted: bool,
    pub redis_version: Option<String>,
    pub storage_path: Option<String>,
    pub size_on_disk: Option<u64>,
    pub key_count: Option<u64>,
    pub created_at: u64,
    pub last_access: Option<u64>,
}

/// Access permissions for database keys
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Permissions {
    Read,
    ReadWrite,
}

/// Access key information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccessKey {
    pub hash: String,
    pub permissions: Permissions,
    pub created_at: u64,
}

/// Database metadata containing access keys
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabaseMeta {
    pub public: bool,
    pub keys: HashMap<String, AccessKey>,
}

/// Access key information returned by RPC
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccessKeyInfo {
    pub hash: String,
    pub permissions: Permissions,
    pub created_at: u64,
}

/// Hash a plaintext key using SHA-256
pub fn hash_key(key: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(key.as_bytes());
    format!("{:x}", hasher.finalize())
}

/// RPC trait for HeroDB management
#[rpc(server, client, namespace = "herodb")]
pub trait Rpc {
    /// Create a new database with specified configuration
    #[method(name = "createDatabase")]
    async fn create_database(
        &self,
        backend: BackendType,
        config: DatabaseConfig,
        encryption_key: Option<String>,
    ) -> RpcResult<u64>;

    /// Set encryption for an existing database (write-only key)
    #[method(name = "setEncryption")]
    async fn set_encryption(&self, db_id: u64, encryption_key: String) -> RpcResult<bool>;

    /// List all managed databases
    #[method(name = "listDatabases")]
    async fn list_databases(&self) -> RpcResult<Vec<DatabaseInfo>>;

    /// Get detailed information about a specific database
    #[method(name = "getDatabaseInfo")]
    async fn get_database_info(&self, db_id: u64) -> RpcResult<DatabaseInfo>;

    /// Delete a database
    #[method(name = "deleteDatabase")]
    async fn delete_database(&self, db_id: u64) -> RpcResult<bool>;

    /// Get server statistics
    #[method(name = "getServerStats")]
    async fn get_server_stats(&self) -> RpcResult<HashMap<String, serde_json::Value>>;

    /// Add an access key to a database
    #[method(name = "addAccessKey")]
    async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult<bool>;

    /// Delete an access key from a database
    #[method(name = "deleteAccessKey")]
    async fn delete_access_key(&self, db_id: u64, key_hash: String) -> RpcResult<bool>;

    /// List all access keys for a database
    #[method(name = "listAccessKeys")]
    async fn list_access_keys(&self, db_id: u64) -> RpcResult<Vec<AccessKeyInfo>>;

    /// Set database public/private status
    #[method(name = "setDatabasePublic")]
    async fn set_database_public(&self, db_id: u64, public: bool) -> RpcResult<bool>;
}

/// RPC Server implementation
pub struct RpcServerImpl {
    /// Base directory for database files
    base_dir: String,
    /// Managed database servers
    servers: Arc<RwLock<HashMap<u64, Arc<Server>>>>,
    /// Next unencrypted database ID to assign
    next_unencrypted_id: Arc<RwLock<u64>>,
    /// Next encrypted database ID to assign
    next_encrypted_id: Arc<RwLock<u64>>,
    /// Default backend type
    backend: crate::options::BackendType,
    /// Encryption keys for databases
    encryption_keys: Arc<RwLock<HashMap<u64, Option<String>>>>,
}

impl RpcServerImpl {
    /// Create a new RPC server instance
    pub fn new(base_dir: String, backend: crate::options::BackendType) -> Self {
        Self {
            base_dir,
            servers: Arc::new(RwLock::new(HashMap::new())),
            next_unencrypted_id: Arc::new(RwLock::new(0)),
            next_encrypted_id: Arc::new(RwLock::new(10)),
            backend,
            encryption_keys: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Get or create a server instance for the given database ID
    async fn get_or_create_server(&self, db_id: u64) -> Result<Arc<Server>, jsonrpsee::types::ErrorObjectOwned> {
        // Check if server already exists
        {
            let servers = self.servers.read().await;
            if let Some(server) = servers.get(&db_id) {
                return Ok(server.clone());
            }
        }

        // Check if database file exists
        let db_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}.db", db_id));
        if !db_path.exists() {
            return Err(jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Database {} not found", db_id),
                None::<()>
            ));
        }

        // Create server instance with default options
        let db_option = DBOption {
            dir: self.base_dir.clone(),
            port: 0, // Not used for RPC-managed databases
            debug: false,
            encryption_key: None,
            encrypt: false,
            backend: self.backend.clone(),
        };

        let mut server = Server::new(db_option).await;

        // Set the selected database to the db_id for proper file naming
        server.selected_db = db_id;

        // Store the server
        let mut servers = self.servers.write().await;
        servers.insert(db_id, Arc::new(server.clone()));

        Ok(Arc::new(server))
    }

    /// Discover existing database files in the base directory
    async fn discover_databases(&self) -> Vec<u64> {
        let mut db_ids = Vec::new();

        if let Ok(entries) = std::fs::read_dir(&self.base_dir) {
            for entry in entries.flatten() {
                if let Ok(file_name) = entry.file_name().into_string() {
                    // Check if it's a database file (ends with .db)
                    if file_name.ends_with(".db") {
                        // Extract database ID from filename (e.g., "11.db" -> 11)
                        if let Some(id_str) = file_name.strip_suffix(".db") {
                            if let Ok(db_id) = id_str.parse::<u64>() {
                                db_ids.push(db_id);
                            }
                        }
                    }
                }
            }
        }

        db_ids
    }

    /// Get the next available database ID
    async fn get_next_db_id(&self, is_encrypted: bool) -> u64 {
        if is_encrypted {
            let mut id = self.next_encrypted_id.write().await;
            let current_id = *id;
            *id += 1;
            current_id
        } else {
            let mut id = self.next_unencrypted_id.write().await;
            let current_id = *id;
            *id += 1;
            current_id
        }
    }

    /// Load database metadata from file (static version)
    pub async fn load_meta_static(base_dir: &str, db_id: u64) -> Result<DatabaseMeta, jsonrpsee::types::ErrorObjectOwned> {
        let meta_path = std::path::PathBuf::from(base_dir).join(format!("{}_meta.json", db_id));

        // If meta file doesn't exist, create and persist default
        if !meta_path.exists() {
            let default_meta = DatabaseMeta {
                public: true,
                keys: HashMap::new(),
            };
            // Persist default metadata to disk
            Self::save_meta_static(base_dir, db_id, &default_meta).await?;
            return Ok(default_meta);
        }

        // Read file as UTF-8 JSON
        let json_str = std::fs::read_to_string(&meta_path)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to read meta file: {}", e),
                None::<()>
            ))?;

        serde_json::from_str(&json_str)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to parse meta JSON: {}", e),
                None::<()>
            ))
    }

    /// Load database metadata from file
    async fn load_meta(&self, db_id: u64) -> Result<DatabaseMeta, jsonrpsee::types::ErrorObjectOwned> {
        let meta_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}_meta.json", db_id));

        // If meta file doesn't exist, create and persist default
        if !meta_path.exists() {
            let default_meta = DatabaseMeta {
                public: true,
                keys: HashMap::new(),
            };
            self.save_meta(db_id, &default_meta).await?;
            return Ok(default_meta);
        }

        // Read file as UTF-8 JSON (meta files are always plain JSON)
        let json_str = std::fs::read_to_string(&meta_path)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to read meta file: {}", e),
                None::<()>
            ))?;

        serde_json::from_str(&json_str)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to parse meta JSON: {}", e),
                None::<()>
            ))
    }

    /// Save database metadata to file (static version)
    pub async fn save_meta_static(base_dir: &str, db_id: u64, meta: &DatabaseMeta) -> Result<(), jsonrpsee::types::ErrorObjectOwned> {
        let meta_path = std::path::PathBuf::from(base_dir).join(format!("{}_meta.json", db_id));

        let json_str = serde_json::to_string_pretty(meta)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to serialize meta: {}", e),
                None::<()>
            ))?;

        std::fs::write(&meta_path, json_str)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to write meta file: {}", e),
                None::<()>
            ))?;

        Ok(())
    }

    /// Save database metadata to file
    async fn save_meta(&self, db_id: u64, meta: &DatabaseMeta) -> Result<(), jsonrpsee::types::ErrorObjectOwned> {
        let meta_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}_meta.json", db_id));

        let json_str = serde_json::to_string_pretty(meta)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to serialize meta: {}", e),
                None::<()>
            ))?;

        // Meta files are always stored as plain JSON (even when data DB is encrypted)
        std::fs::write(&meta_path, json_str)
            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                format!("Failed to write meta file: {}", e),
                None::<()>
            ))?;

        Ok(())
    }
}

#[jsonrpsee::core::async_trait]
impl RpcServer for RpcServerImpl {
    async fn create_database(
        &self,
        backend: BackendType,
        config: DatabaseConfig,
        encryption_key: Option<String>,
    ) -> RpcResult<u64> {
        let db_id = self.get_next_db_id(encryption_key.is_some()).await;

        // Handle both Redb and Sled backends
        match backend {
            BackendType::Redb | BackendType::Sled => {
                // Create database directory
                let db_dir = if let Some(path) = &config.storage_path {
                    std::path::PathBuf::from(path)
                } else {
                    std::path::PathBuf::from(&self.base_dir).join(format!("rpc_db_{}", db_id))
                };

                // Ensure directory exists
                std::fs::create_dir_all(&db_dir)
                    .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
                        -32000,
                        format!("Failed to create directory: {}", e),
                        None::<()>
                    ))?;

                // Create DB options
                let encrypt = encryption_key.is_some();
                let option = DBOption {
                    dir: db_dir.to_string_lossy().to_string(),
                    port: 0, // Not used for RPC-managed databases
                    debug: false,
                    encryption_key: encryption_key.clone(),
                    encrypt,
                    backend: match backend {
                        BackendType::Redb => crate::options::BackendType::Redb,
                        BackendType::Sled => crate::options::BackendType::Sled,
                    },
                };

                // Create server instance
                let mut server = Server::new(option).await;

                // Set the selected database to the db_id for proper file naming
                server.selected_db = db_id;

                // Initialize the storage to create the database file
                let _ = server.current_storage();

                // Store the encryption key
                {
                    let mut keys = self.encryption_keys.write().await;
                    keys.insert(db_id, encryption_key.clone());
                }

                // Initialize meta file
                let meta = DatabaseMeta {
                    public: true,
                    keys: HashMap::new(),
                };
                self.save_meta(db_id, &meta).await?;

                // Store the server
                let mut servers = self.servers.write().await;
                servers.insert(db_id, Arc::new(server));

                Ok(db_id)
            }
        }
    }

    async fn set_encryption(&self, db_id: u64, _encryption_key: String) -> RpcResult<bool> {
        // Note: In a real implementation, we'd need to modify the existing database
        // For now, return false as encryption can only be set during creation
        let _servers = self.servers.read().await;
        // TODO: Implement encryption setting for existing databases
        Ok(false)
    }

    async fn list_databases(&self) -> RpcResult<Vec<DatabaseInfo>> {
        let db_ids = self.discover_databases().await;
        let mut result = Vec::new();

        for db_id in db_ids {
            // Try to get or create server for this database
            if let Ok(server) = self.get_or_create_server(db_id).await {
                let backend = match server.option.backend {
                    crate::options::BackendType::Redb => BackendType::Redb,
                    crate::options::BackendType::Sled => BackendType::Sled,
                };

                let info = DatabaseInfo {
                    id: db_id,
                    name: None, // TODO: Store name in server metadata
                    backend,
                    encrypted: server.option.encrypt,
                    redis_version: Some("7.0".to_string()), // Default Redis compatibility
                    storage_path: Some(server.option.dir.clone()),
                    size_on_disk: None, // TODO: Calculate actual size
                    key_count: None, // TODO: Get key count from storage
                    created_at: std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap()
                        .as_secs(),
                    last_access: None,
                };
                result.push(info);
            }
        }

        Ok(result)
    }

    async fn get_database_info(&self, db_id: u64) -> RpcResult<DatabaseInfo> {
        let server = self.get_or_create_server(db_id).await?;

        let backend = match server.option.backend {
            crate::options::BackendType::Redb => BackendType::Redb,
            crate::options::BackendType::Sled => BackendType::Sled,
        };

        Ok(DatabaseInfo {
            id: db_id,
            name: None,
            backend,
            encrypted: server.option.encrypt,
            redis_version: Some("7.0".to_string()),
            storage_path: Some(server.option.dir.clone()),
            size_on_disk: None,
            key_count: None,
            created_at: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs(),
            last_access: None,
        })
    }

    async fn delete_database(&self, db_id: u64) -> RpcResult<bool> {
        let mut servers = self.servers.write().await;

        if let Some(_server) = servers.remove(&db_id) {
            // Clean up database files
            let db_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}.db", db_id));
            if db_path.exists() {
                if db_path.is_dir() {
                    std::fs::remove_dir_all(&db_path).ok();
                } else {
                    std::fs::remove_file(&db_path).ok();
                }
            }
            Ok(true)
        } else {
            Ok(false)
        }
    }

    async fn get_server_stats(&self) -> RpcResult<HashMap<String, serde_json::Value>> {
        let db_ids = self.discover_databases().await;
        let mut stats = HashMap::new();

        stats.insert("total_databases".to_string(), serde_json::json!(db_ids.len()));
        stats.insert("uptime".to_string(), serde_json::json!(
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs()
        ));

        Ok(stats)
    }

    async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult<bool> {
        let mut meta = self.load_meta(db_id).await?;

        let perms = match permissions.to_lowercase().as_str() {
            "read" => Permissions::Read,
            "readwrite" => Permissions::ReadWrite,
            _ => return Err(jsonrpsee::types::ErrorObjectOwned::owned(
                -32000,
                "Invalid permissions: use 'read' or 'readwrite'",
                None::<()>
            )),
        };

        let hash = hash_key(&key);
        let access_key = AccessKey {
            hash: hash.clone(),
            permissions: perms,
            created_at: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs(),
        };

        meta.keys.insert(hash, access_key);
        self.save_meta(db_id, &meta).await?;
        Ok(true)
    }

    async fn delete_access_key(&self, db_id: u64, key_hash: String) -> RpcResult<bool> {
        let mut meta = self.load_meta(db_id).await?;

        if meta.keys.remove(&key_hash).is_some() {
            // If no keys left, make database public
            if meta.keys.is_empty() {
                meta.public = true;
            }
            self.save_meta(db_id, &meta).await?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    async fn list_access_keys(&self, db_id: u64) -> RpcResult<Vec<AccessKeyInfo>> {
        let meta = self.load_meta(db_id).await?;
        let keys: Vec<AccessKeyInfo> = meta.keys.values()
            .map(|k| AccessKeyInfo {
                hash: k.hash.clone(),
                permissions: k.permissions.clone(),
                created_at: k.created_at,
            })
            .collect();
        Ok(keys)
    }

    async fn set_database_public(&self, db_id: u64, public: bool) -> RpcResult<bool> {
        let mut meta = self.load_meta(db_id).await?;
        meta.public = public;
        self.save_meta(db_id, &meta).await?;
        Ok(true)
    }
}