//! Main PersistentVectorIndex implementation.

use super::config::PersistentVectorIndexConfig;
use crate::checkpoint_format::{read_checkpoint, write_checkpoint};
use crate::error::{Result, StorageError};
use crate::vector_index::persistence::format::{
    IdMappingEntry, IdMappingTable, SnapshotMetadata, DATA_MAGIC, FORMAT_VERSION, GRAPH_MAGIC,
};
use crate::vector_index::{merge_search_results, HnswVectorIndex, VectorIndex};
use parking_lot::RwLock;
use std::collections::HashMap;
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use stemedb_core::types::Hash;
use tracing::{debug, info, instrument, warn};

/// Cold index loaded from disk.
struct ColdIndex {
    /// Metadata about this snapshot.
    #[allow(dead_code)]
    metadata: SnapshotMetadata,
    /// In-memory HNSW (loaded from disk, not mmap'd for simplicity).
    hnsw: HnswVectorIndex,
    /// Timestamp cutoff - vectors before this time are in cold.
    #[allow(dead_code)]
    cutoff_timestamp: u64,
}

impl ColdIndex {
    /// Load cold index from disk.
    #[instrument(skip_all, fields(path = %data_dir.display()))]
    fn load(data_dir: &Path, version: u64, dimension: usize) -> Result<Self> {
        let graph_path = data_dir.join(format!("{}.hnsw.graph", version));
        let data_path = data_dir.join(format!("{}.hnsw.data", version));

        if !graph_path.exists() || !data_path.exists() {
            return Err(StorageError::NotFound(format!(
                "Cold index version {} not found",
                version
            )));
        }

        // Load metadata from graph file
        let mut graph_file = File::open(&graph_path)?;
        let meta_bytes = read_checkpoint(&mut graph_file, GRAPH_MAGIC, FORMAT_VERSION)?;
        let metadata: SnapshotMetadata = crate::serde_helpers::deserialize(&meta_bytes)?;

        if metadata.dimension != dimension {
            return Err(StorageError::InputValidation(format!(
                "Dimension mismatch: expected {}, got {}",
                dimension, metadata.dimension
            )));
        }

        // Load ID mapping from data file
        let mut data_file = File::open(&data_path)?;
        let data_bytes = read_checkpoint(&mut data_file, DATA_MAGIC, FORMAT_VERSION)?;
        let _mapping: IdMappingTable = crate::serde_helpers::deserialize(&data_bytes)?;

        // LIMITATION: Current implementation only persists ID mappings, not actual vectors.
        //
        // The cold index is created empty here. After restart, vector search on
        // cold data will return no results until vectors are re-ingested via WAL
        // replay. This is intentional for Phase 5C - full vector persistence with
        // mmap'd HNSW graphs is planned for a future phase.
        //
        // The mapping table tracks which hashes exist in cold storage, enabling
        // deduplication and contains() checks, but search requires re-ingestion.
        let hnsw = HnswVectorIndex::new(dimension);

        info!(
            version = metadata.version,
            vector_count = metadata.vector_count,
            dimension = metadata.dimension,
            "Loaded cold index metadata"
        );

        Ok(Self { cutoff_timestamp: metadata.created_at, metadata, hnsw })
    }
}

/// Persistent vector index with hot/cold architecture.
///
/// The hot index is in-memory and handles recent vectors. The cold index
/// is loaded from disk and stores historical data. Background checkpointing
/// periodically persists hot data to disk.
pub struct PersistentVectorIndex {
    /// Configuration.
    config: PersistentVectorIndexConfig,

    /// Hot index (in-memory, recent vectors).
    hot: RwLock<HnswVectorIndex>,

    /// Cold index (disk-backed, historical vectors).
    cold: RwLock<Option<ColdIndex>>,

    /// Current cold version number.
    cold_version: AtomicU64,

    /// Maps hash to insertion timestamp for hot/cold partitioning.
    hash_timestamps: RwLock<HashMap<Hash, u64>>,

    /// Shutdown flag.
    shutdown: AtomicBool,

    /// Number of inserts since last checkpoint.
    inserts_since_checkpoint: AtomicU64,
}

