use crate::models::*;
use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::*;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;

pub struct LogDecisionHandler;

#[async_trait]
impl McpTool for LogDecisionHandler {
    fn name(&self) -> &'static str {
        "log_decision"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<LogDecisionTool>("log_decision", "Execute log_decision")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LogDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?;

        let idx = state.get_search_index();
        let mut final_id = String::new();

        state.adrs.modify(|adrs| {
            if let Some(superseded_id) = &req.supersedes {
                for old_adr in adrs.iter_mut() {
                    if old_adr.id == *superseded_id {
                        old_adr.status = "superseded".to_string();
                        // Re-index the modified old ADR
                        drop(idx.index_adr(old_adr));
                        break;
                    }
                }
            }

            final_id = format!("ADR-{:04}", adrs.len() + 1);
            let a = Adr {
                id: final_id.clone(),
                title: req.title,
                context: req.context,
                decision: req.decision,
                consequence: req.consequence,
                status: "accepted".to_string(),
                supersedes: req.supersedes,
                timestamp: crate::handlers::utils::now_secs(),
            };

            drop(idx.index_adr(&a));
            adrs.push(a);
        });

        Ok(format!("Decision logged as {}", final_id))
    }
}

pub struct QueryDecisionsHandler;

#[async_trait]
impl McpTool for QueryDecisionsHandler {
    fn name(&self) -> &'static str {
        "query_decisions"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<QueryDecisionsTool>("query_decisions", "Execute query_decisions")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: QueryDecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.adrs.read_with(|adrs| {
            let filtered: Vec<_> = adrs
                .iter()
                .filter(|a| {
                    if let Some(q) = &req.query {
                        contains_ignore_ascii_case(&a.title, q)
                            || contains_ignore_ascii_case(&a.context, q)
                            || contains_ignore_ascii_case(&a.decision, q)
                    } else {
                        true
                    }
                })
                .collect();
            serde_json::to_string(&filtered).map_err(|e| e.to_string())
        })?;
        Ok(data)
    }
}

pub struct DeleteDecisionHandler;

#[async_trait]
impl McpTool for DeleteDecisionHandler {
    fn name(&self) -> &'static str {
        "delete_decision"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<crate::tools::DeleteDecisionTool>(
            "delete_decision",
            "Delete an architectural decision record",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: crate::tools::DeleteDecisionTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut found = false;
        state.adrs.modify(|adrs| {
            if let Some(pos) = adrs.iter().position(|a| a.id == req.id) {
                adrs.remove(pos);
                found = true;
            }
        });

        if found {
            state.rebuild_index().await;
            Ok("Decision deleted successfully".to_string())
        } else {
            Err("Decision not found".to_string())
        }
    }
}

pub struct LogErrorFixHandler;

#[async_trait]
impl McpTool for LogErrorFixHandler {
    fn name(&self) -> &'static str {
        "log_error_fix"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<LogErrorFixTool>("log_error_fix", "Execute log_error_fix")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LogErrorFixTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let text_to_embed = format!("Signature: {}\nSolution: {}", req.signature, req.solution);
        let embedding = crate::embedding::generate_embedding_async(text_to_embed).await.ok();
        state.error_fixes.modify(|fixes| {
            fixes.push(crate::models::ErrorFix {
                signature: req.signature,
                solution: req.solution,
                timestamp: crate::handlers::utils::now_secs(),
                git_commit: req.git_commit,
                git_branch: req.git_branch,
                embedding,
            })
        });
        Ok("Error fix logged".to_string())
    }
}

pub struct SearchErrorFixesHandler;

#[async_trait]
impl McpTool for SearchErrorFixesHandler {
    fn name(&self) -> &'static str {
        "search_error_fixes"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<SearchErrorFixesTool>(
            "search_error_fixes",
            "Execute search_error_fixes",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: SearchErrorFixesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let q = req.query;
        let data = state.error_fixes.read_with(|fixes| {
            let filtered: Vec<_> = fixes
                .iter()
                .filter(|f| {
                    contains_ignore_ascii_case(&f.signature, &q)
                        || contains_ignore_ascii_case(&f.solution, &q)
                })
                .collect();
            serde_json::to_string(&filtered).map_err(|e| e.to_string())
        })?;
        Ok(data)
    }
}

pub struct LogCodeChangeHandler;

