//! Local Episteme instance for persistent storage and alias management.
//!
//! Provides ingestion, conflict checking, and auto-alias creation backed by
//! write-ahead log and KV store.

mod queries;
mod store;

use std::path::Path;
use std::sync::Arc;

use ed25519_dalek::SigningKey;
use stemedb_ingest::Ingestor;
use stemedb_storage::{
    GenericAliasStore, GenericPackSourceStore, GenericPredicateAliasStore,
    GenericPredicateIndexStore, HybridStore, KVStore, PredicateAliasStore, StoredPredicateAliasSet,
};
use stemedb_wal::Journal;
use tokio::sync::Mutex;
use tracing::{info, instrument};

use crate::bridge::load_or_generate_key;
use crate::config::AphoriaConfig;
use crate::types::PredicateAliasSet;
use crate::AphoriaError;

/// Local Episteme instance for Aphoria.
pub struct LocalEpisteme {
    pub(super) journal: Arc<Mutex<Journal>>,
    pub(super) store: Arc<HybridStore>, // KV store for assertions
    pub(super) ingestor: Ingestor<HybridStore>,
    pub(super) signing_key: SigningKey,
    pub(super) alias_store: GenericAliasStore<Arc<HybridStore>>,
    pub(super) predicate_index_store: GenericPredicateIndexStore<Arc<HybridStore>>,
    pub(super) pack_source_store: GenericPackSourceStore<Arc<HybridStore>>,
    /// Predicate alias store for persisting semantic predicate equivalences.
    pub(super) predicate_alias_store: GenericPredicateAliasStore<Arc<HybridStore>>,
    /// Predicate aliases from imported Trust Packs (loaded from storage on startup).
    pub(super) predicate_aliases: Vec<PredicateAliasSet>,
}

impl LocalEpisteme {
    /// Open or create a local Episteme instance.
    #[instrument(skip(config), fields(data_dir = %config.episteme.data_dir.display()))]
    pub async fn open(config: &AphoriaConfig, project_root: &Path) -> Result<Self, AphoriaError> {
        let data_dir = &config.episteme.data_dir;

        // Create directories if needed
        std::fs::create_dir_all(data_dir)?;

        // Canonicalize paths (required by fjall/lsm-tree)
        let data_dir = data_dir.canonicalize().map_err(|e| {
            AphoriaError::Storage(format!("Failed to canonicalize data_dir: {}", e))
        })?;

        let wal_dir = data_dir.join("wal");
        let store_dir = data_dir.join("store");
        std::fs::create_dir_all(&wal_dir)?;
        std::fs::create_dir_all(&store_dir)?;

        info!("Opening local Episteme at {}", data_dir.display());

        // Open WAL
        let journal = Arc::new(Mutex::new(Journal::open(&wal_dir).map_err(|e| {
            AphoriaError::Storage(format!("Failed to open WAL at {}: {e}", wal_dir.display()))
        })?));

        // Open store
        let store = Arc::new(HybridStore::open(&store_dir).map_err(|e| {
            AphoriaError::Storage(format!("Failed to open store at {}: {e}", store_dir.display()))
        })?);

        // Create ingestor
        let mut ingestor = Ingestor::new(journal.clone(), store.clone())
            .await
            .map_err(|e| AphoriaError::Storage(format!("Failed to create ingestor: {e}")))?;
        ingestor.start();

        // Load or generate signing key
        let signing_key = load_or_generate_key(project_root).map_err(|e| {
            AphoriaError::Storage(format!(
                "Failed to load/generate signing key at {}: {e}",
                project_root.display()
            ))
        })?;

        // Create alias store for auto-alias persistence
        let alias_store = GenericAliasStore::new(store.clone());

        // Create predicate index store for predicate-based queries
        let predicate_index_store = GenericPredicateIndexStore::new(store.clone());

        // Create pack source store for policy attribution
        let pack_source_store = GenericPackSourceStore::new(store.clone());

        // Create predicate alias store for semantic predicate matching
        let predicate_alias_store = GenericPredicateAliasStore::new(store.clone());

        // Load persisted predicate aliases from storage
        let stored_aliases = predicate_alias_store
            .list_all_predicate_aliases()
            .await
            .map_err(|e| AphoriaError::Storage(format!("Failed to load predicate aliases: {e}")))?;
        let predicate_aliases: Vec<PredicateAliasSet> = stored_aliases
            .into_iter()
            .map(|s| PredicateAliasSet::new(s.canonical, s.aliases))
            .collect();

        if !predicate_aliases.is_empty() {
            info!(count = predicate_aliases.len(), "Loaded predicate aliases from storage");
        }

        Ok(Self {
            journal,
            store,
            ingestor,
            signing_key,
            alias_store,
            predicate_index_store,
            pack_source_store,
            predicate_alias_store,
            predicate_aliases,
        })
    }

