//! Memory-mapped storage for vectors and index data
//!
//! Provides efficient persistent storage using memory-mapped files,
//! allowing the OS to handle caching and paging.

use std::fs::{self, File, OpenOptions};
use std::io::{self, Read, Write, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::collections::HashMap;

use memmap2::{MmapMut, MmapOptions};
use parking_lot::RwLock;

use crate::types::StorageConfig;

/// File header for vector storage
/// Note: This struct uses manual serialization via to_bytes/from_bytes
/// because serde doesn't support arrays larger than 32 elements by default.
#[derive(Debug, Clone)]
struct VectorFileHeader {
    /// Magic number for file validation
    magic: [u8; 4],
    /// Version number
    version: u32,
    /// Vector dimension
    dimension: usize,
    /// Number of vectors stored
    count: usize,
    /// Reserved for future use
    #[allow(dead_code)]
    reserved: [u8; 48],
}

impl VectorFileHeader {
    const MAGIC: [u8; 4] = *b"AFDB";
    const VERSION: u32 = 1;
    const SIZE: usize = 64; // Fixed header size
    
    fn new(dimension: usize) -> Self {
        Self {
            magic: Self::MAGIC,
            version: Self::VERSION,
            dimension,
            count: 0,
            reserved: [0; 48],
        }
    }
    
    fn validate(&self) -> io::Result<()> {
        if self.magic != Self::MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid magic number"));
        }
        if self.version != Self::VERSION {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "Unsupported version"));
        }
        Ok(())
    }
    
    fn to_bytes(&self) -> [u8; Self::SIZE] {
        let mut buf = [0u8; Self::SIZE];
        buf[0..4].copy_from_slice(&self.magic);
        buf[4..8].copy_from_slice(&self.version.to_le_bytes());
        buf[8..16].copy_from_slice(&self.dimension.to_le_bytes());
        buf[16..24].copy_from_slice(&self.count.to_le_bytes());
        buf
    }
    
    fn from_bytes(buf: &[u8; Self::SIZE]) -> Self {
        let mut magic = [0u8; 4];
        magic.copy_from_slice(&buf[0..4]);
        
        Self {
            magic,
            version: u32::from_le_bytes(
                buf[4..8].try_into().expect("slice length guaranteed by type")
            ),
            dimension: usize::from_le_bytes(
                buf[8..16].try_into().expect("slice length guaranteed by type")
            ),
            count: usize::from_le_bytes(
                buf[16..24].try_into().expect("slice length guaranteed by type")
            ),
            reserved: [0; 48],
        }
    }
}

/// Memory-mapped storage for vector data
pub struct MmapStorage {
    /// Configuration
    config: StorageConfig,
    
    /// Path to vectors file
    vectors_path: PathBuf,
    
    /// Path to index file (reserved for future HNSW index persistence)
    #[allow(dead_code)]
    index_path: PathBuf,
    
    /// Path to metadata file
    meta_path: PathBuf,
    
    /// Memory-mapped vector data
    vector_mmap: RwLock<Option<MmapMut>>,
    
    /// Vector file
    vector_file: RwLock<Option<File>>,
    
    /// Vector dimension
    dimension: usize,
    
    /// Current vector count
    count: RwLock<usize>,
    
    /// ID to offset mapping
    id_to_offset: RwLock<HashMap<String, usize>>,
}

impl MmapStorage {
    /// Create a new mmap storage with the given configuration
    pub fn new(config: StorageConfig, dimension: usize) -> io::Result<Self> {
        let base_path = Path::new(&config.path);
        fs::create_dir_all(base_path)?;
        
        let vectors_path = base_path.join("vectors.aifastdb");
        let index_path = base_path.join("index.aifastdb");
        let meta_path = base_path.join("meta.aifastdb");
        
        let storage = Self {
            config,
            vectors_path,
            index_path,
            meta_path,
            vector_mmap: RwLock::new(None),
            vector_file: RwLock::new(None),
            dimension,
            count: RwLock::new(0),
            id_to_offset: RwLock::new(HashMap::new()),
        };
        
        Ok(storage)
    }
    
    /// Open an existing storage or create a new one
    pub fn open_or_create(config: StorageConfig, dimension: usize) -> io::Result<Self> {
        let mut storage = Self::new(config, dimension)?;
        
        if storage.vectors_path.exists() {
            storage.load()?;
        } else {
            storage.initialize()?;
        }
        
        Ok(storage)
    }
    
    /// Initialize a new storage
    fn initialize(&mut self) -> io::Result<()> {
        // Create vectors file with header
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(&self.vectors_path)?;
        
        let header = VectorFileHeader::new(self.dimension);
        file.write_all(&header.to_bytes())?;
        file.flush()?;
        
        *self.vector_file.write() = Some(file);
        *self.count.write() = 0;
        
        Ok(())
    }
    
    /// Load existing storage
    fn load(&mut self) -> io::Result<()> {
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&self.vectors_path)?;
        
        // Read and validate header
        let mut header_buf = [0u8; VectorFileHeader::SIZE];
        file.read_exact(&mut header_buf)?;
        