#[async_trait]
impl McpTool for LogCodeChangeHandler {
    fn name(&self) -> &'static str {
        "log_code_change"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<LogCodeChangeTool>("log_code_change", "Execute log_code_change")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LogCodeChangeTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.ledger.modify(|ledger| {
            ledger.push(CodeChange {
                timestamp: crate::handlers::utils::now_secs(),
                file_path: req.file_path,
                description: req.description,
                git_commit: req.git_commit,
                git_branch: req.git_branch,
            });
        });
        Ok("Code change logged".to_string())
    }
}

pub struct QueryRecentChangesHandler;

#[async_trait]
impl McpTool for QueryRecentChangesHandler {
    fn name(&self) -> &'static str {
        "query_recent_changes"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<QueryRecentChangesTool>(
            "query_recent_changes",
            "Execute query_recent_changes",
        )
    }

    async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let data = state
            .ledger
            .read_with(|l| serde_json::to_string(l).map_err(|e| e.to_string()))?;
        Ok(data)
    }
}

pub struct LearnPreferenceHandler;

#[async_trait]
impl McpTool for LearnPreferenceHandler {
    fn name(&self) -> &'static str {
        "learn_preference"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<LearnPreferenceTool>("learn_preference", "Execute learn_preference")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LearnPreferenceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.prefs.modify(|prefs| {
            prefs.insert(
                req.key.clone(),
                crate::models::Preference {
                    key: req.key.clone(),
                    value: req.value,
                    updated_at: crate::handlers::utils::now_secs(),
                },
            );
        });
        Ok("Preference learned".to_string())
    }
}

pub struct ReadPreferencesHandler;

#[async_trait]
impl McpTool for ReadPreferencesHandler {
    fn name(&self) -> &'static str {
        "read_preferences"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ReadPreferencesTool>("read_preferences", "Execute read_preferences")
    }

    async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        state
            .prefs
            .read_with(|prefs| serde_json::to_string(prefs).map_err(|e| e.to_string()))
    }
}

pub struct LogTechDebtHandler;

#[async_trait]
impl McpTool for LogTechDebtHandler {
    fn name(&self) -> &'static str {
        "log_tech_debt"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<LogTechDebtTool>("log_tech_debt", "Execute log_tech_debt")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LogTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let text_to_embed = format!("Description: {}\nIdeal Solution: {}", req.description, req.ideal_solution);
        let embedding = crate::embedding::generate_embedding_async(text_to_embed).await.ok();
        state.tech_debts.modify(|debts| {
            debts.push(crate::models::TechDebt {
                id: uuid::Uuid::new_v4().to_string(),
                namespace: req.namespace,
                description: req.description,
                ideal_solution: req.ideal_solution,
                is_resolved: false,
                created_at: crate::handlers::utils::now_secs(),
                git_commit: req.git_commit,
                git_branch: req.git_branch,
                embedding,
            })
        });
        Ok("Tech debt logged".to_string())
    }
}

pub struct ResolveTechDebtHandler;

#[async_trait]
impl McpTool for ResolveTechDebtHandler {
    fn name(&self) -> &'static str {
        "resolve_tech_debt"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ResolveTechDebtTool>(
            "resolve_tech_debt",
            "Execute resolve_tech_debt",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: ResolveTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut found = false;
        state.tech_debts.modify(|debts| {
            for d in debts.iter_mut() {
                if d.id == req.id {
                    d.is_resolved = true;
                    found = true;
                    break;
                }
            }
        });
        if found {
            Ok("Tech debt resolved".to_string())
        } else {
            Err(
                "Tech debt not found. Please verify the tech debt ID using list_tech_debt."
                    .to_string(),
            )
        }
    }
}

pub struct ListTechDebtHandler;

#[async_trait]
impl McpTool for ListTechDebtHandler {
    fn name(&self) -> &'static str {
        "list_tech_debt"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ListTechDebtTool>("list_tech_debt", "Execute list_tech_debt")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: ListTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.tech_debts.read_with(|debts| {
            let filtered: Vec<_> = debts
                .iter()
                .filter(|d| {
                    d.namespace == req.namespace && (req.include_resolved || !d.is_resolved)
                })
                .collect();
            serde_json::to_string(&filtered).map_err(|e| e.to_string())
        })?;
        Ok(data)
    }
}

pub struct OmniSearchHandler;

