//! Persistent visual index implementation.

use super::config::PersistentVisualIndexConfig;
use crate::checkpoint_format::{atomic_write, read_checkpoint, write_checkpoint};
use crate::error::Result;
use crate::visual_index::bk_tree::{BkTreeSnapshot, BkTreeVisualIndex};
use crate::visual_index::hamming::VisualIndex;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use stemedb_core::types::{Hash, PHash};
use tracing::{debug, info, instrument, warn};

/// Magic bytes for visual index checkpoint file.
const CHECKPOINT_MAGIC: &[u8; 4] = b"VHCK";

/// Current checkpoint format version.
const CHECKPOINT_VERSION: u8 = 1;

/// Persistent wrapper around BkTreeVisualIndex that checkpoints to disk.
///
/// Uses atomic writes (temp file + fsync + rename) to ensure crash safety.
/// CRC32C checksums detect corruption.
pub struct PersistentVisualIndex {
    /// The underlying in-memory index.
    inner: BkTreeVisualIndex,
    /// Path to the checkpoint file.
    checkpoint_path: PathBuf,
    /// Number of changes since last checkpoint.
    changes_since_checkpoint: AtomicU64,
    /// Minimum changes before checkpoint is worthwhile.
    min_changes_for_checkpoint: usize,
    /// Shutdown flag.
    shutdown: AtomicBool,
}

impl PersistentVisualIndex {
    /// Open a persistent visual index, loading from checkpoint if it exists.
    ///
    /// # Arguments
    /// * `config` - Configuration for the persistent index
    ///
    /// # Errors
    /// Returns error if checkpoint file is corrupt or cannot be read.
    #[instrument(skip(config), fields(checkpoint_path = %config.checkpoint_path.display()))]
    pub fn open(config: PersistentVisualIndexConfig) -> Result<Self> {
        let inner = if config.checkpoint_path.exists() {
            info!(path = %config.checkpoint_path.display(), "Loading visual index from checkpoint");
            Self::load_checkpoint(&config.checkpoint_path)?
        } else {
            info!(path = %config.checkpoint_path.display(), "No checkpoint found, starting fresh");
            BkTreeVisualIndex::new()
        };

        Ok(Self {
            inner,
            checkpoint_path: config.checkpoint_path,
            changes_since_checkpoint: AtomicU64::new(0),
            min_changes_for_checkpoint: config.min_changes_for_checkpoint,
            shutdown: AtomicBool::new(false),
        })
    }

    /// Create a new persistent visual index with default config.
    pub fn new(checkpoint_path: PathBuf) -> Self {
        Self {
            inner: BkTreeVisualIndex::new(),
            checkpoint_path,
            changes_since_checkpoint: AtomicU64::new(0),
            min_changes_for_checkpoint: 100,
            shutdown: AtomicBool::new(false),
        }
    }

    /// Load index from checkpoint file.
    ///
    /// File format:
    /// ```text
    /// [MAGIC:4 "VHCK"][VERSION:1][RESERVED:3][PAYLOAD_LEN:u64_LE][CRC32C:u32][PAYLOAD:N]
    /// ```
    fn load_checkpoint(path: &Path) -> Result<BkTreeVisualIndex> {
        let mut file = File::open(path)?;

        // Read and validate checkpoint using shared format
        let payload = read_checkpoint(&mut file, CHECKPOINT_MAGIC, CHECKPOINT_VERSION)?;

        // Deserialize snapshot and rebuild index
        let snapshot: BkTreeSnapshot = crate::serde_helpers::deserialize(&payload)?;
        let node_count = snapshot.node_count;
        let index = BkTreeVisualIndex::from_snapshot(snapshot)?;
        info!(
            path = %path.display(),
            node_count,
            "Loaded visual index from checkpoint"
        );

        Ok(index)
    }

    /// Write checkpoint to disk with atomic write pattern.
    ///
    /// 1. Serialize snapshot
    /// 2. Write to temp file with CRC32C
    /// 3. fsync temp file
    /// 4. Rename to final path (atomic on POSIX)
    /// 5. fsync parent directory
    #[instrument(skip(self))]
    pub fn checkpoint(&self) -> Result<()> {
        let changes = self.changes_since_checkpoint.load(Ordering::Relaxed);
        if changes == 0 {
            debug!("No changes since last checkpoint, skipping");
            return Ok(());
        }

        // Create snapshot
        let snapshot = self.inner.snapshot()?;

        // Serialize
        let payload = crate::serde_helpers::serialize(&snapshot)?;
        let node_count = snapshot.nodes.len();
        let payload_size = payload.len();

        // Atomic write using shared helper
        atomic_write(&self.checkpoint_path, |file| {
            write_checkpoint(file, CHECKPOINT_MAGIC, CHECKPOINT_VERSION, &payload)
        })?;

        // Reset change counter
        self.changes_since_checkpoint.store(0, Ordering::Relaxed);

        info!(
            path = %self.checkpoint_path.display(),
            node_count,
            payload_size,
            "Wrote visual index checkpoint"
        );

        Ok(())
    }

    /// Check if an assertion hash is already in the index.
    pub fn contains(&self, hash: &Hash) -> bool {
        self.inner.contains(hash)
    }

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

    /// Start background checkpoint task.
    ///
    /// Returns a handle that can be used to stop the 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!("Visual index checkpoint task shutting down");
                    break;
                }

                let changes = this.changes_since_checkpoint.load(Ordering::Relaxed);
                if changes >= this.min_changes_for_checkpoint as u64 {
                    if let Err(e) = this.checkpoint() {
                        warn!(error = %e, "Background checkpoint failed");
                    }
                }
            }
        })
    }
}

impl VisualIndex for PersistentVisualIndex {
    fn insert(&self, hash: &Hash, phash: &PHash) -> Result<()> {
        let was_new = !self.inner.contains(hash);
        self.inner.insert(hash, phash)?;
        if was_new {
            self.changes_since_checkpoint.fetch_add(1, Ordering::Relaxed);
        }
        Ok(())
    }

    fn search(&self, query: &PHash, threshold: u32) -> Result<Vec<(Hash, u32)>> {
        self.inner.search(query, threshold)
    }

    fn len(&self) -> Result<usize> {
        self.inner.len()
    }

    fn checkpoint(&self) -> Result<()> {
        PersistentVisualIndex::checkpoint(self)
    }
}
