//! Storage abstraction for VM persistence

use crate::config::{StorageBackendKind, StorageRuntimeConfig};
use async_trait::async_trait;
use plasmavmc_types::{Image, ImageFormat, Node, VirtualMachine, VmHandle, Volume};
use std::path::PathBuf;
use thiserror::Error;

const FLAREDB_CAS_RETRIES: usize = 16;
const FLAREDB_CAS_RETRY_BASE_DELAY_MS: u64 = 10;

/// Storage backend type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StorageBackend {
    FlareDB,
    File,
}

impl StorageBackend {
    pub fn from_config(config: &StorageRuntimeConfig) -> Self {
        match config.backend {
            StorageBackendKind::Flaredb => Self::FlareDB,
            StorageBackendKind::File => Self::File,
        }
    }
}

/// Storage error
#[derive(Debug, Error)]
pub enum StorageError {
    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),
    #[error("FlareDB error: {0}")]
    FlareDB(String),
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Lock contention: {0}")]
    LockContention(String),
    #[error("Not found: {0}")]
    NotFound(String),
    #[error("Storage unavailable")]
    Unavailable,
}

/// Result type for storage operations
pub type StorageResult<T> = Result<T, StorageError>;

#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ImageUploadPart {
    pub part_number: u32,
    pub etag: String,
    pub size_bytes: u64,
}

#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ImageUploadSession {
    pub session_id: String,
    pub org_id: String,
    pub project_id: String,
    pub image_id: String,
    pub upload_id: String,
    pub staging_key: String,
    pub source_format: ImageFormat,
    #[serde(default)]
    pub parts: Vec<ImageUploadPart>,
    pub committed_size_bytes: u64,
    pub created_at: u64,
    pub updated_at: u64,
}

/// Storage trait for VM persistence
#[async_trait]
pub trait VmStore: Send + Sync {
    /// Save a VM
    async fn save_vm(&self, vm: &VirtualMachine) -> StorageResult<()>;

    /// Load a VM by ID
    async fn load_vm(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
    ) -> StorageResult<Option<VirtualMachine>>;

    /// Delete a VM
    async fn delete_vm(&self, org_id: &str, project_id: &str, vm_id: &str) -> StorageResult<()>;

