markbase/markbase-core/src/webdav_version.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::SystemTime;
use uuid::Uuid;

fn recover_rwlock<T>(result: std::sync::LockResult<T>) -> T {
    match result {
        Ok(guard) => guard,
        Err(e) => {
            log::warn!("RwLock poisoned in webdav_version, recovering");
            e.into_inner()
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VersionInfo {
    pub version_id: String,
    pub file_path: String,
    pub created_at: SystemTime,
    pub size: u64,
    pub checksum: String,
    pub author: Option<String>,
    pub comment: Option<String>,
    pub is_current: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VersionHistory {
    pub file_path: String,
    pub versions: Vec<VersionInfo>,
    pub current_version: String,
    pub total_versions: u64,
}

pub struct WebDavVersioning {
    db: Arc<RwLock<HashMap<String, Vec<u8>>>>,
    version_storage: PathBuf,
    index_path: PathBuf,
}

impl WebDavVersioning {
    pub fn new(version_storage: PathBuf) -> Self {
        let index_path = version_storage.join("version_index.json");
        let db = Arc::new(RwLock::new(HashMap::new()));

        // Load persisted index from disk
        if index_path.exists() {
            if let Ok(json) = std::fs::read_to_string(&index_path) {
                if let Ok(map) = serde_json::from_str::<HashMap<String, Vec<u8>>>(&json) {
                    *recover_rwlock(db.write()) = map;
                }
            }
        }

        Self { db, version_storage, index_path }
    }

    fn save_index(&self) -> Result<(), VersionError> {
        let db = recover_rwlock(self.db.read());
        let json = serde_json::to_string(&*db)?;
        std::fs::write(&self.index_path, json)?;
        Ok(())
    }

    pub fn create_version(
        &self,
        file_path: &str,
        content: &[u8],
        author: Option<&str>,
        comment: Option<&str>,
    ) -> Result<VersionInfo, VersionError> {
        if !self.version_storage.exists() {
            std::fs::create_dir_all(&self.version_storage)?;
        }

        let version_id = Uuid::new_v4().hyphenated().to_string();
        let checksum = Self::calculate_checksum(content);
        let size = content.len() as u64;
        let created_at = SystemTime::now();

        let version_file = self.version_storage.join(&version_id);
        std::fs::write(&version_file, content)?;

        let version_info = VersionInfo {
            version_id: version_id.clone(),
            file_path: file_path.to_string(),
            created_at,
            size,
            checksum,
            author: author.map(|s| s.to_string()),
            comment: comment.map(|s| s.to_string()),
            is_current: true,
        };

        self.mark_previous_versions_not_current(file_path)?;

        let key = Self::version_key(file_path, &version_id);
        let value = serde_json::to_vec(&version_info)?;
        recover_rwlock(self.db.write()).insert(key, value);

        self.update_version_history(file_path, &version_id)?;

        self.save_index()?;

        Ok(version_info)
    }

    pub fn get_version(&self, file_path: &str, version_id: &str) -> Result<Vec<u8>, VersionError> {
        let key = Self::version_key(file_path, version_id);
        let value = recover_rwlock(self.db.read()).get(&key).cloned().ok_or(VersionError::VersionNotFound)?;

        let version_info: VersionInfo = serde_json::from_slice(&value)?;
        let version_file = self.version_storage.join(&version_info.version_id);

        std::fs::read(&version_file).map_err(|e| e.into())
    }

    pub fn get_version_info(&self, file_path: &str, version_id: &str) -> Result<VersionInfo, VersionError> {
        let key = Self::version_key(file_path, version_id);
        let value = recover_rwlock(self.db.read()).get(&key).cloned().ok_or(VersionError::VersionNotFound)?;

        serde_json::from_slice(&value).map_err(|e| e.into())
    }

    pub fn get_version_history(&self, file_path: &str) -> Result<VersionHistory, VersionError> {
        let history_key = Self::history_key(file_path);
        let value = recover_rwlock(self.db.read()).get(&history_key).cloned().ok_or(VersionError::HistoryNotFound)?;

        serde_json::from_slice(&value).map_err(|e| e.into())
    }

    pub fn list_all_versions(&self, file_path: &str) -> Result<Vec<VersionInfo>, VersionError> {
        let prefix = format!("version:{}:", file_path);
        let mut versions = Vec::new();

        let db = recover_rwlock(self.db.read());
        for (key, value) in db.iter() {
            if key.starts_with(&prefix) {
                let version_info: VersionInfo = serde_json::from_slice(value)?;
                versions.push(version_info);
            }
        }

        versions.sort_by(|a, b| b.created_at.cmp(&a.created_at));
        Ok(versions)
    }

    pub fn restore_version(&self, file_path: &str, version_id: &str) -> Result<VersionInfo, VersionError> {
        let old_content = self.get_version(file_path, version_id)?;
        let old_version_info = self.get_version_info(file_path, version_id)?;

        self.mark_previous_versions_not_current(file_path)?;

        let new_version_id = Uuid::new_v4().hyphenated().to_string();
        let new_version_info = VersionInfo {
            version_id: new_version_id.clone(),
            file_path: file_path.to_string(),
            created_at: SystemTime::now(),
            size: old_version_info.size,
            checksum: old_version_info.checksum.clone(),
            author: None,
            comment: Some(format!("Restored from version {}", version_id)),
            is_current: true,
        };

        let version_file = self.version_storage.join(&new_version_id);
        std::fs::write(&version_file, &old_content)?;

        let key = Self::version_key(file_path, &new_version_id);
        let value = serde_json::to_vec(&new_version_info)?;
        recover_rwlock(self.db.write()).insert(key, value);

        self.update_version_history(file_path, &new_version_id)?;

        self.save_index()?;

        Ok(new_version_info)
    }

    pub fn delete_version(&self, file_path: &str, version_id: &str) -> Result<(), VersionError> {
        let version_info = self.get_version_info(file_path, version_id)?;

        if version_info.is_current {
            return Err(VersionError::CannotDeleteCurrentVersion);
        }

        let version_file = self.version_storage.join(version_id);
        if version_file.exists() {
            std::fs::remove_file(&version_file)?;
        }

        let key = Self::version_key(file_path, version_id);
        recover_rwlock(self.db.write()).remove(&key);

        let current = self.get_current_version(file_path)?;
        self.update_version_history(file_path, &current.version_id)?;

        self.save_index()?;

        Ok(())
    }

    pub fn get_current_version(&self, file_path: &str) -> Result<VersionInfo, VersionError> {
        let versions = self.list_all_versions(file_path)?;
        versions
            .into_iter()
            .find(|v| v.is_current)
            .ok_or(VersionError::NoCurrentVersion)
    }

    fn mark_previous_versions_not_current(&self, file_path: &str) -> Result<(), VersionError> {
        let versions = self.list_all_versions(file_path)?;

        for version in versions.iter().filter(|v| v.is_current) {
            let mut updated_version = version.clone();
            updated_version.is_current = false;

            let key = Self::version_key(file_path, &version.version_id);
            let value = serde_json::to_vec(&updated_version)?;
            recover_rwlock(self.db.write()).insert(key, value);
        }

        Ok(())
    }

    fn update_version_history(&self, file_path: &str, current_version_id: &str) -> Result<(), VersionError> {
        let versions = self.list_all_versions(file_path)?;

        let history = VersionHistory {
            file_path: file_path.to_string(),
            versions: versions.clone(),
            current_version: current_version_id.to_string(),
            total_versions: versions.len() as u64,
        };

        let history_key = Self::history_key(file_path);
        let value = serde_json::to_vec(&history)?;
        recover_rwlock(self.db.write()).insert(history_key, value);

        Ok(())
    }

    fn calculate_checksum(content: &[u8]) -> String {
        use sha2::{Sha256, Digest};
        let mut hasher = Sha256::new();
        hasher.update(content);
        format!("{:x}", hasher.finalize())
    }

    fn version_key(file_path: &str, version_id: &str) -> String {
        format!("version:{}:{}", file_path, version_id)
    }

    fn history_key(file_path: &str) -> String {
        format!("history:{}:info", file_path)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum VersionError {
    Io(String),
    Json(String),
    VersionNotFound,
    HistoryNotFound,
    NoCurrentVersion,
    CannotDeleteCurrentVersion,
}

impl From<std::io::Error> for VersionError {
    fn from(e: std::io::Error) -> Self {
        VersionError::Io(e.to_string())
    }
}

impl From<serde_json::Error> for VersionError {
    fn from(e: serde_json::Error) -> Self {
        VersionError::Json(e.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn setup_versioning() -> (WebDavVersioning, TempDir) {
        let version_dir = TempDir::new().unwrap();

        let versioning = WebDavVersioning::new(version_dir.path().to_path_buf());

        (versioning, version_dir)
    }

    #[test]
    fn test_create_version() {
        let (versioning, _) = setup_versioning();
        let content = b"Hello, World!";
        let version_info = versioning.create_version("/test.txt", content, Some("demo"), Some("Initial version")).unwrap();

        assert_eq!(version_info.file_path, "/test.txt");
        assert_eq!(version_info.size, 13);
        assert!(version_info.is_current);
        assert!(version_info.author.is_some());
        assert!(version_info.comment.is_some());
    }

    #[test]
    fn test_get_version() {
        let (versioning, _) = setup_versioning();
        let content = b"Hello, World!";
        let version_info = versioning.create_version("/test.txt", content, None, None).unwrap();

        let retrieved_content = versioning.get_version("/test.txt", &version_info.version_id).unwrap();
        assert_eq!(retrieved_content, content);
    }

    #[test]
    fn test_get_version_history() {
        let (versioning, _) = setup_versioning();
        let content = b"Version 1";
        versioning.create_version("/test.txt", content, None, None).unwrap();

        let content2 = b"Version 2";
        versioning.create_version("/test.txt", content2, None, None).unwrap();

        let history = versioning.get_version_history("/test.txt").unwrap();
        assert_eq!(history.total_versions, 2);
        assert_eq!(history.versions.len(), 2);
    }

    #[test]
    fn test_restore_version() {
        let (versioning, _) = setup_versioning();
        let content1 = b"Original content";
        let version1 = versioning.create_version("/test.txt", content1, None, None).unwrap();

        let content2 = b"Modified content";
        let _version2 = versioning.create_version("/test.txt", content2, None, None).unwrap();

        let restored = versioning.restore_version("/test.txt", &version1.version_id).unwrap();
        assert_eq!(restored.checksum, version1.checksum);
        assert!(restored.is_current);
    }

    #[test]
    fn test_delete_version() {
        let (versioning, _) = setup_versioning();
        let content1 = b"Version 1";
        let version1 = versioning.create_version("/test.txt", content1, None, None).unwrap();

        let content2 = b"Version 2";
        versioning.create_version("/test.txt", content2, None, None).unwrap();

        versioning.delete_version("/test.txt", &version1.version_id).unwrap();

        let history = versioning.get_version_history("/test.txt").unwrap();
        assert_eq!(history.total_versions, 1);
    }

    #[test]
    fn test_cannot_delete_current_version() {
        let (versioning, _) = setup_versioning();
        let content = b"Current version";
        let version_info = versioning.create_version("/test.txt", content, None, None).unwrap();

        let result = versioning.delete_version("/test.txt", &version_info.version_id);
        assert!(matches!(result, Err(VersionError::CannotDeleteCurrentVersion)));
    }

    #[test]
    fn test_get_current_version() {
        let (versioning, _) = setup_versioning();
        let content1 = b"Old version";
        versioning.create_version("/test.txt", content1, None, None).unwrap();

        let content2 = b"Current version";
        let version2 = versioning.create_version("/test.txt", content2, None, None).unwrap();

        let current = versioning.get_current_version("/test.txt").unwrap();
        assert_eq!(current.version_id, version2.version_id);
        assert!(current.is_current);
    }

    #[test]
    fn test_checksum_calculation() {
        let content = b"Hello, World!";
        let checksum = WebDavVersioning::calculate_checksum(content);
        assert_eq!(checksum.len(), 64);
        assert!(checksum.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn test_list_all_versions_sorted() {
        let (versioning, _) = setup_versioning();
        let content1 = b"Version 1";
        versioning.create_version("/test.txt", content1, None, None).unwrap();

        std::thread::sleep(std::time::Duration::from_millis(10));

        let content2 = b"Version 2";
        versioning.create_version("/test.txt", content2, None, None).unwrap();

        let versions = versioning.list_all_versions("/test.txt").unwrap();
        assert_eq!(versions.len(), 2);
        assert!(versions[0].created_at >= versions[1].created_at);
    }

    #[test]
    fn test_version_not_found() {
        let (versioning, _) = setup_versioning();
        let result = versioning.get_version("/nonexistent.txt", "nonexistent-id");
        assert!(matches!(result, Err(VersionError::VersionNotFound)));
    }

    #[test]
    fn test_history_not_found() {
        let (versioning, _) = setup_versioning();
        let result = versioning.get_version_history("/nonexistent.txt");
        assert!(matches!(result, Err(VersionError::HistoryNotFound)));
    }
}