        let header = VectorFileHeader::from_bytes(&header_buf);
        header.validate()?;
        
        if header.dimension != self.dimension {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Dimension mismatch: expected {}, got {}", self.dimension, header.dimension)
            ));
        }
        
        *self.count.write() = header.count;
        *self.vector_file.write() = Some(file);
        
        // Load metadata (ID mappings)
        self.load_metadata()?;
        
        // Create memory map if file is large enough
        if header.count > 0 {
            self.create_mmap()?;
        }
        
        Ok(())
    }
    
    /// Create memory map for the vector file
    fn create_mmap(&self) -> io::Result<()> {
        let file = self.vector_file.read();
        if let Some(ref f) = *file {
            let file_len = f.metadata()?.len();
            if file_len > VectorFileHeader::SIZE as u64 {
                let mmap = unsafe {
                    MmapOptions::new()
                        .offset(VectorFileHeader::SIZE as u64)
                        .len((file_len as usize) - VectorFileHeader::SIZE)
                        .map_mut(f)?
                };
                *self.vector_mmap.write() = Some(mmap);
            }
        }
        Ok(())
    }
    
    /// Load metadata from file
    fn load_metadata(&self) -> io::Result<()> {
        if !self.meta_path.exists() {
            return Ok(());
        }
        
        let data = fs::read_to_string(&self.meta_path)?;
        let mappings: HashMap<String, usize> = serde_json::from_str(&data)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        
        *self.id_to_offset.write() = mappings;
        Ok(())
    }
    
    /// Save metadata to file
    fn save_metadata(&self) -> io::Result<()> {
        let mappings = self.id_to_offset.read();
        let data = serde_json::to_string(&*mappings)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        
        fs::write(&self.meta_path, data)?;
        Ok(())
    }
    
    /// Get vector dimension
    pub fn dimension(&self) -> usize {
        self.dimension
    }
    
    /// Get number of stored vectors
    pub fn len(&self) -> usize {
        *self.count.read()
    }
    
    /// Check if storage is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    
    /// Calculate byte size for a vector
    fn vector_byte_size(&self) -> usize {
        self.dimension * std::mem::size_of::<f32>()
    }
    
    /// Store a vector with an ID
    pub fn store(&self, id: impl Into<String>, vector: &[f32]) -> io::Result<usize> {
        let id = id.into();
        
        if vector.len() != self.dimension {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("Vector dimension mismatch: expected {}, got {}", self.dimension, vector.len())
            ));
        }
        
        // Check for duplicate
        {
            let mappings = self.id_to_offset.read();
            if mappings.contains_key(&id) {
                return Err(io::Error::new(io::ErrorKind::AlreadyExists, "ID already exists"));
            }
        }
        
        let offset = {
            let mut count = self.count.write();
            let current_offset = *count;
            *count += 1;
            current_offset
        };
        
        // Write vector to file
        {
            let mut file_guard = self.vector_file.write();
            if let Some(ref mut file) = *file_guard {
                let byte_offset = VectorFileHeader::SIZE + offset * self.vector_byte_size();
                file.seek(SeekFrom::Start(byte_offset as u64))?;
                
                // Write f32 values as bytes
                for &val in vector {
                    file.write_all(&val.to_le_bytes())?;
                }
                
                if self.config.sync_on_write {
                    file.flush()?;
                }
            }
        }
        
        // Update header
        self.update_header()?;
        
        // Update ID mapping
        {
            let mut mappings = self.id_to_offset.write();
            mappings.insert(id, offset);
        }
        
        Ok(offset)
    }
    
    /// Update file header with current count
    fn update_header(&self) -> io::Result<()> {
        let mut file_guard = self.vector_file.write();
        if let Some(ref mut file) = *file_guard {
            file.seek(SeekFrom::Start(0))?;
            
            let count = *self.count.read();
            let header = VectorFileHeader {
                magic: VectorFileHeader::MAGIC,
                version: VectorFileHeader::VERSION,
                dimension: self.dimension,
                count,
                reserved: [0; 48],
            };
            
            file.write_all(&header.to_bytes())?;
            file.flush()?;
        }
        Ok(())
    }
    
    /// Retrieve a vector by offset
    pub fn get_by_offset(&self, offset: usize) -> io::Result<Vec<f32>> {
        let count = *self.count.read();
        if offset >= count {
            return Err(io::Error::new(io::ErrorKind::NotFound, "Offset out of bounds"));
        }
        
        let byte_offset = offset * self.vector_byte_size();
        
        // Try mmap first
        {
            let mmap_guard = self.vector_mmap.read();
            if let Some(ref mmap) = *mmap_guard {
                let start = byte_offset;
                let end = start + self.vector_byte_size();
                
                if end <= mmap.len() {
                    let bytes = &mmap[start..end];
                    return Ok(bytes_to_f32_vec(bytes));
                }
            }
        }
        
        // Fall back to file read
        let mut file_guard = self.vector_file.write();
        if let Some(ref mut file) = *file_guard {
            let file_offset = VectorFileHeader::SIZE + byte_offset;
            file.seek(SeekFrom::Start(file_offset as u64))?;
            
            let mut vec = Vec::with_capacity(self.dimension);
            for _ in 0..self.dimension {
                let mut buf = [0u8; 4];
                file.read_exact(&mut buf)?;
                vec.push(f32::from_le_bytes(buf));
            }
            
            return Ok(vec);
        }
        
        Err(io::Error::new(io::ErrorKind::NotFound, "Storage not initialized"))
    }
    
    /// Retrieve a vector by ID
    pub fn get(&self, id: &str) -> io::Result<Vec<f32>> {
        let offset = {
            let mappings = self.id_to_offset.read();
            *mappings.get(id).ok_or_else(|| {
                io::Error::new(io::ErrorKind::NotFound, format!("ID not found: {}", id))
            })?
        };
        
        self.get_by_offset(offset)
    }
    
    /// Check if an ID exists
    pub fn contains(&self, id: &str) -> bool {
        self.id_to_offset.read().contains_key(id)
    }
    
    /// Get all IDs
    pub fn ids(&self) -> Vec<String> {
        self.id_to_offset.read().keys().cloned().collect()
    }
    
    /// Flush all pending writes to disk
    pub fn flush(&self) -> io::Result<()> {
        // Flush vector file
        {
            let mut file_guard = self.vector_file.write();
            if let Some(ref mut file) = *file_guard {
                file.flush()?;
            }
        }
        
        // Flush mmap
        {
            let mmap_guard = self.vector_mmap.read();
            if let Some(ref mmap) = *mmap_guard {
                mmap.flush()?;
            }
        }
        
        // Save metadata
        self.save_metadata()?;
        
        Ok(())
    }
    
    /// Close the storage
    pub fn close(&self) -> io::Result<()> {
        self.flush()?;
        
        *self.vector_mmap.write() = None;
        *self.vector_file.write() = None;
        
        Ok(())
    }
}