impl PersistentVectorIndex {
    /// Open a persistent vector index.
    ///
    /// Loads cold index from disk if it exists, otherwise starts fresh.
    #[instrument(skip(config), fields(data_dir = %config.data_dir.display(), dimension = config.dimension))]
    pub fn open(config: PersistentVectorIndexConfig) -> Result<Self> {
        // Create data directory if needed
        fs::create_dir_all(&config.data_dir)?;

        // Try to load cold index
        let (cold, cold_version) = match Self::find_latest_version(&config.data_dir) {
            Some(version) => match ColdIndex::load(&config.data_dir, version, config.dimension) {
                Ok(index) => {
                    info!(version, "Loaded cold index");
                    (Some(index), version)
                }
                Err(e) => {
                    warn!(error = %e, version, "Failed to load cold index, starting fresh");
                    (None, 0)
                }
            },
            None => {
                info!("No cold index found, starting fresh");
                (None, 0)
            }
        };

        let hot = HnswVectorIndex::with_params(
            config.dimension,
            config.max_nb_connection,
            config.ef_construction,
        );

        Ok(Self {
            config,
            hot: RwLock::new(hot),
            cold: RwLock::new(cold),
            cold_version: AtomicU64::new(cold_version),
            hash_timestamps: RwLock::new(HashMap::new()),
            shutdown: AtomicBool::new(false),
            inserts_since_checkpoint: AtomicU64::new(0),
        })
    }

    /// Create a new persistent vector index with default config.
    pub fn new(data_dir: PathBuf, dimension: usize) -> Result<Self> {
        let config = PersistentVectorIndexConfig { data_dir, dimension, ..Default::default() };
        Self::open(config)
    }

    /// Find the latest snapshot version in the data directory.
    fn find_latest_version(data_dir: &Path) -> Option<u64> {
        let mut max_version = None;

        if let Ok(entries) = fs::read_dir(data_dir) {
            for entry in entries.flatten() {
                let name = entry.file_name();
                let name_str = name.to_string_lossy();
                if name_str.ends_with(".hnsw.graph") {
                    if let Some(version_str) = name_str.strip_suffix(".hnsw.graph") {
                        if let Ok(version) = version_str.parse::<u64>() {
                            max_version =
                                Some(max_version.map_or(version, |m: u64| m.max(version)));
                        }
                    }
                }
            }
        }

        max_version
    }

    /// Write a checkpoint of the current hot index to disk.
    #[instrument(skip(self))]
    pub fn checkpoint(&self) -> Result<()> {
        let inserts = self.inserts_since_checkpoint.load(Ordering::Relaxed);
        if inserts == 0 {
            debug!("No changes since last checkpoint, skipping");
            return Ok(());
        }

        let new_version = self.cold_version.load(Ordering::Relaxed) + 1;

        // Build ID mapping from hot index
        let mapping = {
            let timestamps = self.hash_timestamps.read();
            let hot = self.hot.read();

            let mut table = IdMappingTable::new();
            // We need to iterate through all hashes in hot index
            // For now, we use the timestamps map as our source of truth
            for (hash, &timestamp) in timestamps.iter() {
                if hot.contains(hash) {
                    table.push(IdMappingEntry { internal_id: table.len(), hash: *hash, timestamp });
                }
            }
            table
        };

        // Serialize mapping
        let mapping_bytes = crate::serde_helpers::serialize(&mapping)?;
        let data_hash = blake3::hash(&mapping_bytes);

        // For now, graph is empty - in full impl, we'd serialize HNSW graph
        let graph_bytes: Vec<u8> = Vec::new();
        let graph_hash = blake3::hash(&graph_bytes);

        // Create metadata
        let hot = self.hot.read();
        let metadata = SnapshotMetadata::new(
            new_version,
            0, // WAL offset start - not used yet
            0, // WAL offset end - not used yet
            hot.len() as u64,
            self.config.dimension,
            *graph_hash.as_bytes(),
            *data_hash.as_bytes(),
        );
        drop(hot);

        // Write files atomically
        self.write_snapshot_files(new_version, &metadata, &graph_bytes, &mapping_bytes)?;

        // Update cold version
        self.cold_version.store(new_version, Ordering::Release);
        self.inserts_since_checkpoint.store(0, Ordering::Relaxed);

        info!(
            version = new_version,
            vector_count = metadata.vector_count,
            "Wrote vector index checkpoint"
        );

        Ok(())
    }