#[async_trait]
impl McpTool for OmniSearchHandler {
    fn name(&self) -> &'static str {
        "omni_search"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<OmniSearchTool>("omni_search", "Execute omni_search")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: OmniSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let limit = req.limit.unwrap_or(5);
        let include_body = req.include_body.unwrap_or(false);
        let matches = match state
            .get_search_index()
            .search(&req.query, req.namespace.as_deref())
        {
            Ok(m) => m,
            Err(e) => {
                return Err(format!(
                    "Search query failed (possibly malformed Lucene syntax). Error: {}",
                    e
                ));
            }
        };
        println!("OMNI SEARCH MATCHES: {:?}", matches);
        let q = req.query.clone();
        let query_emb = crate::embedding::generate_embedding_async(q.clone()).await.unwrap_or_default();

        let kg_json = state.read_graph(|full| {
            let mut kg_entities = std::collections::HashMap::new();
            let mut count = 0;
            for (id, doc_type, _, _, _) in &matches {
                if doc_type == "entity"
                    && let Some(e) = full.entities.get(id)
                {
                    if count >= limit {
                        continue;
                    }
                    count += 1;
                    if !include_body {
                        let mut summary = e.clone();
                        summary.observations = vec![];
                        kg_entities.insert(id.clone(), summary);
                    } else {
                        kg_entities.insert(id.clone(), e.clone());
                    }
                }
            }
            serde_json::to_value(&kg_entities).map_err(|e| e.to_string())
        })?;

        let mut matched_tasks = std::collections::HashSet::new();
        let mut matched_snippets = std::collections::HashSet::new();
        let mut matched_adrs = std::collections::HashSet::new();
        for (id, typ, _, _, _) in &matches {
            match typ.as_str() {
                "task" => {
                    matched_tasks.insert(id.as_str());
                }
                "snippet" => {
                    matched_snippets.insert(id.as_str());
                }
                "adr" => {
                    matched_adrs.insert(id.as_str());
                }
                _ => {}
            }
        }

        let tasks_json = state.tasks.read_with(|all_tasks| {
            let filtered: Vec<_> = all_tasks
                .iter()
                .filter(|t| matched_tasks.contains(t.id.as_str()))
                .take(limit)
                .map(|t| {
                    if !include_body {
                        let mut summary = t.clone();
                        summary.description = "".to_string();
                        summary.acceptance_criteria = vec![];
                        summary
                    } else {
                        t.clone()
                    }
                })
                .collect();
            serde_json::to_value(&filtered).map_err(|e| e.to_string())
        })?;

        let snippets_json = state.snippets.read_with(|all_snippets| {
            let mut scored: Vec<_> = all_snippets.iter().map(|s| {
                let mut score = 0.0;
                if matched_snippets.contains(s.name.as_str()) || contains_ignore_ascii_case(&s.name, &q) || contains_ignore_ascii_case(&s.description, &q) {
                    score += 1.0;
                }
                if let Some(emb) = &s.embedding {
                    score += crate::embedding::cosine_similarity(&query_emb, emb);
                }
                (s, score)
            }).filter(|(_, score)| *score > 0.4).collect();
            scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

            let filtered: Vec<_> = scored.into_iter().take(limit).map(|(s, _)| {
                if !include_body {
                    let mut summary = s.clone();
                    summary.code = "".to_string();
                    summary
                } else {
                    s.clone()
                }
            }).collect();
            serde_json::to_value(&filtered).map_err(|e| e.to_string())
        })?;

        let adrs_json = state.adrs.read_with(|all_adrs| {
            let filtered: Vec<_> = all_adrs
                .iter()
                .filter(|a| matched_adrs.contains(a.id.as_str()))
                .take(limit)
                .map(|a| {
                    if !include_body {
                        let mut summary = a.clone();
                        summary.context = "".to_string();
                        summary.decision = "".to_string();
                        summary.consequence = "".to_string();
                        summary
                    } else {
                        a.clone()
                    }
                })
                .collect();
            serde_json::to_value(&filtered).map_err(|e| e.to_string())
        })?;

