use crate::models::*;
use crate::search::MemoryIndex;
use crate::store::Store;
use crate::vector_db::VectorDB;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct GenericEvent {
    pub topic: String,
    pub session_id: Option<String>,
    pub payload: serde_json::Value,
}

pub struct ProjectStores {
    pub tasks: Store<Vec<Task>>,
    pub milestones: Store<Vec<Milestone>>,
    pub pr_checklists: Store<Vec<PrChecklistItem>>,
    pub context_workspaces: Store<Vec<ContextWorkspace>>,
    pub pinned_files: Store<Vec<PinnedFile>>,
    pub snapshots: Store<Vec<StateSnapshot>>,
}

pub struct CodeStores {
    pub ledger: Store<Vec<CodeChange>>,
    pub snippets: Store<Vec<Snippet>>,
    pub adrs: Store<Vec<Adr>>,
    pub error_fixes: Store<Vec<ErrorFix>>,
    pub tech_debts: Store<Vec<TechDebt>>,
    pub sticky: Store<Vec<StickyNote>>,
    pub hypotheses: Store<Vec<Hypothesis>>,
}

pub struct EnvironmentStores {
    pub env_fingerprints: Store<HashMap<String, EnvFingerprint>>,
    pub env_requirements: Store<Vec<EnvRequirement>>,
    pub environments: Store<Vec<EnvironmentDetail>>,
    pub gates: Store<Vec<GateRecord>>,
    pub prefs: Store<HashMap<String, Preference>>,
}

pub struct TelemetryStores {
    pub session_summaries: Store<Vec<SessionSummary>>,
    pub handoff_memos: Store<Vec<HandoffMemo>>,
    pub recent_activities: Store<std::collections::VecDeque<serde_json::Value>>,
    pub terminal_history: Store<std::collections::VecDeque<TerminalHistory>>,
    pub agent_signals: Store<Vec<AgentSignal>>,
}

pub struct MemoryState {
    pub base_dir: PathBuf,
    pub clipboard_watch_mode: tokio::sync::RwLock<bool>,
    pub graph: Store<KnowledgeGraph>,
    pub search_index: RwLock<MemoryIndex>,
    pub vector_db: tokio::sync::RwLock<Option<VectorDB>>,

    pub project: ProjectStores,
    pub code: CodeStores,
    pub env: EnvironmentStores,
    pub telemetry: TelemetryStores,

    pub activity_tx: tokio::sync::broadcast::Sender<String>,
    pub event_bus_tx: tokio::sync::broadcast::Sender<GenericEvent>,
    pub ollama: Arc<crate::ollama::OllamaClient>,
}

impl MemoryState {
    pub fn new(base_dir_str: &str) -> Self {
        let base = std::path::PathBuf::from(base_dir_str);
        std::fs::create_dir_all(&base).expect("Failed to create store dir");

        let db = crate::db::init_redb(&base);

        let state = Self {
            ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
            clipboard_watch_mode: tokio::sync::RwLock::new(false),
            graph: Store::new("knowledge_graph_master", db.clone()),
            base_dir: base.clone(),
            search_index: RwLock::new(match crate::search::MemoryIndex::new(&base) {
                Ok(idx) => idx,
                Err(e) => {
                    let log_path = dirs::home_dir()
                        .unwrap_or_default()
                        .join(".gemini/mcp_memory/daemon_error.log");
                    let _ =
                        std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
                    std::process::exit(1);
                }
            }),
            vector_db: tokio::sync::RwLock::new(None),

            project: ProjectStores {
                tasks: Store::new("tasks", db.clone()),
                milestones: Store::new("milestones", db.clone()),
                pr_checklists: Store::new("pr_checklists", db.clone()),
                context_workspaces: Store::new("context_workspaces", db.clone()),
                pinned_files: Store::new("pinned_files", db.clone()),
                snapshots: Store::new("state_snapshots", db.clone()),
            },
            code: CodeStores {
                ledger: Store::new("audit_ledger", db.clone()),
                snippets: Store::new("snippets", db.clone()),
                adrs: Store::new("adrs", db.clone()),
                error_fixes: Store::new("error_fixes", db.clone()),
                tech_debts: Store::new("tech_debts", db.clone()),
                sticky: Store::new("sticky_notes", db.clone()),
                hypotheses: Store::new("hypotheses", db.clone()),
            },
            env: EnvironmentStores {
                env_fingerprints: Store::new("env_fingerprints", db.clone()),
                env_requirements: Store::new("env_requirements", db.clone()),
                environments: Store::new("environments", db.clone()),
                gates: Store::new("gates", db.clone()),
                prefs: Store::new("preferences", db.clone()),
            },
            telemetry: TelemetryStores {
                session_summaries: Store::new("session_summaries", db.clone()),
                handoff_memos: Store::new("handoff_memos", db.clone()),
                recent_activities: Store::new("recent_activities", db.clone()),
                terminal_history: Store::new("terminal_history", db.clone()),
                agent_signals: Store::new("agent_signals", db.clone()),
            },

            activity_tx: tokio::sync::broadcast::channel(100).0,
            event_bus_tx: tokio::sync::broadcast::channel(1000).0,
        };

        // Normalize pre-existing graph entity and relation types
        state.graph.modify(|g| {
            for entity in g.entities.values_mut() {
                entity.entity_type = crate::models::normalize_entity_type(&entity.entity_type);
            }
            for relation in g.relations.iter_mut() {
                relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
            }
        });

        state
    }

