Comicviewer/src/zip_cache.rs

use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use std::rc::{Rc, Weak};

use anyhow::{Context, Result, bail};
use filetime::FileTime;
use gtk::gio;
use gtk::gio::prelude::*;
use gtk::glib;
use sha2::{Digest, Sha256};

use crate::browser::BrowserEntry;

#[derive(Clone)]
pub struct ArchiveSource {
    pub local_path: PathBuf,
    pub source_uri: String,
    pub display_name: String,
    _temporary: Option<Rc<tempfile::TempPath>>,
    _pin: Option<Rc<CachePin>>,
}

pub struct ZipCache {
    directory: PathBuf,
    quota_bytes: Cell<u64>,
    pinned: RefCell<HashMap<PathBuf, usize>>,
}

#[derive(Debug, Default, PartialEq, Eq)]
pub struct CacheCleanupReport {
    pub removed_files: usize,
    pub removed_bytes: u64,
    pub retained_open_files: usize,
    pub remaining_bytes: u64,
}

struct CachePin {
    path: PathBuf,
    cache: Weak<ZipCache>,
}

impl Drop for CachePin {
    fn drop(&mut self) {
        if let Some(cache) = self.cache.upgrade() {
            let mut pinned = cache.pinned.borrow_mut();
            let mut released = false;
            if let Some(count) = pinned.get_mut(&self.path) {
                *count -= 1;
                if *count == 0 {
                    pinned.remove(&self.path);
                    released = true;
                }
            }
            drop(pinned);
            if released {
                let _ = cache.enforce_quota(None);
            }
        }
    }
}

impl ZipCache {
    pub fn new(directory: PathBuf, quota_gb: u64) -> Result<Rc<Self>> {
        fs::create_dir_all(&directory)?;
        for entry in fs::read_dir(&directory)? {
            let path = entry?.path();
            let incomplete_download = path
                .extension()
                .is_some_and(|extension| extension == "part");
            let stale_temporary = path
                .file_name()
                .and_then(|name| name.to_str())
                .is_some_and(|name| name.starts_with("temporary-"));
            if incomplete_download || stale_temporary {
                let _ = fs::remove_file(path);
            }
        }
        let cache = Rc::new(Self {
            directory,
            quota_bytes: Cell::new(quota_gb.saturating_mul(1024 * 1024 * 1024)),
            pinned: RefCell::new(HashMap::new()),
        });
        cache.enforce_quota(None)?;
        Ok(cache)
    }

    pub fn prepare<F, P>(self: &Rc<Self>, entry: BrowserEntry, progress: P, callback: F)
    where
        F: FnOnce(Result<ArchiveSource>) + 'static,
        P: FnMut(i64, i64) + 'static,
    {
        let source = gio::File::for_uri(&entry.uri);
        if let Some(path) = native_path(&source) {
            callback(Ok(ArchiveSource {
                local_path: path,
                source_uri: entry.uri,
                display_name: entry.name,
                _temporary: None,
                _pin: None,
            }));
            return;
        }

        let quota_bytes = self.quota_bytes.get();
        if quota_bytes == 0 || entry.size.is_some_and(|size| size > quota_bytes) {
            self.download_temporary(source, entry, true, progress, callback);
        } else {
            self.download_persistent(source, entry, progress, callback);
        }
    }

    pub fn prepare_media<F, P>(self: &Rc<Self>, entry: BrowserEntry, progress: P, callback: F)
    where
        F: FnOnce(Result<ArchiveSource>) + 'static,
        P: FnMut(i64, i64) + 'static,
    {
        let source = gio::File::for_uri(&entry.uri);
        if let Some(path) = native_path(&source) {
            callback(Ok(ArchiveSource {
                local_path: path,
                source_uri: entry.uri,
                display_name: entry.name,
                _temporary: None,
                _pin: None,
            }));
            return;
        }
        self.download_temporary(source, entry, false, progress, callback);
    }

