aria2-rust-pro/crates/aria2-rust-pro-storage/src/store.rs

use std::{fs, io, path::PathBuf};

use crate::{
    control::{
        ControlMetadata, decode_control_metadata, encode_control_metadata, read_aria2_control_file,
        write_aria2_control_file,
    },
    model::{PieceIndex, PieceMap, PieceState},
    resume::{ResumeData, ResumeStore},
    session::{Aria2MetadataState, SessionFile, load_session_file, save_session_file},
};

/// Persists control metadata by download gid.
pub trait ControlStore {
    /// Concrete error type returned by the store implementation.
    type Error;

    /// Loads persisted control metadata for a download gid.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when control metadata cannot be read.
    fn load_control(&self, gid: &str) -> Result<Option<ControlMetadata>, Self::Error>;
    /// Persists control metadata for a download gid.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when control metadata cannot be written.
    fn save_control(&self, gid: &str, value: &ControlMetadata) -> Result<(), Self::Error>;
    /// Deletes control metadata for a download gid.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when control metadata cannot be removed.
    fn delete_control(&self, gid: &str) -> Result<(), Self::Error>;
}

/// Persists the session file.
pub trait SessionStore {
    /// Concrete error type returned by the store implementation.
    type Error;

    /// Loads the persisted session file.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when the session file cannot be read.
    fn load_session(&self) -> Result<SessionFile, Self::Error>;
    /// Persists the session file.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when the session file cannot be written.
    fn save_session(&self, session: &SessionFile) -> Result<(), Self::Error>;
}

/// Persists metadata and resume state keyed by gid.
pub trait MetadataStore {
    /// Concrete error type returned by the store implementation.
    type Error;

    /// Loads persisted metadata state for a gid.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when metadata state cannot be read.
    fn load_metadata_state(&self, gid: &str) -> Result<Option<Aria2MetadataState>, Self::Error>;
    /// Persists metadata state for a gid.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when metadata state cannot be written.
    fn save_metadata_state(&self, state: &Aria2MetadataState) -> Result<(), Self::Error>;
    /// Persists resume data through the metadata-oriented store surface.
    ///
    /// # Errors
    ///
    /// Returns the store-specific error when resume data cannot be written.
    fn save_resume_data(&self, resume: &ResumeData) -> Result<(), Self::Error>;
}

/// Error type used by the local filesystem-backed store.
#[derive(Debug)]
pub enum LocalStoreError {
    /// Raw filesystem I/O failure.
    Io(io::Error),
    /// Format or decoding failure.
    Parse(String),
}

impl From<io::Error> for LocalStoreError {
    fn from(value: io::Error) -> Self {
        Self::Io(value)
    }
}

/// Local filesystem-backed implementation of the storage traits.
#[derive(Debug)]
pub struct LocalFileStore {
    /// Root directory that contains the store layout.
    root: PathBuf,
}

impl LocalFileStore {
    /// Creates a store rooted at the provided directory.
    #[must_use]
    pub const fn new(root: PathBuf) -> Self {
        Self { root }
    }

    /// Creates the on-disk directory layout used by the local store.
    ///
    /// # Errors
    ///
    /// Returns [`LocalStoreError`] when any required directory cannot be created.
    pub fn ensure_layout(&self) -> Result<(), LocalStoreError> {
        fs::create_dir_all(self.controls_dir())?;
        fs::create_dir_all(self.metadata_dir())?;
        fs::create_dir_all(self.resume_dir())?;
        if let Some(parent) = self.session_path().parent() {
            fs::create_dir_all(parent)?;
        }
        Ok(())
    }

    /// Returns the control-file directory.
    fn controls_dir(&self) -> PathBuf {
        self.root.join("control")
    }

    /// Returns the metadata directory.
    fn metadata_dir(&self) -> PathBuf {
        self.root.join("metadata")
    }

    /// Returns the resume directory.
    fn resume_dir(&self) -> PathBuf {
        self.root.join("resume")
    }

    /// Returns the persisted session file path.
    fn session_path(&self) -> PathBuf {
        self.root.join("session").join("session.txt")
    }

    /// Returns the control-file path for one gid.
    fn control_path(&self, gid: &str) -> PathBuf {
        self.controls_dir().join(format!("{gid}.aria2"))
    }

    /// Returns the metadata path for one gid.
    fn metadata_path(&self, gid: &str) -> PathBuf {
        self.metadata_dir().join(format!("{gid}.meta"))
    }