    pub fn broadcast_activity(&self, category: &str, message: &str) {
        self.record_activity(category, message, None);
    }

    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 fn get_search_index(&self) -> MemoryIndex {
        self.search_index
            .read()
            .unwrap_or_else(|e| e.into_inner())
            .clone()
    }

    pub fn search(self: &Arc<Self>) -> SearchService {
        SearchService::new(self.clone())
    }

    pub async fn rebuild_index(self: &Arc<Self>) {
        let idx = self.search_index.read().unwrap().clone();
        idx.delete_all();

        let entities: Vec<_> = self
            .graph
            .read_with(|g| g.entities.values().cloned().collect());
        let tasks = self.project.tasks.read_with(|t| t.clone());
        let snippets = self.code.snippets.read_with(|s| s.clone());
        let adrs = self.code.adrs.read_with(|a| a.clone());

        tracing::info!(
            "rebuild_index: found {} entities, {} tasks",
            entities.len(),
            tasks.len()
        );

        let idx_clone = idx.clone();
        tokio::task::spawn_blocking(move || {
            // tracing::info!("spawn_blocking started in rebuild_index");
            for e in entities {
                idx_clone.add_entity_sync(&e);
            }
            for task in tasks {
                idx_clone.add_task_sync(&task);
            }
            for snippet in snippets {
                idx_clone.add_snippet_sync(&snippet);
            }
            for adr in adrs {
                idx_clone.add_adr_sync(&adr);
            }
        })
        .await
        .unwrap_or_else(|e| {
            tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
        });

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

    pub fn record_activity(&self, category: &str, summary: &str, details: Option<&str>) {
        let ts = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;

        let category_upper = category.to_uppercase();

        let activity = ActivityRecord {
            timestamp: ts,
            category: category_upper,
            summary: summary.to_string(),
            details: details.map(|s| s.to_string()),
        };

        let record_val = serde_json::to_value(&activity).unwrap_or_default();

        self.telemetry.recent_activities.modify(|activities| {
            activities.push_front(record_val.clone());
            if activities.len() > 100 {
                activities.pop_back();
            }
        });

        let payload = serde_json::json!({
            "jsonrpc": "2.0",
            "method": "notifications/activity",
            "params": activity
        })
        .to_string();

        let _ = self.activity_tx.send(payload);
    }

    pub fn record_terminal_history(&self, payload: TerminalHistory) {
        self.telemetry.terminal_history.modify(|history| {
            history.push_front(payload);
            if history.len() > 100 {
                history.pop_back();
            }
        });
    }
}

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

    #[tokio::test]
    async fn test_memory_state_initialization() {
        let dir = tempdir().unwrap();
        let state = MemoryState::new(dir.path().to_str().unwrap());

        // Ensure state fields are properly initialized
        assert_eq!(state.base_dir, dir.path());

        // Write a test value
        state.project.tasks.modify(|tasks| {
            tasks.push(Task {
                id: "123".to_string(),
                title: "Test Task".to_string(),
                description: "".to_string(),
                status: "active".to_string(),
                created_at: 0,
                updated_at: 0,
                dependencies: vec![],
                acceptance_criteria: vec![],
                git_branch: None,
                parent_id: None,
                expires_at: None,
            });
        });

        // Ensure it is saved
        state.project.tasks.read_with(|tasks| {
            assert_eq!(tasks.len(), 1);
            assert_eq!(tasks[0].id, "123");
        });

        // Test rebuild index
        let arc_state = Arc::new(state);
        arc_state.rebuild_index().await;

        // Check search index initialization
        let idx = arc_state.search_index.read().unwrap();
        // Force reload reader to ensure it sees the commit made by rebuild_index
        idx.reader.reload().unwrap();
        // tracing::info!(
        //     "Index reader doc count: {}",
        //     idx.reader.searcher().num_docs()
        // );

        let _all_docs = idx.search("Test", None).expect("Search failed");
        // tracing::info!("All docs for 'Test': {:?}", all_docs);

        // Verify the task added synchronously is actually searchable
        let results = idx.search("Test", None).expect("Search failed");
        assert_eq!(results.len(), 1, "Expected exactly 1 search result");
        assert_eq!(
            results[0].0, "123",
            "Expected the result to be the task we just added"
        );
        assert_eq!(results[0].1, "task", "Expected document type to be task");
    }