        let q = req.query;
        let tech_debts_json = state.tech_debts.read_with(|debts| {
            let mut scored: Vec<_> = debts.iter().map(|d| {
                if req.namespace.as_ref().is_some_and(|ns| d.namespace != *ns) {
                    return (d, 0.0);
                }
                let mut score = 0.0;
                if contains_ignore_ascii_case(&d.description, &q) || contains_ignore_ascii_case(&d.ideal_solution, &q) {
                    score += 1.0;
                }
                if let Some(emb) = &d.embedding {
                    score += crate::embedding::cosine_similarity(&query_emb, emb);
                }
                (d, score)
            }).filter(|(_, score)| *score > 0.4).collect();
            scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

            let filtered: Vec<_> = scored.into_iter().take(limit).map(|(d, _)| {
                if !include_body {
                    let mut summary = d.clone();
                    summary.description = "".to_string();
                    summary.ideal_solution = "".to_string();
                    summary
                } else {
                    d.clone()
                }
            }).collect();
            serde_json::to_value(&filtered).map_err(|e| e.to_string())
        })?;

        let memos_json = state.handoff_memos.read_with(|memos| {
            let filtered: Vec<_> = memos
                .iter()
                .filter(|m| {
                    req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
                        && contains_ignore_ascii_case(&m.content, &q)
                })
                .take(limit)
                .map(|m| {
                    if !include_body {
                        let mut summary = m.clone();
                        summary.content = "".to_string();
                        summary
                    } else {
                        m.clone()
                    }
                })
                .collect();
            serde_json::to_value(&filtered).map_err(|e| e.to_string())
        })?;

        let error_fixes_json = state.error_fixes.read_with(|fixes| {
            let mut scored: Vec<_> = fixes.iter().map(|f| {
                let mut score = 0.0;
                if contains_ignore_ascii_case(&f.signature, &q) || contains_ignore_ascii_case(&f.solution, &q) {
                    score += 1.0;
                }
                if let Some(emb) = &f.embedding {
                    score += crate::embedding::cosine_similarity(&query_emb, emb);
                }
                (f, score)
            }).filter(|(_, score)| *score > 0.4).collect();
            scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            
            let filtered: Vec<_> = scored.into_iter().take(limit).map(|(f, _)| f.clone()).collect();
            serde_json::to_value(&filtered).map_err(|e| e.to_string())
        })?;

        let report = serde_json::json!({
            "knowledge_graph": kg_json,
            "tasks": tasks_json,
            "snippets": snippets_json,
            "adrs": adrs_json,
            "tech_debts": tech_debts_json,
            "handoff_memos": memos_json,
            "error_fixes": error_fixes_json
        });
        Ok(report.to_string())
    }
}

pub struct GetProjectHealthHandler;

#[async_trait]
impl McpTool for GetProjectHealthHandler {
    fn name(&self) -> &'static str {
        "get_project_health"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<GetProjectHealthTool>(
            "get_project_health",
            "Execute get_project_health",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: GetProjectHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let active_tasks = state
            .tasks
            .read_with(|tasks| tasks.iter().filter(|t| t.status != "done").count());
        let unresolved_debt = state.tech_debts.read_with(|debts| {
            debts
                .iter()
                .filter(|d| d.namespace == req.namespace && !d.is_resolved)
                .count()
        });
        let unread_memos = state.handoff_memos.read_with(|memos| {
            memos
                .iter()
                .filter(|m| m.namespace == req.namespace)
                .count()
        });
        let active_milestones = state.milestones.read_with(|milestones| {
            milestones
                .iter()
                .filter(|m| m.namespace == req.namespace && m.status != "done")
                .count()
        });
        let remaining_checklists = state.pr_checklists.read_with(|checklists| {
            checklists
                .iter()
                .filter(|c| c.namespace == req.namespace)
                .count()
        });

        let report = serde_json::json!({
            "active_tasks": active_tasks,
            "unresolved_tech_debt": unresolved_debt,
            "unread_handoff_memos": unread_memos,
            "active_milestones": active_milestones,
            "remaining_pr_checklist_items": remaining_checklists
        });
        Ok(report.to_string())
    }
}

