vflight 0.9.2

//! Resume state management for interrupted downloads.
//!
//! This module provides persistence for download progress, enabling interrupted
//! transfers to resume from where they left off instead of starting over.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::fs;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::{debug, info, warn};

use crate::protocol::FileMetadata;

/// Version of the resume file format for forward compatibility.
pub const RESUME_VERSION: u32 = 1;

/// Resume state directory name (hidden).
const RESUME_DIR_NAME: &str = ".vflight-resume";

/// Persistent state for a resumable download.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResumeState {
    /// Version of the resume file format.
    pub version: u32,
    /// DHT key being fetched (for validation).
    pub dht_key: String,
    /// Snapshot of file metadata for validation on resume.
    pub metadata: ResumeMetadata,
    /// Set of chunk indices successfully downloaded and verified.
    pub completed_chunks: HashSet<u64>,
    /// Unix timestamp of last update.
    pub last_updated: u64,
}

/// Snapshot of file metadata for validation on resume.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ResumeMetadata {
    /// Original filename.
    pub name: String,
    /// Total file size in bytes.
    pub size: u64,
    /// Number of chunks.
    pub total_chunks: u64,
    /// BLAKE3 hashes for each chunk (for verification).
    pub chunk_hashes: Vec<String>,
    /// Whether the file is encrypted.
    pub encrypted: bool,
}

impl ResumeState {
    /// Create a new resume state from DHT key and file metadata.
    pub fn new(dht_key: &str, metadata: &FileMetadata) -> Self {
        Self {
            version: RESUME_VERSION,
            dht_key: dht_key.to_string(),
            metadata: ResumeMetadata::from(metadata),
            completed_chunks: HashSet::new(),
            last_updated: current_timestamp(),
        }
    }

    /// Returns indices of chunks still needed.
    pub fn pending_chunks(&self) -> Vec<u64> {
        (0..self.metadata.total_chunks)
            .filter(|i| !self.completed_chunks.contains(i))
            .collect()
    }

    /// Check if download is complete.
    pub fn is_complete(&self) -> bool {
        self.completed_chunks.len() as u64 == self.metadata.total_chunks
    }

    /// Number of completed chunks.
    pub fn completed_count(&self) -> u64 {
        self.completed_chunks.len() as u64
    }
}

impl From<&FileMetadata> for ResumeMetadata {
    fn from(metadata: &FileMetadata) -> Self {
        Self {
            name: metadata.name.clone(),
            size: metadata.size,
            total_chunks: metadata.total_chunks,
            chunk_hashes: metadata.chunk_hashes.clone(),
            encrypted: metadata.encryption_salt.is_some(),
        }
    }
}

impl ResumeMetadata {
    /// Check if this metadata matches current file metadata.
    pub fn matches(&self, metadata: &FileMetadata) -> bool {
        self.name == metadata.name
            && self.size == metadata.size
            && self.total_chunks == metadata.total_chunks
            && self.chunk_hashes == metadata.chunk_hashes
    }
}

/// Manages resume state and chunk files on disk.
pub struct ResumeManager {
    /// Base directory for this download's resume data.
    resume_dir: PathBuf,
    /// Path to resume.json.
    state_path: PathBuf,
    /// Directory containing chunk files.
    chunks_dir: PathBuf,
    /// Current state.
    state: ResumeState,
}