    /// Shut down the Episteme instance gracefully.
    pub async fn shutdown(&mut self) {
        info!("Shutting down local Episteme");
        self.ingestor.shutdown(std::time::Duration::from_secs(2)).await;

        // Flush the store to ensure all data is persisted to disk.
        // This is critical for pack_source data written during policy import.
        if let Err(e) = self.store.as_ref().flush().await {
            tracing::warn!(error = %e, "Failed to flush store during shutdown");
        }
    }

    /// Get the signing key's public key bytes for alias creation.
    pub fn agent_id(&self) -> [u8; 32] {
        self.signing_key.verifying_key().to_bytes()
    }

    /// Get a reference to the alias store for querying created aliases.
    #[allow(dead_code)]
    pub fn alias_store(&self) -> &GenericAliasStore<Arc<HybridStore>> {
        &self.alias_store
    }

    /// Get a reference to the underlying KV store.
    ///
    /// Used for direct storage operations like importing policies.
    pub fn store(&self) -> &Arc<HybridStore> {
        &self.store
    }

    /// Get a reference to the pack source store for policy attribution.
    pub fn pack_source_store(&self) -> &GenericPackSourceStore<Arc<HybridStore>> {
        &self.pack_source_store
    }

    /// Get the current predicate aliases from imported Trust Packs.
    pub fn predicate_aliases(&self) -> &[PredicateAliasSet] {
        &self.predicate_aliases
    }

    /// Persist predicate aliases to storage and update in-memory cache.
    ///
    /// This is called during policy import to ensure aliases survive restarts.
    /// Uses merge semantics: if aliases for the same canonical predicate already
    /// exist, the new aliases are added to the existing set.
    ///
    /// # Arguments
    /// * `aliases` - The predicate alias sets to persist
    pub async fn persist_predicate_aliases(
        &mut self,
        aliases: Vec<PredicateAliasSet>,
    ) -> Result<(), AphoriaError> {
        for alias in &aliases {
            let stored = StoredPredicateAliasSet {
                canonical: alias.canonical.clone(),
                aliases: alias.aliases.clone(),
            };
            self.predicate_alias_store.set_predicate_alias_set(&stored).await.map_err(|e| {
                AphoriaError::Storage(format!("Failed to persist predicate alias: {e}"))
            })?;
        }

        // Update in-memory cache (merge with existing)
        for new_alias in aliases {
            if let Some(existing) =
                self.predicate_aliases.iter_mut().find(|a| a.canonical == new_alias.canonical)
            {
                // Merge aliases
                for alias in new_alias.aliases {
                    if !existing.aliases.contains(&alias) {
                        existing.aliases.push(alias);
                    }
                }
            } else {
                self.predicate_aliases.push(new_alias);
            }
        }

        Ok(())
    }

    /// Add predicate aliases from an imported Trust Pack (in-memory only).
    ///
    /// Deprecated: Use `persist_predicate_aliases` instead to ensure aliases
    /// survive restarts.
    #[deprecated(note = "Use persist_predicate_aliases instead")]
    #[allow(dead_code)]
    pub fn add_predicate_aliases(&mut self, aliases: Vec<PredicateAliasSet>) {
        self.predicate_aliases.extend(aliases);
    }
}