    /// List all VMs for a tenant
    async fn list_vms(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<VirtualMachine>>;

    /// List all VMs across every tenant.
    async fn list_all_vms(&self) -> StorageResult<Vec<VirtualMachine>>;

    /// Save a VM handle
    async fn save_handle(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
        handle: &VmHandle,
    ) -> StorageResult<()>;

    /// Load a VM handle
    async fn load_handle(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
    ) -> StorageResult<Option<VmHandle>>;

    /// Delete a VM handle
    async fn delete_handle(&self, org_id: &str, project_id: &str, vm_id: &str)
        -> StorageResult<()>;

    /// Save a node
    async fn save_node(&self, node: &Node) -> StorageResult<()>;

    /// Load a node by ID
    async fn load_node(&self, node_id: &str) -> StorageResult<Option<Node>>;

    /// Delete a node
    async fn delete_node(&self, node_id: &str) -> StorageResult<()>;

    /// List all nodes
    async fn list_nodes(&self) -> StorageResult<Vec<Node>>;

    /// Save an image
    async fn save_image(&self, image: &Image) -> StorageResult<()>;

    /// Load an image by ID
    async fn load_image(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<Option<Image>>;

    /// Delete an image
    async fn delete_image(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<()>;

    /// List images for a tenant
    async fn list_images(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<Image>>;

    /// Save an in-progress image upload session.
    async fn save_image_upload_session(&self, session: &ImageUploadSession) -> StorageResult<()>;

    /// Load an in-progress image upload session.
    async fn load_image_upload_session(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
        session_id: &str,
    ) -> StorageResult<Option<ImageUploadSession>>;

    /// Delete an image upload session.
    async fn delete_image_upload_session(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
        session_id: &str,
    ) -> StorageResult<()>;

    /// List image upload sessions for an image.
    async fn list_image_upload_sessions(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<Vec<ImageUploadSession>>;

    /// Save a persistent volume
    async fn save_volume(&self, volume: &Volume) -> StorageResult<()>;

    /// Conditionally save a persistent volume if the currently stored value still matches `expected`.
    async fn compare_and_swap_volume(
        &self,
        expected: Option<&Volume>,
        volume: &Volume,
    ) -> StorageResult<bool>;

    /// Load a volume by ID
    async fn load_volume(
        &self,
        org_id: &str,
        project_id: &str,
        volume_id: &str,
    ) -> StorageResult<Option<Volume>>;

    /// Delete a volume
    async fn delete_volume(
        &self,
        org_id: &str,
        project_id: &str,
        volume_id: &str,
    ) -> StorageResult<()>;

    /// List volumes for a tenant
    async fn list_volumes(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<Volume>>;
}

/// Build key for VM metadata
fn vm_key(org_id: &str, project_id: &str, vm_id: &str) -> String {
    format!("/plasmavmc/vms/{}/{}/{}", org_id, project_id, vm_id)
}

/// Build key for VM handle
fn handle_key(org_id: &str, project_id: &str, vm_id: &str) -> String {
    format!("/plasmavmc/handles/{}/{}/{}", org_id, project_id, vm_id)
}

/// Build prefix for tenant VM listing
fn vm_prefix(org_id: &str, project_id: &str) -> String {
    format!("/plasmavmc/vms/{}/{}/", org_id, project_id)
}

/// Build key for node metadata
fn node_key(node_id: &str) -> String {
    format!("/plasmavmc/nodes/{}", node_id)
}

/// Build prefix for node listing
fn node_prefix() -> String {
    "/plasmavmc/nodes/".to_string()
}

/// Build key for image metadata
fn image_key(org_id: &str, project_id: &str, image_id: &str) -> String {
    format!("/plasmavmc/images/{}/{}/{}", org_id, project_id, image_id)
}

/// Build prefix for tenant image listing
fn image_prefix(org_id: &str, project_id: &str) -> String {
    format!("/plasmavmc/images/{}/{}/", org_id, project_id)
}

/// Build key for image upload session metadata
fn image_upload_session_key(
    org_id: &str,
    project_id: &str,
    image_id: &str,
    session_id: &str,
) -> String {
    format!(
        "/plasmavmc/image-upload-sessions/{}/{}/{}/{}",
        org_id, project_id, image_id, session_id
    )
}

/// Build prefix for image upload session listing
fn image_upload_session_prefix(org_id: &str, project_id: &str, image_id: &str) -> String {
    format!(
        "/plasmavmc/image-upload-sessions/{}/{}/{}/",
        org_id, project_id, image_id
    )
}

/// Build key for volume metadata
fn volume_key(org_id: &str, project_id: &str, volume_id: &str) -> String {
    format!("/plasmavmc/volumes/{}/{}/{}", org_id, project_id, volume_id)
}

/// Build prefix for tenant volume listing
fn volume_prefix(org_id: &str, project_id: &str) -> String {
    format!("/plasmavmc/volumes/{}/{}/", org_id, project_id)
}

/// FlareDB-backed storage
pub struct FlareDBStore {
    client: tokio::sync::Mutex<flaredb_client::RdbClient>,
}

impl FlareDBStore {
    pub async fn new_with_config(config: &StorageRuntimeConfig) -> StorageResult<Self> {
        Self::new_with_endpoints(
            config.flaredb_endpoint.clone(),
            config.chainfire_endpoint.clone(),
        )
        .await
    }

    pub async fn new_with_endpoints(
        endpoint: Option<String>,
        pd_endpoint: Option<String>,
    ) -> StorageResult<Self> {
        let endpoint = endpoint.unwrap_or_else(|| "127.0.0.1:2479".to_string());
        let pd_endpoint = pd_endpoint
            .map(|value| normalize_transport_addr(&value))
            .unwrap_or_else(|| endpoint.clone());

        let client = flaredb_client::RdbClient::connect_with_pd_namespace(
            endpoint,
            pd_endpoint,
            "plasmavmc",
        )
        .await
        .map_err(|e| StorageError::FlareDB(format!("Failed to connect to FlareDB: {}", e)))?;

        Ok(Self {
            client: tokio::sync::Mutex::new(client),
        })
    }

    fn prefix_end(prefix: &[u8]) -> Vec<u8> {
        let mut end_key = prefix.to_vec();
        if let Some(last) = end_key.last_mut() {
            if *last == 0xff {
                end_key.push(0x00);
            } else {
                *last += 1;
            }
        } else {
            end_key.push(0xff);
        }
        end_key
    }

    async fn cas_put(&self, key: &str, value: Vec<u8>) -> StorageResult<()> {
        let key = key.as_bytes().to_vec();
        let mut attempts = 0;

        loop {
            let expected_version = {
                let mut client = self.client.lock().await;
                client
                    .cas_get(key.clone())
                    .await
                    .map_err(|e| {
                        StorageError::FlareDB(format!(
                            "FlareDB put failed during version lookup: {}",
                            e
                        ))
                    })?
                    .map(|(version, _)| version)
                    .unwrap_or(0)
            };

            let (success, current_version, _) = {
                let mut client = self.client.lock().await;
                client
                    .cas(key.clone(), value.clone(), expected_version)
                    .await
                    .map_err(|e| StorageError::FlareDB(format!("FlareDB put failed: {}", e)))?
            };

            if success {
                return Ok(());
            }

            attempts += 1;
            if attempts >= FLAREDB_CAS_RETRIES {
                return Err(StorageError::FlareDB(format!(
                    "FlareDB put failed after {} CAS retries (expected version {}, actual version {})",
                    FLAREDB_CAS_RETRIES, expected_version, current_version
                )));
            }

            tokio::time::sleep(std::time::Duration::from_millis(
                FLAREDB_CAS_RETRY_BASE_DELAY_MS.saturating_mul(attempts as u64),
            ))
            .await;
        }
    }

    async fn cas_get(&self, key: &str) -> StorageResult<Option<Vec<u8>>> {
        let mut client = self.client.lock().await;
        Ok(client
            .cas_get(key.as_bytes().to_vec())
            .await
            .map_err(|e| StorageError::FlareDB(format!("FlareDB get failed: {}", e)))?
            .map(|(_, value)| value))
    }

    async fn cas_get_versioned(&self, key: &str) -> StorageResult<Option<(u64, Vec<u8>)>> {
        let mut client = self.client.lock().await;
        client
            .cas_get(key.as_bytes().to_vec())
            .await
            .map_err(|e| StorageError::FlareDB(format!("FlareDB get failed: {}", e)))
    }

    async fn cas_delete(&self, key: &str) -> StorageResult<()> {
        let key = key.as_bytes().to_vec();
        let mut attempts = 0;

        loop {
            let current = {
                let mut client = self.client.lock().await;
                client.cas_get(key.clone()).await.map_err(|e| {
                    StorageError::FlareDB(format!(
                        "FlareDB delete failed during version lookup: {}",
                        e
                    ))
                })?
            };

            let Some((expected_version, _)) = current else {
                return Ok(());
            };

            let (success, current_version, _) = {
                let mut client = self.client.lock().await;
                client
                    .cas_delete(key.clone(), expected_version)
                    .await
                    .map_err(|e| StorageError::FlareDB(format!("FlareDB delete failed: {}", e)))?
            };

            if success {
                return Ok(());
            }

            attempts += 1;
            if attempts >= FLAREDB_CAS_RETRIES {
                return Err(StorageError::FlareDB(format!(
                    "FlareDB delete failed after {} CAS retries (expected version {}, actual version {})",
                    FLAREDB_CAS_RETRIES, expected_version, current_version
                )));
            }

            tokio::time::sleep(std::time::Duration::from_millis(
                FLAREDB_CAS_RETRY_BASE_DELAY_MS.saturating_mul(attempts as u64),
            ))
            .await;
        }
    }

    async fn cas_scan_values(&self, prefix: &str) -> StorageResult<Vec<Vec<u8>>> {
        let end_key = Self::prefix_end(prefix.as_bytes());
        let mut values = Vec::new();
        let mut start_key = prefix.as_bytes().to_vec();

        loop {
            let (entries, next) = {
                let mut client = self.client.lock().await;
                client
                    .cas_scan(start_key.clone(), end_key.clone(), 1000)
                    .await
                    .map_err(|e| StorageError::FlareDB(format!("FlareDB scan failed: {}", e)))?
            };

            for (_, value, _) in entries {
                values.push(value);
            }

            if let Some(next_key) = next {
                start_key = next_key;
            } else {
                break;
            }
        }

        Ok(values)
    }
}

fn normalize_transport_addr(endpoint: &str) -> String {
    endpoint
        .trim()
        .trim_start_matches("http://")
        .trim_start_matches("https://")
        .trim_end_matches('/')
        .to_string()
}

#[async_trait]
impl VmStore for FlareDBStore {
    async fn save_vm(&self, vm: &VirtualMachine) -> StorageResult<()> {
        let key = vm_key(&vm.org_id, &vm.project_id, &vm.id.to_string());
        let value = serde_json::to_vec(vm)?;
        self.cas_put(&key, value).await
    }

    async fn load_vm(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
    ) -> StorageResult<Option<VirtualMachine>> {
        let key = vm_key(org_id, project_id, vm_id);
        match self.cas_get(&key).await? {
            Some(data) => {
                let vm: VirtualMachine = serde_json::from_slice(&data)?;
                Ok(Some(vm))
            }
            None => Ok(None),
        }
    }

    async fn delete_vm(&self, org_id: &str, project_id: &str, vm_id: &str) -> StorageResult<()> {
        let key = vm_key(org_id, project_id, vm_id);
        self.cas_delete(&key).await
    }

    async fn list_vms(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<VirtualMachine>> {
        let prefix = vm_prefix(org_id, project_id);
        let mut vms = Vec::new();
        for value in self.cas_scan_values(&prefix).await? {
            if let Ok(vm) = serde_json::from_slice::<VirtualMachine>(&value) {
                vms.push(vm);
            }
        }

        Ok(vms)
    }

    async fn list_all_vms(&self) -> StorageResult<Vec<VirtualMachine>> {
        let mut vms = Vec::new();
        for value in self.cas_scan_values("/plasmavmc/vms/").await? {
            if let Ok(vm) = serde_json::from_slice::<VirtualMachine>(&value) {
                vms.push(vm);
            }
        }

        Ok(vms)
    }

    async fn save_handle(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
        handle: &VmHandle,
    ) -> StorageResult<()> {
        let key = handle_key(org_id, project_id, vm_id);
        let value = serde_json::to_vec(handle)?;
        self.cas_put(&key, value).await
    }

    async fn load_handle(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
    ) -> StorageResult<Option<VmHandle>> {
        let key = handle_key(org_id, project_id, vm_id);
        match self.cas_get(&key).await? {
            Some(data) => {
                let handle: VmHandle = serde_json::from_slice(&data)?;
                Ok(Some(handle))
            }
            None => Ok(None),
        }
    }

    async fn delete_handle(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
    ) -> StorageResult<()> {
        let key = handle_key(org_id, project_id, vm_id);
        self.cas_delete(&key).await
    }

    async fn save_node(&self, node: &Node) -> StorageResult<()> {
        let key = node_key(node.id.as_str());
        let value = serde_json::to_vec(node)?;
        self.cas_put(&key, value).await
    }

    async fn load_node(&self, node_id: &str) -> StorageResult<Option<Node>> {
        let key = node_key(node_id);
        match self.cas_get(&key).await? {
            Some(data) => {
                let node: Node = serde_json::from_slice(&data)?;
                Ok(Some(node))
            }
            None => Ok(None),
        }
    }

    async fn delete_node(&self, node_id: &str) -> StorageResult<()> {
        let key = node_key(node_id);
        self.cas_delete(&key).await
    }

    async fn list_nodes(&self) -> StorageResult<Vec<Node>> {
        let prefix = node_prefix();
        let mut nodes = Vec::new();
        for value in self.cas_scan_values(&prefix).await? {
            if let Ok(node) = serde_json::from_slice::<Node>(&value) {
                nodes.push(node);
            }
        }

        Ok(nodes)
    }

    async fn save_image(&self, image: &Image) -> StorageResult<()> {
        let key = image_key(&image.org_id, &image.project_id, &image.id);
        let value = serde_json::to_vec(image)?;
        self.cas_put(&key, value).await
    }

    async fn load_image(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<Option<Image>> {
        let key = image_key(org_id, project_id, image_id);
        match self.cas_get(&key).await? {
            Some(data) => {
                let image: Image = serde_json::from_slice(&data)?;
                Ok(Some(image))
            }
            None => Ok(None),
        }
    }

    async fn delete_image(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<()> {
        let key = image_key(org_id, project_id, image_id);
        self.cas_delete(&key).await
    }

    async fn list_images(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<Image>> {
        let prefix = image_prefix(org_id, project_id);
        let mut images = Vec::new();
        for value in self.cas_scan_values(&prefix).await? {
            if let Ok(image) = serde_json::from_slice::<Image>(&value) {
                images.push(image);
            }
        }

        Ok(images)
    }

    async fn save_image_upload_session(&self, session: &ImageUploadSession) -> StorageResult<()> {
        let key = image_upload_session_key(
            &session.org_id,
            &session.project_id,
            &session.image_id,
            &session.session_id,
        );
        let value = serde_json::to_vec(session)?;
        self.cas_put(&key, value).await
    }

    async fn load_image_upload_session(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
        session_id: &str,
    ) -> StorageResult<Option<ImageUploadSession>> {
        let key = image_upload_session_key(org_id, project_id, image_id, session_id);
        match self.cas_get(&key).await? {
            Some(data) => Ok(Some(serde_json::from_slice(&data)?)),
            None => Ok(None),
        }
    }

    async fn delete_image_upload_session(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
        session_id: &str,
    ) -> StorageResult<()> {
        let key = image_upload_session_key(org_id, project_id, image_id, session_id);
        self.cas_delete(&key).await
    }

    async fn list_image_upload_sessions(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<Vec<ImageUploadSession>> {
        let prefix = image_upload_session_prefix(org_id, project_id, image_id);
        let mut sessions = Vec::new();
        for value in self.cas_scan_values(&prefix).await? {
            if let Ok(session) = serde_json::from_slice::<ImageUploadSession>(&value) {
                sessions.push(session);
            }
        }
        Ok(sessions)
    }

    async fn save_volume(&self, volume: &Volume) -> StorageResult<()> {
        let key = volume_key(&volume.org_id, &volume.project_id, &volume.id);
        let value = serde_json::to_vec(volume)?;
        self.cas_put(&key, value).await
    }

    async fn compare_and_swap_volume(
        &self,
        expected: Option<&Volume>,
        volume: &Volume,
    ) -> StorageResult<bool> {
        let key = volume_key(&volume.org_id, &volume.project_id, &volume.id);
        let current = self.cas_get_versioned(&key).await?;

        match (expected, &current) {
            (None, None) => {}
            (None, Some(_)) => return Ok(false),
            (Some(_), None) => return Ok(false),
            (Some(expected), Some((_, data))) => {
                let current_volume: Volume = serde_json::from_slice(data)?;
                if &current_volume != expected {
                    return Ok(false);
                }
            }
        }

        let expected_version = current.as_ref().map(|(version, _)| *version).unwrap_or(0);
        let value = serde_json::to_vec(volume)?;
        let (success, _, _) = {
            let mut client = self.client.lock().await;
            client
                .cas(key.as_bytes().to_vec(), value, expected_version)
                .await
                .map_err(|e| {
                    StorageError::FlareDB(format!("FlareDB CAS volume save failed: {}", e))
                })?
        };
        Ok(success)
    }

    async fn load_volume(
        &self,
        org_id: &str,
        project_id: &str,
        volume_id: &str,
    ) -> StorageResult<Option<Volume>> {
        let key = volume_key(org_id, project_id, volume_id);
        match self.cas_get(&key).await? {
            Some(data) => Ok(Some(serde_json::from_slice(&data)?)),
            None => Ok(None),
        }
    }

    async fn delete_volume(
        &self,
        org_id: &str,
        project_id: &str,
        volume_id: &str,
    ) -> StorageResult<()> {
        let key = volume_key(org_id, project_id, volume_id);
        self.cas_delete(&key).await
    }

    async fn list_volumes(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<Volume>> {
        let prefix = volume_prefix(org_id, project_id);
        let mut volumes = Vec::new();
        for value in self.cas_scan_values(&prefix).await? {
            if let Ok(volume) = serde_json::from_slice::<Volume>(&value) {
                volumes.push(volume);
            }
        }

        Ok(volumes)
    }
}

/// File-backed storage with atomic writes
pub struct FileStore {
    state_path: PathBuf,
}

#[derive(Default, serde::Serialize, serde::Deserialize)]
struct PersistedState {
    #[serde(default)]
    vms: Vec<VirtualMachine>,
    #[serde(default)]
    handles: Vec<VmHandle>,
    #[serde(default)]
    nodes: Vec<Node>,
    #[serde(default)]
    images: Vec<Image>,
    #[serde(default)]
    image_upload_sessions: Vec<ImageUploadSession>,
    #[serde(default)]
    volumes: Vec<Volume>,
}

impl FileStore {
    /// Create a new file store
    pub fn new(path: Option<PathBuf>) -> Self {
        let state_path = path.unwrap_or_else(|| {
            std::env::var("PLASMAVMC_STATE_PATH")
                .map(PathBuf::from)
                .unwrap_or_else(|_| PathBuf::from("/var/run/plasmavmc/state.json"))
        });
        Self { state_path }
    }

    /// Load state from file
    fn load_state(&self) -> StorageResult<PersistedState> {
        let data = std::fs::read(&self.state_path)?;
        let state: PersistedState = serde_json::from_slice(&data)?;
        Ok(state)
    }

    /// Save state to file atomically
    fn save_state(&self, state: &PersistedState) -> StorageResult<()> {
        let serialized = serde_json::to_vec_pretty(state)?;
        if let Some(parent) = self.state_path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        // Atomic write: write to temp file, then rename
        let temp_path = self.state_path.with_extension("json.tmp");
        std::fs::write(&temp_path, serialized)?;
        std::fs::rename(&temp_path, &self.state_path)?;
        Ok(())
    }
}

#[async_trait]
impl VmStore for FileStore {
    async fn save_vm(&self, vm: &VirtualMachine) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        // Remove existing VM if present
        state.vms.retain(|v| v.id.to_string() != vm.id.to_string());
        state.vms.push(vm.clone());
        self.save_state(&state)?;
        Ok(())
    }

    async fn load_vm(
        &self,
        org_id: &str,
        project_id: &str,
        vm_id: &str,
    ) -> StorageResult<Option<VirtualMachine>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state.vms.into_iter().find(|v| {
            v.org_id == org_id && v.project_id == project_id && v.id.to_string() == vm_id
        }))
    }

    async fn delete_vm(&self, org_id: &str, project_id: &str, vm_id: &str) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.vms.retain(|v| {
            !(v.org_id == org_id && v.project_id == project_id && v.id.to_string() == vm_id)
        });
        state.handles.retain(|h| h.vm_id.to_string() != vm_id);
        self.save_state(&state)?;
        Ok(())
    }

    async fn list_vms(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<VirtualMachine>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state
            .vms
            .into_iter()
            .filter(|v| v.org_id == org_id && v.project_id == project_id)
            .collect())
    }

    async fn list_all_vms(&self) -> StorageResult<Vec<VirtualMachine>> {
        Ok(self.load_state().unwrap_or_default().vms)
    }

    async fn save_handle(
        &self,
        _org_id: &str,
        _project_id: &str,
        vm_id: &str,
        handle: &VmHandle,
    ) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.handles.retain(|h| h.vm_id.to_string() != vm_id);
        state.handles.push(handle.clone());
        self.save_state(&state)?;
        Ok(())
    }

    async fn load_handle(
        &self,
        _org_id: &str,
        _project_id: &str,
        vm_id: &str,
    ) -> StorageResult<Option<VmHandle>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state
            .handles
            .into_iter()
            .find(|h| h.vm_id.to_string() == vm_id))
    }

    async fn delete_handle(
        &self,
        _org_id: &str,
        _project_id: &str,
        vm_id: &str,
    ) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.handles.retain(|h| h.vm_id.to_string() != vm_id);
        self.save_state(&state)?;
        Ok(())
    }

    async fn save_node(&self, node: &Node) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.nodes.retain(|n| n.id.as_str() != node.id.as_str());
        state.nodes.push(node.clone());
        self.save_state(&state)?;
        Ok(())
    }

    async fn load_node(&self, node_id: &str) -> StorageResult<Option<Node>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state.nodes.into_iter().find(|n| n.id.as_str() == node_id))
    }