    fn write_snapshot_files(
        &self,
        version: u64,
        metadata: &SnapshotMetadata,
        _graph_bytes: &[u8],
        data_bytes: &[u8],
    ) -> Result<()> {
        let graph_path = self.config.data_dir.join(format!("{}.hnsw.graph", version));
        let data_path = self.config.data_dir.join(format!("{}.hnsw.data", version));
        let graph_tmp = self.config.data_dir.join(format!("{}.hnsw.graph.tmp", version));
        let data_tmp = self.config.data_dir.join(format!("{}.hnsw.data.tmp", version));

        // Write graph file (contains metadata only for now)
        {
            let mut file = File::create(&graph_tmp)?;
            let meta_bytes = crate::serde_helpers::serialize(metadata)?;
            write_checkpoint(&mut file, GRAPH_MAGIC, FORMAT_VERSION, &meta_bytes)?;
            file.sync_all()?;
        }

        // Write data file
        {
            let mut file = File::create(&data_tmp)?;
            write_checkpoint(&mut file, DATA_MAGIC, FORMAT_VERSION, data_bytes)?;
            file.sync_all()?;
        }

        // Atomic renames (clean up temp files on failure)
        if let Err(e) = fs::rename(&graph_tmp, &graph_path) {
            let _ = fs::remove_file(&graph_tmp);
            let _ = fs::remove_file(&data_tmp);
            return Err(e.into());
        }
        if let Err(e) = fs::rename(&data_tmp, &data_path) {
            // Graph was renamed successfully, but data failed - clean up data temp
            let _ = fs::remove_file(&data_tmp);
            // Note: graph_path now exists but data_path doesn't - partial state
            // On next load, this will fail validation and fall back to hot-only
            return Err(e.into());
        }

        // Sync parent directory
        if let Ok(dir) = File::open(&self.config.data_dir) {
            let _ = dir.sync_all();
        }

        Ok(())
    }

    /// Check if a hash exists in the index.
    pub fn contains(&self, hash: &Hash) -> bool {
        if self.hot.read().contains(hash) {
            return true;
        }
        if let Some(ref cold) = *self.cold.read() {
            return cold.hnsw.contains(hash);
        }
        false
    }

    /// Signal shutdown and perform final checkpoint.
    pub fn shutdown(&self) -> Result<()> {
        self.shutdown.store(true, Ordering::SeqCst);
        self.checkpoint()
    }

    /// Start background checkpoint task.
    pub fn start_background_checkpoint(
        self: &Arc<Self>,
        interval_secs: u64,
    ) -> tokio::task::JoinHandle<()> {
        let this = Arc::clone(self);
        tokio::spawn(async move {
            let interval = tokio::time::Duration::from_secs(interval_secs);
            loop {
                tokio::time::sleep(interval).await;

                if this.shutdown.load(Ordering::SeqCst) {
                    info!("Vector index checkpoint task shutting down");
                    break;
                }

                if let Err(e) = this.checkpoint() {
                    warn!(error = %e, "Background checkpoint failed");
                }
            }
        })
    }
}

impl VectorIndex for PersistentVectorIndex {
    #[instrument(skip(self, vector), fields(hash = %hex::encode(hash), dim = vector.len()))]
    fn insert(&self, hash: &Hash, vector: &[f32]) -> Result<()> {
        // Insert into hot index
        self.hot.read().insert(hash, vector)?;

        // Track timestamp
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs());
        self.hash_timestamps.write().insert(*hash, now);

        self.inserts_since_checkpoint.fetch_add(1, Ordering::Relaxed);