use crate::handlers::utils::*;

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

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

        let handler = LogErrorFixHandler;
        let args = json!({
            "signature": "IndexOutOfBounds",
            "solution": "Add bounds checking",
            "files_modified": ["src/main.rs"],
            "git_commit": "abcdef",
            "git_branch": "main"
        });

        let res = handler.execute(args, state.clone()).await.unwrap();
        assert!(res.contains("Error fix logged"));
    }

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

        let handler = GetProjectHealthHandler;
        let args = json!({"namespace": "global"});

        let res = handler.execute(args, state.clone()).await.unwrap();
        assert!(res.contains("unresolved_tech_debt"));
    }

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

        let decision_handler = LogDecisionHandler;
        let args_dec = json!({
            "title": "Architecture",
            "context": "Needs DB",
            "decision": "Use SQLite",
            "consequence": "Simple",
        });
        let res1 = decision_handler
            .execute(args_dec, state.clone())
            .await
            .unwrap();
        assert_eq!(res1, "Decision logged as ADR-0001");

        let debt_handler = LogTechDebtHandler;
        let args_debt = json!({
            "title": "Hardcoded path",
            "description": "Hardcoded path",
            "location": "main.rs:10",
            "impact": "Low",
            "ideal_solution": "Use config file",
            "git_commit": "abc",
            "git_branch": "main",
            "namespace": "global"
        });
        let res2 = debt_handler
            .execute(args_debt, state.clone())
            .await
            .unwrap();
        assert_eq!(res2, "Tech debt logged");

        let list_debt = ListTechDebtHandler;
        let res3 = list_debt
            .execute(
                json!({"namespace": "global", "include_resolved": false}),
                state.clone(),
            )
            .await
            .unwrap();
        assert!(res3.contains("Hardcoded path"));

        let pref_handler = LearnPreferenceHandler;
        let args_pref = json!({
            "key": "formatting",
            "value": "use spaces",
        });
        let res4 = pref_handler
            .execute(args_pref, state.clone())
            .await
            .unwrap();
        assert_eq!(res4, "Preference learned");

        let read_pref = ReadPreferencesHandler;
        let res5 = read_pref.execute(json!({}), state.clone()).await.unwrap();
        assert!(res5.contains("use spaces"));
    }

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

        let code_handler = LogCodeChangeHandler;
        let args_code = json!({
            "file_path": "main.rs",
            "description": "refactor",
            "git_commit": "def",
            "git_branch": "main"
        });
        code_handler
            .execute(args_code, state.clone())
            .await
            .unwrap();

        let query_changes = QueryRecentChangesHandler;
        let res_changes = query_changes
            .execute(json!({}), state.clone())
            .await
            .unwrap();
        assert!(res_changes.contains("main.rs"));

        let debt_handler = LogTechDebtHandler;
        let args_debt = json!({
            "title": "Debt 1",
            "description": "Needs refactor",
            "location": "main.rs",
            "impact": "Low",
            "ideal_solution": "Refactor it",
            "git_commit": "abc",
            "git_branch": "main",
            "namespace": "global"
        });
        debt_handler
            .execute(args_debt, state.clone())
            .await
            .unwrap();

        // resolve it
        let list_debt = ListTechDebtHandler;
        let debt_list = list_debt
            .execute(
                json!({"namespace": "global", "include_resolved": false}),
                state.clone(),
            )
            .await
            .unwrap();
        let uuid_start = debt_list.find("id\":\"").unwrap() + 5;
        let uuid = &debt_list[uuid_start..uuid_start + 36];

        let resolve_debt = ResolveTechDebtHandler;
        resolve_debt
            .execute(json!({"id": uuid}), state.clone())
            .await
            .unwrap();
    }

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

        let task = crate::models::Task {
            id: "omni-1".to_string(),
            title: "Omni Task".to_string(),
            description: "Testing omni search functionality".to_string(),
            status: "open".to_string(),
            created_at: 0,
            updated_at: 0,
            git_branch: None,
            parent_id: None,
            expires_at: None,
            dependencies: vec![],
            acceptance_criteria: vec![],
        };

        {
            state.tasks.modify(|t| {
                t.push(task.clone());
            });
        }

        state.rebuild_index().await;
        state.get_search_index().reader.reload().unwrap();

        let omni = OmniSearchHandler;
        let omni_res = omni
            .execute(json!({"query": "Omni"}), state.clone())
            .await
            .unwrap();
        println!("OMNI RES: {}", omni_res);
        assert!(
            omni_res.contains("omni-1"),
            "omni search should return results containing the task id"
        );
    }

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

        let omni = OmniSearchHandler;

        // Pass a malformed Lucene query (unclosed parenthesis)
        let omni_res = omni
            .execute(json!({"query": "title: (unclosed"}), state.clone())
            .await;

        assert!(omni_res.is_err());
        let err_msg = omni_res.unwrap_err();
        assert!(err_msg.contains("malformed Lucene syntax"));
    }
}