    /// Returns the resume-file path for one gid.
    fn resume_path(&self, gid: &str) -> PathBuf {
        self.resume_dir().join(format!("{gid}.resume"))
    }
}

impl ControlStore for LocalFileStore {
    type Error = LocalStoreError;
    fn load_control(&self, gid: &str) -> Result<Option<ControlMetadata>, Self::Error> {
        let path = self.control_path(gid);
        if !path.exists() {
            return Ok(None);
        }
        read_aria2_control_file(&path)
            .map(Some)
            .map_err(|e| LocalStoreError::Parse(e.to_string()))
    }
    fn save_control(&self, gid: &str, value: &ControlMetadata) -> Result<(), Self::Error> {
        self.ensure_layout()?;
        write_aria2_control_file(&self.control_path(gid), value)
            .map_err(|e| LocalStoreError::Parse(e.to_string()))
    }
    fn delete_control(&self, gid: &str) -> Result<(), Self::Error> {
        let path = self.control_path(gid);
        if path.exists() {
            fs::remove_file(path)?;
        }
        Ok(())
    }
}

impl SessionStore for LocalFileStore {
    type Error = LocalStoreError;
    fn load_session(&self) -> Result<SessionFile, Self::Error> {
        let path = self.session_path();
        if !path.exists() {
            return Ok(SessionFile::default());
        }
        load_session_file(&path).map_err(LocalStoreError::Io)
    }
    fn save_session(&self, session: &SessionFile) -> Result<(), Self::Error> {
        self.ensure_layout()?;
        save_session_file(&self.session_path(), session).map_err(LocalStoreError::Io)
    }
}

impl MetadataStore for LocalFileStore {
    type Error = LocalStoreError;
    fn load_metadata_state(&self, gid: &str) -> Result<Option<Aria2MetadataState>, Self::Error> {
        let path = self.metadata_path(gid);
        if !path.exists() {
            return Ok(None);
        }
        let content = fs::read_to_string(path)?;
        let mut kv = std::collections::BTreeMap::new();
        for line in content.lines() {
            let mut parts = line.splitn(2, '=');
            let k = parts.next().unwrap_or_default();
            let v = parts.next().unwrap_or_default();
            if !k.is_empty() {
                kv.insert(k.to_owned(), v.to_owned());
            }
        }
        Ok(Some(Aria2MetadataState {
            gid: gid.to_owned(),
            kv,
        }))
    }
    fn save_metadata_state(&self, state: &Aria2MetadataState) -> Result<(), Self::Error> {
        self.ensure_layout()?;
        let body = state
            .kv
            .iter()
            .map(|(k, v)| format!("{k}={v}"))
            .collect::<Vec<_>>()
            .join("\n");
        fs::write(self.metadata_path(&state.gid), body)?;
        Ok(())
    }
    fn save_resume_data(&self, resume: &ResumeData) -> Result<(), Self::Error> {
        <Self as ResumeStore>::save(self, resume)
    }
}

impl ResumeStore for LocalFileStore {
    type Error = LocalStoreError;
    fn load(&self, gid: &str) -> Result<Option<ResumeData>, Self::Error> {
        let path = self.resume_path(gid);
        if !path.exists() {
            return Ok(None);
        }
        let content = fs::read_to_string(path)?;
        let mut download_path = None;
        let mut verified = Vec::new();
        let mut inflight = Vec::new();
        let mut failed = Vec::new();
        let mut metadata_blob = String::new();
        let mut in_metadata = false;
        for line in content.lines() {
            if line == "metadata<<" {
                in_metadata = true;
                continue;
            }
            if line == ">>metadata" {
                in_metadata = false;
                continue;
            }
            if in_metadata {
                metadata_blob.push_str(line);
                metadata_blob.push('\n');
                continue;
            }
            if let Some(v) = line.strip_prefix("download_path=") {
                download_path = Some(PathBuf::from(v));
            }
            if let Some(v) = line.strip_prefix("verified=") {
                verified = parse_csv_u32(v)?;
            }
            if let Some(v) = line.strip_prefix("inflight=") {
                inflight = parse_csv_u32(v)?;
            }
            if let Some(v) = line.strip_prefix("failed=") {
                failed = parse_csv_u32(v)?;
            }
        }
        let mut piece_map = PieceMap::default();
        for idx in verified {
            piece_map.set_state(PieceIndex(idx), PieceState::Verified);
        }
        for idx in inflight {
            piece_map.set_state(PieceIndex(idx), PieceState::InFlight);
        }
        for idx in failed {
            piece_map.set_state(PieceIndex(idx), PieceState::Failed);
        }
        let metadata = if metadata_blob.trim().is_empty() {
            None
        } else {
            Some(
                decode_control_metadata(metadata_blob.trim_end())
                    .map_err(|e| LocalStoreError::Parse(e.to_string()))?,
            )
        };
        Ok(Some(ResumeData {
            gid: gid.to_owned(),
            download_path: download_path.unwrap_or_default(),
            metadata,
            piece_map,
        }))
    }
    fn save(&self, resume: &ResumeData) -> Result<(), Self::Error> {
        self.ensure_layout()?;
        let mut lines = Vec::new();
        lines.push(format!(
            "download_path={}",
            resume.download_path.to_string_lossy()
        ));
        lines.push(format!(
            "verified={}",
            join_piece_indexes(&resume.piece_map, PieceState::Verified)
        ));
        lines.push(format!(
            "inflight={}",
            join_piece_indexes(&resume.piece_map, PieceState::InFlight)
        ));
        lines.push(format!(
            "failed={}",
            join_piece_indexes(&resume.piece_map, PieceState::Failed)
        ));
        if let Some(metadata) = &resume.metadata {
            lines.push("metadata<<".to_owned());
            lines.push(encode_control_metadata(metadata));
            lines.push(">>metadata".to_owned());
        }
        fs::write(self.resume_path(&resume.gid), lines.join("\n"))?;
        Ok(())
    }
    fn remove(&self, gid: &str) -> Result<(), Self::Error> {
        let path = self.resume_path(gid);
        if path.exists() {
            fs::remove_file(path)?;
        }
        Ok(())
    }
}