impl ResumeManager {
    /// Initialize resume manager, loading existing state if valid or creating fresh.
    ///
    /// This is the primary entry point. It will:
    /// 1. Check for existing resume state
    /// 2. Validate it matches the current DHT key and metadata
    /// 3. Return existing state if valid, or create fresh state
    pub fn init(output_dir: &Path, dht_key: &str, metadata: &FileMetadata) -> Result<Self> {
        let resume_dir = Self::resume_dir_path(output_dir, dht_key);
        let state_path = resume_dir.join("resume.json");
        let chunks_dir = resume_dir.join("chunks");

        // Try to load existing state
        if state_path.exists() {
            match Self::load_and_validate(&state_path, dht_key, metadata) {
                Ok(state) => {
                    info!(
                        completed = state.completed_count(),
                        total = state.metadata.total_chunks,
                        "Loaded existing resume state"
                    );
                    return Ok(Self {
                        resume_dir,
                        state_path,
                        chunks_dir,
                        state,
                    });
                }
                Err(e) => {
                    warn!(error = %e, "Invalid resume state, starting fresh");
                    // Clean up invalid state
                    if resume_dir.exists() {
                        fs::remove_dir_all(&resume_dir).ok();
                    }
                }
            }
        }

        // Create fresh state
        Self::new_fresh(output_dir, dht_key, metadata)
    }

    /// Create a new fresh resume manager, ignoring any existing state.
    pub fn new_fresh(output_dir: &Path, dht_key: &str, metadata: &FileMetadata) -> Result<Self> {
        let resume_dir = Self::resume_dir_path(output_dir, dht_key);
        let state_path = resume_dir.join("resume.json");
        let chunks_dir = resume_dir.join("chunks");

        // Clean up any existing state
        if resume_dir.exists() {
            fs::remove_dir_all(&resume_dir)
                .context("Failed to remove existing resume directory")?;
        }

        // Create directories
        fs::create_dir_all(&chunks_dir).context("Failed to create resume chunks directory")?;

        let state = ResumeState::new(dht_key, metadata);

        let manager = Self {
            resume_dir,
            state_path,
            chunks_dir,
            state,
        };

        // Save initial state
        manager.save_state()?;

        debug!("Created fresh resume state");
        Ok(manager)
    }

    /// Get list of pending (not yet downloaded) chunk indices.
    pub fn pending_chunks(&self) -> Vec<u64> {
        self.state.pending_chunks()
    }

    /// Get number of completed chunks.
    pub fn completed_count(&self) -> u64 {
        self.state.completed_count()
    }

    /// Get total chunk count.
    pub fn total_chunks(&self) -> u64 {
        self.state.metadata.total_chunks
    }

    /// Check if file was encrypted.
    pub fn is_encrypted(&self) -> bool {
        self.state.metadata.encrypted
    }

    /// Save a downloaded chunk to disk.
    pub fn save_chunk(&self, index: u64, data: &[u8]) -> Result<()> {
        let chunk_path = self.chunk_path(index);

        // Write to temp file then rename for atomicity
        let temp_path = chunk_path.with_extension("tmp");
        fs::write(&temp_path, data).context("Failed to write chunk file")?;
        fs::rename(&temp_path, &chunk_path).context("Failed to rename chunk file")?;

        debug!(chunk = index, "Saved chunk to disk");
        Ok(())
    }

    /// Mark a chunk as completed and update state file.
    pub fn mark_chunk_complete(&mut self, index: u64) -> Result<()> {
        self.state.completed_chunks.insert(index);
        self.state.last_updated = current_timestamp();
        self.save_state()?;
        Ok(())
    }

    /// Save chunk and mark complete in one operation.
    pub fn save_and_mark_complete(&mut self, index: u64, data: &[u8]) -> Result<()> {
        self.save_chunk(index, data)?;
        self.mark_chunk_complete(index)?;
        Ok(())
    }

    /// Load a chunk from disk.
    pub fn load_chunk(&self, index: u64) -> Result<Vec<u8>> {
        let chunk_path = self.chunk_path(index);
        fs::read(&chunk_path).with_context(|| format!("Failed to read chunk {} from cache", index))
    }

    /// Verify a cached chunk against its expected hash.
    pub fn verify_chunk(&self, index: u64, expected_hash: &str) -> Result<bool> {
        let chunk_path = self.chunk_path(index);
        if !chunk_path.exists() {
            return Ok(false);
        }

        let data = fs::read(&chunk_path)?;
        let computed_hash = blake3::hash(&data).to_hex().to_string();
        Ok(computed_hash == expected_hash)
    }