        Ok(())
    }

    #[instrument(skip(self, query), fields(dim = query.len(), k = k))]
    fn search(&self, query: &[f32], k: usize) -> Result<Vec<(Hash, f32)>> {
        // Search hot index
        let hot_results = self.hot.read().search(query, k)?;

        // Search cold index if present
        let cold_results = if let Some(ref cold) = *self.cold.read() {
            cold.hnsw.search(query, k)?
        } else {
            Vec::new()
        };

        // Merge results
        let merged = merge_search_results(hot_results, cold_results, k);

        debug!(results_count = merged.len(), "Persistent vector search complete");

        Ok(merged)
    }

    fn dimension(&self) -> usize {
        self.config.dimension
    }

    fn len(&self) -> usize {
        let hot_len = self.hot.read().len();
        let cold_len = self.cold.read().as_ref().map_or(0, |c| c.hnsw.len());
        hot_len + cold_len
    }
}

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

    fn random_vector(dim: usize, seed: u64) -> Vec<f32> {
        let mut state = seed;
        (0..dim)
            .map(|_| {
                state = state.wrapping_mul(1103515245).wrapping_add(12345);
                ((state >> 16) as f32 / 32768.0) - 1.0
            })
            .collect()
    }

    #[test]
    fn test_create_persistent_index() {
        let tmp = TempDir::new().expect("create temp dir");
        let index =
            PersistentVectorIndex::new(tmp.path().to_path_buf(), 128).expect("create index");

        assert_eq!(index.dimension(), 128);
        assert!(index.is_empty());
    }

    #[test]
    fn test_insert_and_search() {
        let tmp = TempDir::new().expect("create temp dir");
        let index = PersistentVectorIndex::new(tmp.path().to_path_buf(), 4).expect("create index");

        let hash1: Hash = [1u8; 32];
        let hash2: Hash = [2u8; 32];

        index.insert(&hash1, &[1.0, 0.0, 0.0, 0.0]).expect("insert 1");
        index.insert(&hash2, &[0.9, 0.1, 0.0, 0.0]).expect("insert 2");

        assert_eq!(index.len(), 2);
        assert!(index.contains(&hash1));
        assert!(index.contains(&hash2));

        let results = index.search(&[1.0, 0.0, 0.0, 0.0], 2).expect("search");
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].0, hash1);
    }

    #[test]
    fn test_checkpoint_and_reload() {
        let tmp = TempDir::new().expect("create temp dir");
        let path = tmp.path().to_path_buf();

        // Create and populate index
        {
            let index = PersistentVectorIndex::new(path.clone(), 4).expect("create index");

            let hash1: Hash = [1u8; 32];
            index.insert(&hash1, &[1.0, 0.0, 0.0, 0.0]).expect("insert");

            index.checkpoint().expect("checkpoint");
        }

        // Reopen and verify
        {
            let config =
                PersistentVectorIndexConfig { data_dir: path, dimension: 4, ..Default::default() };
            let index = PersistentVectorIndex::open(config).expect("reopen index");

            // Cold index should have loaded
            assert!(index.cold.read().is_some());
        }
    }

    #[test]
    fn test_dimension_mismatch() {
        let tmp = TempDir::new().expect("create temp dir");
        let index = PersistentVectorIndex::new(tmp.path().to_path_buf(), 4).expect("create index");

        let hash: Hash = [1u8; 32];
        let result = index.insert(&hash, &[1.0, 0.0, 0.0]); // Wrong dimension
        assert!(result.is_err());
    }

    #[test]
    fn test_larger_scale() {
        let tmp = TempDir::new().expect("create temp dir");
        let dim = 128;
        let index =
            PersistentVectorIndex::new(tmp.path().to_path_buf(), dim).expect("create index");

        // Insert 100 vectors
        for i in 0..100u64 {
            let mut hash: Hash = [0u8; 32];
            hash[0..8].copy_from_slice(&i.to_le_bytes());
            let vector = random_vector(dim, i);
            index.insert(&hash, &vector).expect("insert");
        }

        assert_eq!(index.len(), 100);

        // Search should find nearest neighbors
        let query = random_vector(dim, 50);
        let results = index.search(&query, 5).expect("search");

        assert_eq!(results.len(), 5);
        // Exact match should be first
        let mut expected_hash: Hash = [0u8; 32];
        expected_hash[0..8].copy_from_slice(&50u64.to_le_bytes());
        assert_eq!(results[0].0, expected_hash);
    }
}