/// Parses a comma-separated list of piece indexes.
fn parse_csv_u32(raw: &str) -> Result<Vec<u32>, LocalStoreError> {
    if raw.trim().is_empty() {
        return Ok(Vec::new());
    }
    raw.split(',')
        .map(|s| {
            s.parse::<u32>()
                .map_err(|_| LocalStoreError::Parse(format!("invalid piece index: {s}")))
        })
        .collect()
}

/// Joins piece indexes in the provided state into a comma-separated string.
fn join_piece_indexes(piece_map: &PieceMap, target: PieceState) -> String {
    piece_map
        .iter()
        .filter(|(_, state)| **state == target)
        .map(|(idx, _)| idx.0.to_string())
        .collect::<Vec<_>>()
        .join(",")
}

#[cfg(test)]
mod tests {
    use std::{
        collections::BTreeMap,
        time::{SystemTime, UNIX_EPOCH},
    };

    use super::*;
    use crate::session::SessionFileEntry;

    fn temp_root() -> PathBuf {
        let millis = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_millis();
        std::env::temp_dir().join(format!("aria2-rust-pro-storage-test-{millis}"))
    }

    #[test]
    fn local_store_session_roundtrip() {
        let root = temp_root();
        let store = LocalFileStore::new(root.clone());
        let mut metadata = BTreeMap::new();
        metadata.insert("bt.name".to_owned(), "ubuntu".to_owned());
        let session = SessionFile {
            entries: vec![SessionFileEntry {
                gid: "gid-1".to_owned(),
                uri: "https://mirror-1.example/file.iso".to_owned(),
                uris: vec![
                    "https://mirror-1.example/file.iso".to_owned(),
                    "https://mirror-2.example/file.iso".to_owned(),
                ],
                target_path: PathBuf::from("D:/downloads/file.iso"),
                metadata_path: Some(PathBuf::from("D:/downloads/file.meta")),
                metadata: Some(metadata),
            }],
        };
        store.save_session(&session).unwrap();
        let loaded = store.load_session().unwrap();
        assert_eq!(loaded, session);
        let _ = fs::remove_dir_all(root);
    }

    #[test]
    fn local_store_resume_roundtrip() {
        let root = temp_root();
        let store = LocalFileStore::new(root.clone());
        let mut piece_map = PieceMap::default();
        piece_map.set_state(PieceIndex(0), PieceState::Verified);
        piece_map.set_state(PieceIndex(1), PieceState::InFlight);
        piece_map.set_state(PieceIndex(2), PieceState::Failed);
        let resume = ResumeData {
            gid: "gid-resume-1".to_owned(),
            download_path: PathBuf::from("D:/downloads/file.iso"),
            metadata: None,
            piece_map,
        };
        store.save(&resume).unwrap();
        let loaded = store.load(&resume.gid).unwrap().unwrap();
        assert_eq!(loaded, resume);
        let _ = fs::remove_dir_all(root);
    }
}