use crate::models::*;
use crate::search::MemoryIndex;
use crate::store::Store;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::RwLock;

pub struct MemoryState {
    pub base_dir: PathBuf,
    pub graph: Store<KnowledgeGraph>,
    pub search_index: RwLock<MemoryIndex>,
    pub ledger: Store<Vec<CodeChange>>,
    pub sticky: Store<Vec<StickyNote>>,
    pub tasks: Store<Vec<Task>>,
    pub snippets: Store<Vec<Snippet>>,
    pub adrs: Store<Vec<Adr>>,
    pub prefs: Store<HashMap<String, Preference>>,
    pub error_fixes: Store<Vec<ErrorFix>>,
    pub pinned_files: Store<Vec<PinnedFile>>,
    pub session_summaries: Store<Vec<SessionSummary>>,
    pub handoff_memos: Store<Vec<HandoffMemo>>,
    pub env_fingerprints: Store<HashMap<String, EnvFingerprint>>,
    pub env_requirements: Store<Vec<EnvRequirement>>,
    pub milestones: Store<Vec<Milestone>>,
    pub environments: Store<Vec<EnvironmentDetail>>,
    pub pr_checklists: Store<Vec<PrChecklistItem>>,
    pub tech_debts: Store<Vec<TechDebt>>,
    pub gates: Store<Vec<GateRecord>>,
    pub context_workspaces: Store<Vec<ContextWorkspace>>,
    pub activity_tx: tokio::sync::broadcast::Sender<String>,
}

impl MemoryState {
    pub fn deduplicate<T: Eq + std::hash::Hash + Clone>(input: &mut Vec<T>) {
        let mut keys = std::collections::HashSet::new();
        input.retain(|entry| keys.insert(entry.clone()));
    }

    pub fn broadcast_activity(&self, message: &str) {
        let payload = serde_json::json!({
            "type": "activity",
            "data": message
        })
        .to_string();
        let _ = self.activity_tx.send(payload);
    }

    pub fn get_full_graph(&self) -> KnowledgeGraph {
        self.graph.read()
    }

    pub fn read_graph<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&KnowledgeGraph) -> R,
    {
        self.graph.read_with(f)
    }

    pub fn modify_graph<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
        self.graph.modify(update_fn);
    }

    pub async fn rebuild_index(&self) {
        if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
            
            let entities = self.graph.read().entities.values().cloned().collect();
            let tasks = self.tasks.read().clone();
            let snippets = self.snippets.read().clone();
            let adrs = self.adrs.read().clone();

            let handle = new_idx.index_batch(entities, tasks, snippets, adrs);
            let _ = handle.await;

            let _ = new_idx.commit().await;
            if let Ok(mut w) = self.search_index.write() {
                *w = new_idx;
            }
        }
    }
}