    fn download_temporary<F, P>(
        self: &Rc<Self>,
        source: gio::File,
        entry: BrowserEntry,
        validate_zip: bool,
        progress: P,
        callback: F,
    ) where
        F: FnOnce(Result<ArchiveSource>) + 'static,
        P: FnMut(i64, i64) + 'static,
    {
        let suffix = Path::new(&entry.name)
            .extension()
            .and_then(|extension| extension.to_str())
            .map_or_else(
                || ".part".to_owned(),
                |extension| format!(".part.{extension}"),
            );
        let temporary = match tempfile::Builder::new()
            .prefix("temporary-")
            .suffix(&suffix)
            .tempfile_in(&self.directory)
        {
            Ok(file) => Rc::new(file.into_temp_path()),
            Err(error) => {
                callback(Err(error.into()));
                return;
            }
        };
        let destination = gio::File::for_path(temporary.as_ref());
        let retained = Rc::clone(&temporary);
        let expected_size = entry.size;
        source.copy_async(
            &destination,
            gio::FileCopyFlags::OVERWRITE,
            glib::Priority::DEFAULT,
            gio::Cancellable::NONE,
            Some(Box::new(progress)),
            move |result| {
                callback(result.map_err(Into::into).and_then(|()| {
                    validate_size(retained.as_ref(), expected_size)?;
                    if validate_zip {
                        validate_zip_file(retained.as_ref())?;
                    }
                    Ok(ArchiveSource {
                        local_path: retained.to_path_buf(),
                        source_uri: entry.uri,
                        display_name: entry.name,
                        _temporary: Some(retained),
                        _pin: None,
                    })
                }));
            },
        );
    }

    fn download_persistent<F, P>(
        self: &Rc<Self>,
        source: gio::File,
        entry: BrowserEntry,
        progress: P,
        callback: F,
    ) where
        F: FnOnce(Result<ArchiveSource>) + 'static,
        P: FnMut(i64, i64) + 'static,
    {
        let key = cache_key(&entry);
        let final_path = self.directory.join(format!("{key}.zip"));
        if final_path.is_file() {
            if validate_download(&final_path, entry.size).is_ok() {
                let _ = filetime::set_file_mtime(&final_path, FileTime::now());
                callback(Ok(self.lease(final_path, entry.uri, entry.name)));
                return;
            }
            let _ = fs::remove_file(&final_path);
        }

        let part_path = self.directory.join(format!("{key}.part"));
        if let Some(required_bytes) = entry.size
            && let Err(error) = self.make_room_for(required_bytes)
        {
            callback(Err(error));
            return;
        }
        let destination = gio::File::for_path(&part_path);
        let cache = Rc::clone(self);
        source.copy_async(
            &destination,
            gio::FileCopyFlags::OVERWRITE,
            glib::Priority::DEFAULT,
            gio::Cancellable::NONE,
            Some(Box::new(progress)),
            move |result| {
                let result = result.map_err(Into::into).and_then(|()| {
                    validate_download(&part_path, entry.size)?;
                    fs::rename(&part_path, &final_path).with_context(|| {
                        format!(
                            "다운로드 파일을 완료 처리할 수 없습니다: {}",
                            final_path.display()
                        )
                    })?;
                    cache.enforce_quota(Some(&final_path))?;
                    Ok(cache.lease(final_path, entry.uri, entry.name))
                });
                if result.is_err() {
                    let _ = fs::remove_file(&part_path);
                }
                callback(result);
            },
        );
    }

    pub fn clear_persistent(&self) -> Result<CacheCleanupReport> {
        self.cleanup_to_limit(0, None)
    }

    fn make_room_for(&self, required_bytes: u64) -> Result<()> {
        let limit = self.quota_bytes.get().saturating_sub(required_bytes);
        let report = self.cleanup_to_limit(limit, None)?;
        if report.remaining_bytes > limit {
            bail!("열려 있는 ZIP 때문에 다운로드에 필요한 캐시 공간을 확보할 수 없습니다");
        }
        Ok(())
    }

    fn enforce_quota(&self, protected: Option<&Path>) -> Result<CacheCleanupReport> {
        self.cleanup_to_limit(self.quota_bytes.get(), protected)
    }

