//! Local chunk storage

use dashmap::DashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use thiserror::Error;
use tokio::fs;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::debug;

/// Errors from chunk storage operations
#[derive(Debug, Error)]
pub enum StorageError {
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("Chunk not found: {0}")]
    NotFound(String),

    #[error("Storage capacity exceeded")]
    CapacityExceeded,
}

pub type StorageResult<T> = Result<T, StorageError>;

/// Local filesystem-based chunk storage
pub struct LocalChunkStore {
    /// Data directory
    data_dir: PathBuf,

    /// In-memory index of chunk sizes for fast lookups
    chunk_sizes: DashMap<String, u64>,

    /// Total bytes stored
    total_bytes: AtomicU64,

    /// Maximum capacity (0 = unlimited)
    max_capacity: u64,

    /// Number of chunks stored
    chunk_count: AtomicU64,
}

impl LocalChunkStore {
    /// Create a new local chunk store
    pub async fn new(data_dir: PathBuf, max_capacity: u64) -> StorageResult<Self> {
        // Ensure data directory exists
        fs::create_dir_all(&data_dir).await?;

        let store = Self {
            data_dir,
            chunk_sizes: DashMap::new(),
            total_bytes: AtomicU64::new(0),
            max_capacity,
            chunk_count: AtomicU64::new(0),
        };

        // Scan existing chunks
        store.scan_existing_chunks().await?;

        Ok(store)
    }