    /// Verify all completed chunks and remove any that are corrupted.
    /// Returns list of chunk indices that were invalidated.
    pub fn verify_and_invalidate_corrupted(&mut self) -> Result<Vec<u64>> {
        let mut invalidated = Vec::new();

        // Collect completed chunks to check (avoid borrowing issues)
        let completed: Vec<u64> = self.state.completed_chunks.iter().copied().collect();

        for index in completed {
            let expected_hash = self
                .state
                .metadata
                .chunk_hashes
                .get(index as usize)
                .map(|s| s.as_str())
                .unwrap_or("");

            match self.verify_chunk(index, expected_hash) {
                Ok(true) => {
                    // Chunk is valid
                }
                Ok(false) | Err(_) => {
                    // Chunk is missing or corrupted
                    warn!(
                        chunk = index,
                        "Cached chunk corrupted or missing, will re-download"
                    );
                    self.state.completed_chunks.remove(&index);
                    // Remove the corrupted file
                    let chunk_path = self.chunk_path(index);
                    fs::remove_file(&chunk_path).ok();
                    invalidated.push(index);
                }
            }
        }

        if !invalidated.is_empty() {
            self.state.last_updated = current_timestamp();
            self.save_state()?;
        }

        Ok(invalidated)
    }

    /// Clean up all resume state (call on successful completion).
    pub fn cleanup(self) -> Result<()> {
        if self.resume_dir.exists() {
            fs::remove_dir_all(&self.resume_dir).context("Failed to remove resume directory")?;
            debug!("Cleaned up resume state");
        }
        Ok(())
    }

    /// Get the path where a chunk file is stored.
    fn chunk_path(&self, index: u64) -> PathBuf {
        self.chunks_dir.join(format!("chunk_{:06}.bin", index))
    }

    /// Get the resume directory path for a given DHT key.
    fn resume_dir_path(output_dir: &Path, dht_key: &str) -> PathBuf {
        // Use first 16 chars of blake3 hash of DHT key for directory name
        let hash = blake3::hash(dht_key.as_bytes()).to_hex();
        let prefix = &hash[..16];
        output_dir.join(RESUME_DIR_NAME).join(prefix)
    }

    /// Save state to disk atomically.
    fn save_state(&self) -> Result<()> {
        let json = serde_json::to_string_pretty(&self.state)?;
        let temp_path = self.state_path.with_extension("tmp");
        fs::write(&temp_path, &json).context("Failed to write resume state")?;
        fs::rename(&temp_path, &self.state_path).context("Failed to rename resume state")?;
        Ok(())
    }

    /// Load and validate existing state.
    fn load_and_validate(
        state_path: &Path,
        dht_key: &str,
        metadata: &FileMetadata,
    ) -> Result<ResumeState> {
        let json = fs::read_to_string(state_path).context("Failed to read resume state")?;
        let state: ResumeState =
            serde_json::from_str(&json).context("Failed to parse resume state")?;

        // Validate version
        if state.version != RESUME_VERSION {
            anyhow::bail!(
                "Resume state version mismatch: expected {}, got {}",
                RESUME_VERSION,
                state.version
            );
        }

        // Validate DHT key
        if state.dht_key != dht_key {
            anyhow::bail!("Resume state DHT key mismatch");
        }

        // Validate metadata matches
        if !state.metadata.matches(metadata) {
            anyhow::bail!("File metadata has changed since last download attempt");
        }

        Ok(state)
    }
}