    #[tokio::test]
    async fn test_record_and_broadcast_activity() {
        let dir = tempdir().unwrap();
        let state = MemoryState::new(dir.path().to_str().unwrap());
        let mut rx = state.activity_tx.subscribe();

        // 1. Record an activity with details
        state.record_activity("code_change", "Refactored state.rs", Some("Updated ActivityRecord schema"));

        // Verify recent_activities store
        let activities: Vec<ActivityRecord> = state.telemetry.recent_activities.read_with(|act| {
            act.iter()
                .filter_map(|v| serde_json::from_value(v.clone()).ok())
                .collect()
        });

        assert_eq!(activities.len(), 1);
        assert_eq!(activities[0].category, "CODE_CHANGE");
        assert_eq!(activities[0].summary, "Refactored state.rs");
        assert_eq!(activities[0].details, Some("Updated ActivityRecord schema".to_string()));
        assert!(activities[0].timestamp > 1_700_000_000_000, "Timestamp must be in epoch milliseconds");

        // Verify broadcast channel message
        let broadcast_msg = rx.recv().await.expect("Expected broadcast notification");
        let broadcast_val: serde_json::Value = serde_json::from_str(&broadcast_msg).expect("Valid JSON");
        assert_eq!(broadcast_val["jsonrpc"], "2.0");
        assert_eq!(broadcast_val["method"], "notifications/activity");
        assert_eq!(broadcast_val["params"]["category"], "CODE_CHANGE");

        // 2. Broadcast an activity without details
        state.broadcast_activity("task", "Completed live activity fix");

        let activities_updated: Vec<ActivityRecord> = state.telemetry.recent_activities.read_with(|act| {
            act.iter()
                .filter_map(|v| serde_json::from_value(v.clone()).ok())
                .collect()
        });

        assert_eq!(activities_updated.len(), 2);
        assert_eq!(activities_updated[0].category, "TASK");
        assert_eq!(activities_updated[0].summary, "Completed live activity fix");
        assert_eq!(activities_updated[0].details, None);
        assert!(activities_updated[0].timestamp >= activities_updated[1].timestamp);
    }
}

use crate::embedding::{cosine_similarity, generate_embedding_async, generate_embeddings_async};

pub struct UnifiedSearchResult {
    pub id: String,
    pub doc_type: String,
    pub title: String,
    pub body: String,
    pub score: f32,
}

pub struct SearchService {
    state: Arc<MemoryState>,
}

impl SearchService {
    pub fn new(state: Arc<MemoryState>) -> Self {
        Self { state }
    }

    pub async fn semantic_search(
        &self,
        query: &str,
        _filter_namespace: Option<&str>,
        limit: usize,
    ) -> crate::error::Result<Vec<UnifiedSearchResult>> {
        let query_emb = generate_embedding_async(query.to_string())
            .await
            .unwrap_or_default();
        let mut results = Vec::new();

        let mut vdb_search = false;
        if let Some(vdb) = &*self.state.vector_db.read().await {
            vdb_search = true;
            if let Ok(search_results) = vdb.search(query_emb.clone(), limit as u64).await {
                for res in search_results {
                    results.push(UnifiedSearchResult {
                        id: res.id.clone(),
                        doc_type: res.doc_type.clone(),
                        title: res.id,
                        body: res.text,
                        score: res.score,
                    });
                }
            }
        }

        if !vdb_search {
            let mut texts_to_embed = Vec::new();
            let mut metadata = Vec::new();

            let snippets = self.state.code.snippets.read_with(|snips| snips.clone());
            for snippet in snippets {
                let combined = format!("{} {} {}", snippet.name, snippet.description, snippet.code);
                texts_to_embed.push(combined);
                metadata.push((snippet.name, "snippet".to_string(), snippet.description));
            }

            let sticky = self.state.code.sticky.read_with(|s| s.clone());
            for note in sticky {
                texts_to_embed.push(note.content.clone());
                metadata.push((
                    "StickyNote".to_string(),
                    "sticky".to_string(),
                    note.content.chars().take(200).collect::<String>(),
                ));
            }

            if let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await {
                for (emb, meta) in embeddings.into_iter().zip(metadata) {
                    let sim = cosine_similarity(&query_emb, &emb);
                    results.push(UnifiedSearchResult {
                        id: meta.0.clone(),
                        doc_type: meta.1.clone(),
                        title: meta.0,
                        body: meta.2,
                        score: sim,
                    });
                }
            }

            results.sort_by(|a, b| {
                b.score
                    .partial_cmp(&a.score)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            results.truncate(limit);
        }

        Ok(results)
    }

    pub fn keyword_search(
        &self,
        query: &str,
        filter_namespace: Option<&str>,
        limit: usize,
    ) -> crate::error::Result<Vec<UnifiedSearchResult>> {
        let idx = self.state.get_search_index();
        let matches = idx
            .search(query, filter_namespace)
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))?;

        let mut results = Vec::new();
        for (id, doc_type, title, body, score) in matches.into_iter().take(limit) {
            results.push(UnifiedSearchResult {
                id,
                doc_type,
                title,
                body,
                score,
            });
        }
        Ok(results)
    }
}