Arkendro/server/src/controllers/snapshots.rs

use crate::sync::storage::Storage;
use crate::sync::meta::{MetaObj, FsType};
use crate::sync::protocol::MetaType;
use crate::utils::{error::*, models::*, DbPool};
use serde::Serialize;
use chrono::{DateTime, Utc};

#[derive(Debug, Serialize)]
pub struct SnapshotInfo {
    pub id: String, // Use UUID string instead of integer
    pub snapshot_hash: String,
    pub created_at: String,
    pub disks: Vec<DiskInfo>,
}

#[derive(Debug, Serialize)]
pub struct DiskInfo {
    pub serial: String,
    pub size_bytes: u64,
    pub partitions: Vec<PartitionInfo>,
}

#[derive(Debug, Serialize)]
pub struct PartitionInfo {
    pub fs_type: String,
    pub start_lba: u64,
    pub end_lba: u64,
    pub size_bytes: u64,
}

pub struct SnapshotsController;

impl SnapshotsController {
    pub async fn get_machine_snapshots(
        pool: &DbPool,
        machine_id: i64,
        user: &User,
    ) -> AppResult<Vec<SnapshotInfo>> {
        // Verify machine access
        let machine = sqlx::query!(
            "SELECT id, user_id FROM machines WHERE id = ? AND user_id = ?",
            machine_id,
            user.id
        )
        .fetch_optional(pool)
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        if machine.is_none() {
            return Err(AppError::NotFoundError("Machine not found or access denied".to_string()));
        }

        let _machine = machine.unwrap();

        let storage = Storage::new("./data");
        let mut snapshot_infos = Vec::new();

        // List all snapshots for this machine from storage
        match storage.list_snapshots(machine_id).await {
            Ok(snapshot_ids) => {
                for snapshot_id in snapshot_ids {
                    // Load snapshot reference to get hash and timestamp
                    if let Ok(Some((snapshot_hash, created_at_timestamp))) = storage.load_snapshot_ref(machine_id, &snapshot_id).await {
                        // Load snapshot metadata
                        if let Ok(Some(snapshot_meta)) = storage.load_meta(MetaType::Snapshot, &snapshot_hash).await {
                            if let MetaObj::Snapshot(snapshot_obj) = snapshot_meta {
                                let mut disks = Vec::new();

                                for disk_hash in snapshot_obj.disk_hashes {
                                    if let Ok(Some(disk_meta)) = storage.load_meta(MetaType::Disk, &disk_hash).await {
                                        if let MetaObj::Disk(disk_obj) = disk_meta {
                                            let mut partitions = Vec::new();

                                            for partition_hash in disk_obj.partition_hashes {
                                                if let Ok(Some(partition_meta)) = storage.load_meta(MetaType::Partition, &partition_hash).await {
                                                    if let MetaObj::Partition(partition_obj) = partition_meta {
                                                        let fs_type_str = match partition_obj.fs_type_code {
                                                            FsType::Ext => "ext",
                                                            FsType::Ntfs => "ntfs", 
                                                            FsType::Fat32 => "fat32",
                                                            FsType::Unknown => "unknown",
                                                        };

                                                        partitions.push(PartitionInfo {
                                                            fs_type: fs_type_str.to_string(),
                                                            start_lba: partition_obj.start_lba,
                                                            end_lba: partition_obj.end_lba,
                                                            size_bytes: (partition_obj.end_lba - partition_obj.start_lba) * 512,
                                                        });
                                                    }
                                                }
                                            }

                                            disks.push(DiskInfo {
                                                serial: disk_obj.serial,
                                                size_bytes: disk_obj.disk_size_bytes,
                                                partitions,
                                            });
                                        }
                                    }
                                }

                                // Convert timestamp to readable format
                                let created_at_str = DateTime::<Utc>::from_timestamp(created_at_timestamp as i64, 0)
                                    .map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
                                    .unwrap_or_else(|| "Unknown".to_string());

                                snapshot_infos.push(SnapshotInfo {
                                    id: snapshot_id,
                                    snapshot_hash: hex::encode(snapshot_hash),
                                    created_at: created_at_str,
                                    disks,
                                });
                            }
                        }
                    }
                }
            }
            Err(_) => {
                // If no snapshots directory exists, return empty list
                return Ok(Vec::new());
            }
        }

        // Sort snapshots by creation time (newest first)
        snapshot_infos.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        Ok(snapshot_infos)
    }

    pub async fn get_snapshot_details(
        pool: &DbPool,
        machine_id: i64,
        snapshot_id: String,
        user: &User,
    ) -> AppResult<SnapshotInfo> {
        // Verify machine access
        let machine = sqlx::query!(
            "SELECT id, user_id FROM machines WHERE id = ? AND user_id = ?",
            machine_id,
            user.id
        )
        .fetch_optional(pool)
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        if machine.is_none() {
            return Err(AppError::NotFoundError("Machine not found or access denied".to_string()));
        }

        let _machine = machine.unwrap();

        let storage = Storage::new("./data");

        // Load snapshot reference to get hash and timestamp
        let (snapshot_hash, created_at_timestamp) = storage.load_snapshot_ref(machine_id, &snapshot_id).await
            .map_err(|_| AppError::NotFoundError("Snapshot not found".to_string()))?
            .ok_or_else(|| AppError::NotFoundError("Snapshot not found".to_string()))?;

        // Load snapshot metadata
        let snapshot_meta = storage.load_meta(MetaType::Snapshot, &snapshot_hash).await
            .map_err(|_| AppError::NotFoundError("Snapshot metadata not found".to_string()))?
            .ok_or_else(|| AppError::NotFoundError("Snapshot metadata not found".to_string()))?;

        if let MetaObj::Snapshot(snapshot_obj) = snapshot_meta {
            let mut disks = Vec::new();

            for disk_hash in snapshot_obj.disk_hashes {
                if let Ok(Some(disk_meta)) = storage.load_meta(MetaType::Disk, &disk_hash).await {
                    if let MetaObj::Disk(disk_obj) = disk_meta {
                        let mut partitions = Vec::new();

                        for partition_hash in disk_obj.partition_hashes {
                            if let Ok(Some(partition_meta)) = storage.load_meta(MetaType::Partition, &partition_hash).await {
                                if let MetaObj::Partition(partition_obj) = partition_meta {
                                    let fs_type_str = match partition_obj.fs_type_code {
                                        FsType::Ext => "ext",
                                        FsType::Ntfs => "ntfs",
                                        FsType::Fat32 => "fat32", 
                                        FsType::Unknown => "unknown",
                                    };

                                    partitions.push(PartitionInfo {
                                        fs_type: fs_type_str.to_string(),
                                        start_lba: partition_obj.start_lba,
                                        end_lba: partition_obj.end_lba,
                                        size_bytes: (partition_obj.end_lba - partition_obj.start_lba) * 512,
                                    });
                                }
                            }
                        }

                        disks.push(DiskInfo {
                            serial: disk_obj.serial,
                            size_bytes: disk_obj.disk_size_bytes,
                            partitions,
                        });
                    }
                }
            }

            // Convert timestamp to readable format
            let created_at_str = DateTime::<Utc>::from_timestamp(created_at_timestamp as i64, 0)
                .map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
                .unwrap_or_else(|| "Unknown".to_string());

            Ok(SnapshotInfo {
                id: snapshot_id,
                snapshot_hash: hex::encode(snapshot_hash),
                created_at: created_at_str,
                disks,
            })
        } else {
            Err(AppError::ValidationError("Invalid snapshot metadata".to_string()))
        }
    }
}