    fn cleanup_to_limit(&self, limit: u64, protected: Option<&Path>) -> Result<CacheCleanupReport> {
        let mut files = Vec::new();
        let mut total = 0_u64;
        for entry in fs::read_dir(&self.directory)? {
            let path = entry?.path();
            if path
                .file_name()
                .and_then(|name| name.to_str())
                .is_some_and(|name| name.starts_with("temporary-"))
            {
                continue;
            }
            if path.extension().is_none_or(|extension| extension != "zip") {
                continue;
            }
            let metadata = fs::metadata(&path)?;
            total = total.saturating_add(metadata.len());
            files.push((metadata.modified().ok(), metadata.len(), path));
        }
        files.sort_by_key(|(modified, _, path)| (*modified, path.clone()));
        let mut report = CacheCleanupReport::default();
        for (_, size, path) in files {
            if total <= limit {
                break;
            }
            if protected == Some(path.as_path()) {
                continue;
            }
            if self.pinned.borrow().contains_key(&path) {
                report.retained_open_files += 1;
                continue;
            }
            fs::remove_file(&path)
                .with_context(|| format!("캐시 파일을 삭제할 수 없습니다: {}", path.display()))?;
            total = total.saturating_sub(size);
            report.removed_files += 1;
            report.removed_bytes = report.removed_bytes.saturating_add(size);
        }
        report.remaining_bytes = total;
        Ok(report)
    }

    fn lease(
        self: &Rc<Self>,
        path: PathBuf,
        source_uri: String,
        display_name: String,
    ) -> ArchiveSource {
        *self.pinned.borrow_mut().entry(path.clone()).or_default() += 1;
        ArchiveSource {
            local_path: path.clone(),
            source_uri,
            display_name,
            _temporary: None,
            _pin: Some(Rc::new(CachePin {
                path,
                cache: Rc::downgrade(self),
            })),
        }
    }
}

fn native_path(file: &gio::File) -> Option<PathBuf> {
    file.is_native().then(|| file.path()).flatten()
}

fn validate_download(path: &Path, expected_size: Option<u64>) -> Result<()> {
    validate_size(path, expected_size)?;
    validate_zip_file(path)
}

fn validate_size(path: &Path, expected_size: Option<u64>) -> Result<()> {
    let actual_size = fs::metadata(path)?.len();
    if let Some(expected_size) = expected_size
        && actual_size != expected_size
    {
        bail!(
            "다운로드 크기가 원본과 다릅니다: 예상 {expected_size}바이트, 실제 {actual_size}바이트"
        );
    }
    Ok(())
}

fn validate_zip_file(path: &Path) -> Result<()> {
    let file = File::open(path)?;
    zip::ZipArchive::new(file).context("다운로드된 파일이 올바른 ZIP이 아닙니다")?;
    Ok(())
}