/// Get current Unix timestamp.
fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

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

    fn create_test_metadata(chunks: u64) -> FileMetadata {
        let chunk_hashes: Vec<String> = (0..chunks)
            .map(|i| {
                blake3::hash(format!("chunk{}", i).as_bytes())
                    .to_hex()
                    .to_string()
            })
            .collect();

        FileMetadata {
            name: "test_file.bin".to_string(),
            size: chunks * 30000,
            total_chunks: chunks,
            chunk_hashes,
            route_blob: String::new(),
            encryption_salt: None,
            encryption_nonce: None,
            compressed: false,
        }
    }

    #[test]
    fn test_resume_state_new() {
        let metadata = create_test_metadata(10);
        let state = ResumeState::new("VLD0:test_key", &metadata);

        assert_eq!(state.version, RESUME_VERSION);
        assert_eq!(state.dht_key, "VLD0:test_key");
        assert_eq!(state.metadata.total_chunks, 10);
        assert!(state.completed_chunks.is_empty());
        assert!(!state.is_complete());
    }

    #[test]
    fn test_resume_state_pending_chunks() {
        let metadata = create_test_metadata(5);
        let mut state = ResumeState::new("VLD0:test", &metadata);

        // Initially all chunks pending
        assert_eq!(state.pending_chunks(), vec![0, 1, 2, 3, 4]);

        // Mark some complete
        state.completed_chunks.insert(1);
        state.completed_chunks.insert(3);

        assert_eq!(state.pending_chunks(), vec![0, 2, 4]);
        assert_eq!(state.completed_count(), 2);
        assert!(!state.is_complete());

        // Mark all complete
        state.completed_chunks.insert(0);
        state.completed_chunks.insert(2);
        state.completed_chunks.insert(4);

        assert!(state.pending_chunks().is_empty());
        assert!(state.is_complete());
    }

    #[test]
    fn test_resume_state_serialization() {
        let metadata = create_test_metadata(3);
        let mut state = ResumeState::new("VLD0:test", &metadata);
        state.completed_chunks.insert(0);
        state.completed_chunks.insert(2);

        // Serialize
        let json = serde_json::to_string(&state).unwrap();

        // Deserialize
        let restored: ResumeState = serde_json::from_str(&json).unwrap();

        assert_eq!(restored.version, state.version);
        assert_eq!(restored.dht_key, state.dht_key);
        assert_eq!(restored.metadata, state.metadata);
        assert_eq!(restored.completed_chunks, state.completed_chunks);
    }

    #[test]
    fn test_resume_metadata_matches() {
        let metadata = create_test_metadata(5);
        let resume_meta = ResumeMetadata::from(&metadata);

        assert!(resume_meta.matches(&metadata));

        // Different size
        let mut different = metadata.clone();
        different.size = 999;
        assert!(!resume_meta.matches(&different));

        // Different chunks
        let mut different = metadata.clone();
        different.total_chunks = 10;
        assert!(!resume_meta.matches(&different));
    }

    #[test]
    fn test_resume_manager_init_fresh() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(5);

        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();

        assert_eq!(manager.pending_chunks().len(), 5);
        assert_eq!(manager.completed_count(), 0);

        // State file should exist
        assert!(manager.state_path.exists());
    }

    #[test]
    fn test_resume_manager_init_existing() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(5);

        // Create initial manager and save some progress
        {
            let mut manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();
            manager.save_chunk(0, b"chunk0data").unwrap();
            manager.mark_chunk_complete(0).unwrap();
            manager.save_chunk(2, b"chunk2data").unwrap();
            manager.mark_chunk_complete(2).unwrap();
        }

        // Create new manager - should load existing state
        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();

        assert_eq!(manager.completed_count(), 2);
        assert_eq!(manager.pending_chunks(), vec![1, 3, 4]);
    }

    #[test]
    fn test_resume_manager_init_corrupted() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(5);

        // Create resume directory with invalid state at the actual path init() will check
        let hash = blake3::hash("VLD0:test".as_bytes()).to_hex();
        let resume_dir = temp.path().join(RESUME_DIR_NAME).join(&hash[..16]);
        fs::create_dir_all(&resume_dir).unwrap();
        fs::write(resume_dir.join("resume.json"), "invalid json").unwrap();

        // Should start fresh
        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();
        assert_eq!(manager.completed_count(), 0);
    }

    #[test]
    fn test_resume_manager_save_and_load_chunk() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(3);

        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();

        let test_data = b"test chunk data for verification";
        manager.save_chunk(1, test_data).unwrap();

        let loaded = manager.load_chunk(1).unwrap();
        assert_eq!(loaded, test_data);
    }

    #[test]
    fn test_resume_manager_verify_chunk_valid() {
        let temp = tempdir().unwrap();
        let test_data = b"test chunk data";
        let hash = blake3::hash(test_data).to_hex().to_string();

        let metadata = FileMetadata {
            name: "test.bin".to_string(),
            size: test_data.len() as u64,
            total_chunks: 1,
            chunk_hashes: vec![hash.clone()],
            route_blob: String::new(),
            encryption_salt: None,
            encryption_nonce: None,
            compressed: false,
        };

        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();
        manager.save_chunk(0, test_data).unwrap();

        assert!(manager.verify_chunk(0, &hash).unwrap());
    }

    #[test]
    fn test_resume_manager_verify_chunk_corrupted() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(1);

        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();

        // Save wrong data
        manager.save_chunk(0, b"wrong data").unwrap();

        // Should fail verification
        let expected_hash = &metadata.chunk_hashes[0];
        assert!(!manager.verify_chunk(0, expected_hash).unwrap());
    }

    #[test]
    fn test_resume_manager_verify_and_invalidate() {
        let temp = tempdir().unwrap();
        let good_data = b"good chunk data";
        let good_hash = blake3::hash(good_data).to_hex().to_string();

        let metadata = FileMetadata {
            name: "test.bin".to_string(),
            size: 60000,
            total_chunks: 2,
            chunk_hashes: vec![good_hash.clone(), "badhash".to_string()],
            route_blob: String::new(),
            encryption_salt: None,
            encryption_nonce: None,
            compressed: false,
        };

        let mut manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();

        // Save one good chunk and one bad
        manager.save_chunk(0, good_data).unwrap();
        manager.mark_chunk_complete(0).unwrap();
        manager.save_chunk(1, b"bad data").unwrap();
        manager.mark_chunk_complete(1).unwrap();

        assert_eq!(manager.completed_count(), 2);

        // Verify - should invalidate chunk 1
        let invalidated = manager.verify_and_invalidate_corrupted().unwrap();

        assert_eq!(invalidated, vec![1]);
        assert_eq!(manager.completed_count(), 1);
        assert!(manager.state.completed_chunks.contains(&0));
        assert!(!manager.state.completed_chunks.contains(&1));
    }

    #[test]
    fn test_resume_manager_cleanup() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(3);

        let manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();
        manager.save_chunk(0, b"data").unwrap();

        let resume_dir = manager.resume_dir.clone();
        assert!(resume_dir.exists());

        manager.cleanup().unwrap();
        assert!(!resume_dir.exists());
    }

    #[test]
    fn test_resume_manager_atomic_state_update() {
        let temp = tempdir().unwrap();
        let metadata = create_test_metadata(5);

        let mut manager = ResumeManager::init(temp.path(), "VLD0:test", &metadata).unwrap();

        // Multiple updates
        for i in 0..5 {
            manager
                .save_chunk(i, format!("chunk{}", i).as_bytes())
                .unwrap();
            manager.mark_chunk_complete(i).unwrap();

            // Verify state file is valid after each update
            let json = fs::read_to_string(&manager.state_path).unwrap();
            let state: ResumeState = serde_json::from_str(&json).unwrap();
            assert_eq!(state.completed_count(), i + 1);
        }
    }

    #[test]
    fn test_resume_metadata_encrypted() {
        let mut metadata = create_test_metadata(3);
        metadata.encryption_salt = Some("salt".to_string());
        metadata.encryption_nonce = Some("nonce".to_string());

        let resume_meta = ResumeMetadata::from(&metadata);
        assert!(resume_meta.encrypted);

        // Without encryption
        let metadata = create_test_metadata(3);
        let resume_meta = ResumeMetadata::from(&metadata);
        assert!(!resume_meta.encrypted);
    }
}