    async fn delete_node(&self, node_id: &str) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.nodes.retain(|n| n.id.as_str() != node_id);
        self.save_state(&state)?;
        Ok(())
    }

    async fn list_nodes(&self) -> StorageResult<Vec<Node>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state.nodes)
    }

    async fn save_image(&self, image: &Image) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.images.retain(|existing| existing.id != image.id);
        state.images.push(image.clone());
        self.save_state(&state)?;
        Ok(())
    }

    async fn load_image(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<Option<Image>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state.images.into_iter().find(|image| {
            image.org_id == org_id && image.project_id == project_id && image.id == image_id
        }))
    }

    async fn delete_image(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.images.retain(|image| {
            !(image.org_id == org_id && image.project_id == project_id && image.id == image_id)
        });
        self.save_state(&state)?;
        Ok(())
    }

    async fn list_images(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<Image>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state
            .images
            .into_iter()
            .filter(|image| image.org_id == org_id && image.project_id == project_id)
            .collect())
    }

    async fn save_image_upload_session(&self, session: &ImageUploadSession) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.image_upload_sessions.retain(|existing| {
            !(existing.org_id == session.org_id
                && existing.project_id == session.project_id
                && existing.image_id == session.image_id
                && existing.session_id == session.session_id)
        });
        state.image_upload_sessions.push(session.clone());
        self.save_state(&state)?;
        Ok(())
    }

    async fn load_image_upload_session(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
        session_id: &str,
    ) -> StorageResult<Option<ImageUploadSession>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state.image_upload_sessions.into_iter().find(|session| {
            session.org_id == org_id
                && session.project_id == project_id
                && session.image_id == image_id
                && session.session_id == session_id
        }))
    }

    async fn delete_image_upload_session(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
        session_id: &str,
    ) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.image_upload_sessions.retain(|session| {
            !(session.org_id == org_id
                && session.project_id == project_id
                && session.image_id == image_id
                && session.session_id == session_id)
        });
        self.save_state(&state)?;
        Ok(())
    }

    async fn list_image_upload_sessions(
        &self,
        org_id: &str,
        project_id: &str,
        image_id: &str,
    ) -> StorageResult<Vec<ImageUploadSession>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state
            .image_upload_sessions
            .into_iter()
            .filter(|session| {
                session.org_id == org_id
                    && session.project_id == project_id
                    && session.image_id == image_id
            })
            .collect())
    }

    async fn save_volume(&self, volume: &Volume) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.volumes.retain(|existing| existing.id != volume.id);
        state.volumes.push(volume.clone());
        self.save_state(&state)?;
        Ok(())
    }

    async fn compare_and_swap_volume(
        &self,
        expected: Option<&Volume>,
        volume: &Volume,
    ) -> StorageResult<bool> {
        let mut state = self.load_state().unwrap_or_default();
        let current = state.volumes.iter().find(|candidate| {
            candidate.org_id == volume.org_id
                && candidate.project_id == volume.project_id
                && candidate.id == volume.id
        });

        match (expected, current) {
            (None, None) => {}
            (None, Some(_)) => return Ok(false),
            (Some(_), None) => return Ok(false),
            (Some(expected), Some(current)) if current == expected => {}
            (Some(_), Some(_)) => return Ok(false),
        }

        state.volumes.retain(|existing| {
            !(existing.org_id == volume.org_id
                && existing.project_id == volume.project_id
                && existing.id == volume.id)
        });
        state.volumes.push(volume.clone());
        self.save_state(&state)?;
        Ok(true)
    }

    async fn load_volume(
        &self,
        org_id: &str,
        project_id: &str,
        volume_id: &str,
    ) -> StorageResult<Option<Volume>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state.volumes.into_iter().find(|volume| {
            volume.org_id == org_id && volume.project_id == project_id && volume.id == volume_id
        }))
    }

    async fn delete_volume(
        &self,
        org_id: &str,
        project_id: &str,
        volume_id: &str,
    ) -> StorageResult<()> {
        let mut state = self.load_state().unwrap_or_default();
        state.volumes.retain(|volume| {
            !(volume.org_id == org_id && volume.project_id == project_id && volume.id == volume_id)
        });
        self.save_state(&state)?;
        Ok(())
    }

    async fn list_volumes(&self, org_id: &str, project_id: &str) -> StorageResult<Vec<Volume>> {
        let state = self.load_state().unwrap_or_default();
        Ok(state
            .volumes
            .into_iter()
            .filter(|volume| volume.org_id == org_id && volume.project_id == project_id)
            .collect())
    }
}

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

    fn test_volume() -> Volume {
        Volume::new("vol-1", "vol-1", "org-1", "project-1", 8)
    }

    #[tokio::test]
    async fn filestore_compare_and_swap_volume_creates_when_absent() {
        let tempdir = tempdir().unwrap();
        let store = FileStore::new(Some(tempdir.path().join("state.json")));
        let volume = test_volume();

        let created = store.compare_and_swap_volume(None, &volume).await.unwrap();
        assert!(created);

        let loaded = store
            .load_volume(&volume.org_id, &volume.project_id, &volume.id)
            .await
            .unwrap()
            .expect("volume should exist after CAS create");
        assert_eq!(loaded, volume);
    }

    #[tokio::test]
    async fn filestore_compare_and_swap_volume_rejects_stale_expected_value() {
        let tempdir = tempdir().unwrap();
        let store = FileStore::new(Some(tempdir.path().join("state.json")));
        let original = test_volume();
        store.save_volume(&original).await.unwrap();

        let mut current = original.clone();
        current.attached_to_vm = Some("vm-a".to_string());
        current.attached_to_node = Some("node04".to_string());
        current.attachment_generation = 1;
        store.save_volume(&current).await.unwrap();

        let mut stale_update = original.clone();
        stale_update.attached_to_vm = Some("vm-b".to_string());
        stale_update.attached_to_node = Some("node05".to_string());
        stale_update.attachment_generation = 1;

        let swapped = store
            .compare_and_swap_volume(Some(&original), &stale_update)
            .await
            .unwrap();
        assert!(!swapped);

        let loaded = store
            .load_volume(&current.org_id, &current.project_id, &current.id)
            .await
            .unwrap()
            .expect("current volume should remain");
        assert_eq!(loaded, current);
    }

    #[tokio::test]
    async fn filestore_round_trips_image_upload_sessions() {
        let tempdir = tempdir().unwrap();
        let store = FileStore::new(Some(tempdir.path().join("state.json")));
        let session = ImageUploadSession {
            session_id: "11111111-1111-1111-1111-111111111111".to_string(),
            org_id: "org-1".to_string(),
            project_id: "project-1".to_string(),
            image_id: "22222222-2222-2222-2222-222222222222".to_string(),
            upload_id: "multipart-1".to_string(),
            staging_key: "org-1/project-1/uploads/22222222-2222-2222-2222-222222222222.source"
                .to_string(),
            source_format: ImageFormat::Raw,
            parts: vec![ImageUploadPart {
                part_number: 1,
                etag: "etag-1".to_string(),
                size_bytes: 4096,
            }],
            committed_size_bytes: 4096,
            created_at: 1,
            updated_at: 2,
        };

        store.save_image_upload_session(&session).await.unwrap();

        let loaded = store
            .load_image_upload_session(
                &session.org_id,
                &session.project_id,
                &session.image_id,
                &session.session_id,
            )
            .await
            .unwrap()
            .expect("session should exist");
        assert_eq!(loaded, session);

        let listed = store
            .list_image_upload_sessions(&session.org_id, &session.project_id, &session.image_id)
            .await
            .unwrap();
        assert_eq!(listed, vec![session.clone()]);

        store
            .delete_image_upload_session(
                &session.org_id,
                &session.project_id,
                &session.image_id,
                &session.session_id,
            )
            .await
            .unwrap();
        assert!(store
            .load_image_upload_session(
                &session.org_id,
                &session.project_id,
                &session.image_id,
                &session.session_id,
            )
            .await
            .unwrap()
            .is_none());
    }
}