fn cache_key(entry: &BrowserEntry) -> String {
    let mut hasher = Sha256::new();
    hasher.update(entry.uri.as_bytes());
    hasher.update(entry.size.unwrap_or_default().to_le_bytes());
    hasher.update(
        entry
            .modified
            .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
            .map_or(0, |duration| duration.as_secs())
            .to_le_bytes(),
    );
    if let Some(etag) = &entry.etag {
        hasher.update(etag.as_bytes());
    }
    format!("{:x}", hasher.finalize())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::browser::EntryKind;

    #[test]
    fn cache_key_changes_with_remote_metadata() {
        let mut entry = BrowserEntry {
            uri: "sftp://example.invalid/a.zip".into(),
            name: "a.zip".into(),
            list_name: None,
            kind: EntryKind::ZipArchive,
            size: Some(10),
            modified: None,
            created: None,
            etag: None,
        };
        let first = cache_key(&entry);
        entry.size = Some(11);
        assert_ne!(first, cache_key(&entry));
    }

    #[test]
    fn remote_gio_files_are_not_opened_as_local_archives() {
        let remote = gio::File::for_uri("ftp://example.invalid/comic.zip");
        let local = gio::File::for_path("/tmp/comic.zip");

        assert!(native_path(&remote).is_none());
        assert_eq!(native_path(&local), Some(PathBuf::from("/tmp/comic.zip")));
    }

    #[test]
    fn invalid_completed_download_is_rejected() {
        let directory = tempfile::tempdir().unwrap();
        let path = directory.path().join("invalid.zip");
        fs::write(&path, b"not a zip").unwrap();

        assert!(validate_download(&path, Some(9)).is_err());
        assert!(validate_download(&path, Some(1)).is_err());
    }

    #[test]
    fn lru_does_not_remove_an_open_archive() {
        let directory = tempfile::tempdir().unwrap();
        let first = directory.path().join("first.zip");
        let second = directory.path().join("second.zip");
        fs::write(&first, [0_u8; 8]).unwrap();
        fs::write(&second, [0_u8; 8]).unwrap();
        let cache = Rc::new(ZipCache {
            directory: directory.path().to_path_buf(),
            quota_bytes: Cell::new(8),
            pinned: RefCell::new(HashMap::new()),
        });

        let lease = cache.lease(
            first.clone(),
            "file:///first.zip".into(),
            "first.zip".into(),
        );
        cache.enforce_quota(Some(&second)).unwrap();
        assert!(first.exists());
        assert!(second.exists());

        drop(lease);
        cache.enforce_quota(Some(&second)).unwrap();
        assert!(!first.exists());
        assert!(second.exists());
    }

    #[test]
    fn lru_removes_oldest_cache_file_first() {
        let directory = tempfile::tempdir().unwrap();
        let oldest = directory.path().join("oldest.zip");
        let newest = directory.path().join("newest.zip");
        fs::write(&oldest, [0_u8; 8]).unwrap();
        fs::write(&newest, [0_u8; 8]).unwrap();
        filetime::set_file_mtime(&oldest, FileTime::from_unix_time(1, 0)).unwrap();
        filetime::set_file_mtime(&newest, FileTime::from_unix_time(2, 0)).unwrap();
        let cache = ZipCache {
            directory: directory.path().to_path_buf(),
            quota_bytes: Cell::new(8),
            pinned: RefCell::new(HashMap::new()),
        };

        let report = cache.enforce_quota(None).unwrap();

        assert_eq!(report.removed_files, 1);
        assert!(!oldest.exists());
        assert!(newest.exists());
    }

    #[test]
    fn download_reserves_space_before_writing() {
        let directory = tempfile::tempdir().unwrap();
        let oldest = directory.path().join("oldest.zip");
        let newest = directory.path().join("newest.zip");
        fs::write(&oldest, [0_u8; 8]).unwrap();
        fs::write(&newest, [0_u8; 8]).unwrap();
        filetime::set_file_mtime(&oldest, FileTime::from_unix_time(1, 0)).unwrap();
        filetime::set_file_mtime(&newest, FileTime::from_unix_time(2, 0)).unwrap();
        let cache = ZipCache {
            directory: directory.path().to_path_buf(),
            quota_bytes: Cell::new(16),
            pinned: RefCell::new(HashMap::new()),
        };

        cache.make_room_for(8).unwrap();

        assert!(!oldest.exists());
        assert!(newest.exists());
    }

    #[test]
    fn zero_quota_clears_existing_cache_on_startup() {
        let directory = tempfile::tempdir().unwrap();
        let cached = directory.path().join("cached.zip");
        fs::write(&cached, [0_u8; 8]).unwrap();

        let _cache = ZipCache::new(directory.path().to_path_buf(), 0).unwrap();

        assert!(!cached.exists());
    }

    #[test]
    fn startup_removes_stale_temporary_downloads() {
        let directory = tempfile::tempdir().unwrap();
        let temporary = directory.path().join("temporary-abcd.part.mp4");
        fs::write(&temporary, [0_u8; 8]).unwrap();

        let _cache = ZipCache::new(directory.path().to_path_buf(), 10).unwrap();

        assert!(!temporary.exists());
    }

    #[test]
    fn manual_clear_keeps_open_archives() {
        let directory = tempfile::tempdir().unwrap();
        let open = directory.path().join("open.zip");
        let unused = directory.path().join("unused.zip");
        let temporary = directory.path().join("temporary-download.part.zip");
        fs::write(&open, [0_u8; 8]).unwrap();
        fs::write(&unused, [0_u8; 8]).unwrap();
        fs::write(&temporary, [0_u8; 8]).unwrap();
        let cache = Rc::new(ZipCache {
            directory: directory.path().to_path_buf(),
            quota_bytes: Cell::new(16),
            pinned: RefCell::new(HashMap::new()),
        });
        let lease = cache.lease(open.clone(), "file:///open.zip".into(), "open.zip".into());

        let report = cache.clear_persistent().unwrap();

        assert_eq!(report.removed_files, 1);
        assert_eq!(report.retained_open_files, 1);
        assert!(open.exists());
        assert!(!unused.exists());
        assert!(temporary.exists());
        drop(lease);
    }
}