impl Drop for MmapStorage {
    fn drop(&mut self) {
        let _ = self.close();
    }
}

/// Convert bytes to f32 vector
fn bytes_to_f32_vec(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(4)
        .map(|chunk| {
            // chunks_exact(4) guarantees exactly 4 bytes per chunk
            let arr: [u8; 4] = chunk.try_into().expect("chunks_exact guarantees 4 bytes");
            f32::from_le_bytes(arr)
        })
        .collect()
}

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

    #[test]
    fn test_mmap_storage_basic() {
        let dir = tempdir().unwrap();
        let config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            use_mmap: true,
            cache_size: 1024 * 1024,
            sync_on_write: true,
            ..Default::default()
        };
        
        let storage = MmapStorage::open_or_create(config, 4).unwrap();
        
        // Store vectors
        storage.store("v1", &[1.0, 2.0, 3.0, 4.0]).unwrap();
        storage.store("v2", &[5.0, 6.0, 7.0, 8.0]).unwrap();
        
        assert_eq!(storage.len(), 2);
        
        // Retrieve vectors
        let v1 = storage.get("v1").unwrap();
        assert_eq!(v1, vec![1.0, 2.0, 3.0, 4.0]);
        
        let v2 = storage.get("v2").unwrap();
        assert_eq!(v2, vec![5.0, 6.0, 7.0, 8.0]);
        
        // Check contains
        assert!(storage.contains("v1"));
        assert!(!storage.contains("v3"));
    }

    #[test]
    fn test_mmap_storage_persistence() {
        let dir = tempdir().unwrap();
        let config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            use_mmap: true,
            cache_size: 1024 * 1024,
            sync_on_write: true,
            ..Default::default()
        };
        
        // Create and store
        {
            let storage = MmapStorage::open_or_create(config.clone(), 4).unwrap();
            storage.store("test", &[1.0, 2.0, 3.0, 4.0]).unwrap();
            storage.flush().unwrap();
        }
        
        // Reopen and verify
        {
            let storage = MmapStorage::open_or_create(config, 4).unwrap();
            assert_eq!(storage.len(), 1);
            
            let v = storage.get("test").unwrap();
            assert_eq!(v, vec![1.0, 2.0, 3.0, 4.0]);
        }
    }

    #[test]
    fn test_dimension_mismatch() {
        let dir = tempdir().unwrap();
        let config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            use_mmap: true,
            cache_size: 1024 * 1024,
            sync_on_write: false,
            ..Default::default()
        };
        
        let storage = MmapStorage::open_or_create(config, 4).unwrap();
        
        // Wrong dimension should fail
        let result = storage.store("bad", &[1.0, 2.0]);
        assert!(result.is_err());
    }

    #[test]
    fn test_duplicate_id() {
        let dir = tempdir().unwrap();
        let config = StorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            use_mmap: true,
            cache_size: 1024 * 1024,
            sync_on_write: false,
            ..Default::default()
        };
        
        let storage = MmapStorage::open_or_create(config, 4).unwrap();
        
        storage.store("dup", &[1.0, 2.0, 3.0, 4.0]).unwrap();
        let result = storage.store("dup", &[5.0, 6.0, 7.0, 8.0]);
        assert!(result.is_err());
    }
}