    /// Scan existing chunks in the data directory
    async fn scan_existing_chunks(&self) -> StorageResult<()> {
        let mut entries = fs::read_dir(&self.data_dir).await?;
        let mut total_bytes = 0u64;
        let mut chunk_count = 0u64;

        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();
            if path.is_file() {
                if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
                    if let Ok(metadata) = entry.metadata().await {
                        let size = metadata.len();
                        self.chunk_sizes.insert(name.to_string(), size);
                        total_bytes += size;
                        chunk_count += 1;
                    }
                }
            }
        }

        self.total_bytes.store(total_bytes, Ordering::SeqCst);
        self.chunk_count.store(chunk_count, Ordering::SeqCst);

        debug!(
            total_bytes,
            chunk_count,
            "Scanned existing chunks"
        );

        Ok(())
    }

    /// Get the path for a chunk
    fn chunk_path(&self, chunk_id: &str) -> PathBuf {
        // Sanitize chunk_id to be a valid filename
        let safe_id = chunk_id.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "_");
        self.data_dir.join(safe_id)
    }

    /// Store a chunk
    pub async fn put(&self, chunk_id: &str, data: &[u8]) -> StorageResult<u64> {
        let size = data.len() as u64;

        // Check capacity
        if self.max_capacity > 0 {
            let current = self.total_bytes.load(Ordering::SeqCst);
            if current + size > self.max_capacity {
                return Err(StorageError::CapacityExceeded);
            }
        }

        let path = self.chunk_path(chunk_id);

        // Check if replacing existing chunk
        let old_size = self.chunk_sizes.get(chunk_id).map(|v| *v).unwrap_or(0);

        // Write data
        let mut file = fs::File::create(&path).await?;
        file.write_all(data).await?;
        file.sync_all().await?;

        // Update index
        self.chunk_sizes.insert(chunk_id.to_string(), size);

        // Update totals
        if old_size > 0 {
            // Replacing existing chunk
            self.total_bytes.fetch_sub(old_size, Ordering::SeqCst);
        } else {
            // New chunk
            self.chunk_count.fetch_add(1, Ordering::SeqCst);
        }
        self.total_bytes.fetch_add(size, Ordering::SeqCst);

        debug!(chunk_id, size, "Stored chunk");

        Ok(size)
    }

    /// Retrieve a chunk
    pub async fn get(&self, chunk_id: &str) -> StorageResult<Vec<u8>> {
        let path = self.chunk_path(chunk_id);

        if !path.exists() {
            return Err(StorageError::NotFound(chunk_id.to_string()));
        }

        let mut file = fs::File::open(&path).await?;
        let mut data = Vec::new();
        file.read_to_end(&mut data).await?;

        debug!(chunk_id, size = data.len(), "Retrieved chunk");

        Ok(data)
    }

    /// Delete a chunk
    pub async fn delete(&self, chunk_id: &str) -> StorageResult<()> {
        let path = self.chunk_path(chunk_id);

        if let Some((_, size)) = self.chunk_sizes.remove(chunk_id) {
            if path.exists() {
                fs::remove_file(&path).await?;
            }
            self.total_bytes.fetch_sub(size, Ordering::SeqCst);
            self.chunk_count.fetch_sub(1, Ordering::SeqCst);
            debug!(chunk_id, "Deleted chunk");
        }

        Ok(())
    }

    /// Check if a chunk exists
    pub fn exists(&self, chunk_id: &str) -> bool {
        self.chunk_sizes.contains_key(chunk_id)
    }

    /// Get the size of a chunk
    pub fn size(&self, chunk_id: &str) -> Option<u64> {
        self.chunk_sizes.get(chunk_id).map(|v| *v)
    }

    /// Get total bytes stored
    pub fn total_bytes(&self) -> u64 {
        self.total_bytes.load(Ordering::SeqCst)
    }

    /// Get chunk count
    pub fn chunk_count(&self) -> u64 {
        self.chunk_count.load(Ordering::SeqCst)
    }

    /// Get maximum capacity
    pub fn max_capacity(&self) -> u64 {
        self.max_capacity
    }

    /// Get available capacity
    pub fn available_bytes(&self) -> u64 {
        if self.max_capacity == 0 {
            u64::MAX
        } else {
            self.max_capacity.saturating_sub(self.total_bytes())
        }
    }
}

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

    async fn create_test_store() -> (LocalChunkStore, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let store = LocalChunkStore::new(temp_dir.path().to_path_buf(), 0)
            .await
            .unwrap();
        (store, temp_dir)
    }

    #[tokio::test]
    async fn test_put_get() {
        let (store, _temp) = create_test_store().await;

        let chunk_id = "test-chunk-1";
        let data = vec![42u8; 1024];

        let size = store.put(chunk_id, &data).await.unwrap();
        assert_eq!(size, 1024);

        let retrieved = store.get(chunk_id).await.unwrap();
        assert_eq!(retrieved, data);
    }

    #[tokio::test]
    async fn test_delete() {
        let (store, _temp) = create_test_store().await;

        let chunk_id = "test-chunk-2";
        let data = vec![42u8; 512];

        store.put(chunk_id, &data).await.unwrap();
        assert!(store.exists(chunk_id));

        store.delete(chunk_id).await.unwrap();
        assert!(!store.exists(chunk_id));

        let result = store.get(chunk_id).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_size_tracking() {
        let (store, _temp) = create_test_store().await;

        assert_eq!(store.total_bytes(), 0);
        assert_eq!(store.chunk_count(), 0);

        store.put("chunk1", &vec![0u8; 100]).await.unwrap();
        assert_eq!(store.total_bytes(), 100);
        assert_eq!(store.chunk_count(), 1);

        store.put("chunk2", &vec![0u8; 200]).await.unwrap();
        assert_eq!(store.total_bytes(), 300);
        assert_eq!(store.chunk_count(), 2);

        store.delete("chunk1").await.unwrap();
        assert_eq!(store.total_bytes(), 200);
        assert_eq!(store.chunk_count(), 1);
    }

    #[tokio::test]
    async fn test_capacity_limit() {
        let temp_dir = TempDir::new().unwrap();
        let store = LocalChunkStore::new(temp_dir.path().to_path_buf(), 1000)
            .await
            .unwrap();

        // Should succeed
        store.put("chunk1", &vec![0u8; 500]).await.unwrap();

        // Should fail - would exceed capacity
        let result = store.put("chunk2", &vec![0u8; 600]).await;
        assert!(matches!(result, Err(StorageError::CapacityExceeded)));

        // Should succeed - within remaining capacity
        store.put("chunk2", &vec![0u8; 400]).await.unwrap();
    }

    #[tokio::test]
    async fn test_replace_chunk() {
        let (store, _temp) = create_test_store().await;

        let chunk_id = "test-chunk";

        store.put(chunk_id, &vec![0u8; 100]).await.unwrap();
        assert_eq!(store.total_bytes(), 100);
        assert_eq!(store.chunk_count(), 1);

        // Replace with larger data
        store.put(chunk_id, &vec![0u8; 200]).await.unwrap();
        assert_eq!(store.total_bytes(), 200);
        assert_eq!(store.chunk_count(), 1); // Still 1 chunk

        // Replace with smaller data
        store.put(chunk_id, &vec![0u8; 50]).await.unwrap();
        assert_eq!(store.total_bytes(), 50);
        assert_eq!(store.chunk_count(), 1);
    }
}