//! Gandalf: Constraint Extractor Agent
//!
//! Gandalf is responsible for extracting architectural constraints from the
//! Kuzu knowledge graph, enriching SCROLL tasks with constraint context,
//! validating constraints against existing patterns, and resolving pattern
//! examples from the codebase.
//!
//! Named after the wizard who guards the passage -- Gandalf guards the
//! architectural integrity of the codebase by ensuring every code generation
//! task carries the full weight of existing constraints.

use std::sync::Arc;

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tracing::{debug, info, warn};

use crate::db::{KuzuGraph, SqliteDb};
use crate::error::FurnaceError;
use crate::knowledge::KnowledgeManager;
use crate::recipes::manager::RecipeManager;
use crate::search::SearchOrchestrator;
use crate::types::{
    Constraint, ConstraintSource, ConstraintType, LlmProvider, PatternExample, SearchRequest,
    Severity, Scroll, Task, TaskType, ValidationResult,
};

// ---------------------------------------------------------------------------
// Local types
// ---------------------------------------------------------------------------

/// Result of Gandalf's LLM task augmentation pass.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskAugmentation {
    /// Whether the LLM verified the tasks are sufficient.
    pub verified: bool,
    /// Additional tasks suggested by the LLM.
    pub suggestions: Vec<TaskSuggestion>,
    /// Warnings about potential issues.
    pub warnings: Vec<String>,
    /// Groupings of related tasks (e.g. module + tests).
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub groupings: Vec<TaskGrouping>,
}

/// A task suggestion from Gandalf's LLM augmentation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskSuggestion {
    pub description: String,
    pub target_file: String,
    pub task_type: String,
    /// Optional group label for this suggested task.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub group: Option<String>,
}

/// A grouping of related existing tasks.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskGrouping {
    pub group_label: String,
    pub task_ids: Vec<String>,
}

/// Aggregated constraint context for a single code-generation task.
///
/// Built by [`Gandalf::build_constraint_context`], this struct is the primary
/// input that Gimli (and other downstream agents) use to understand *what
/// rules the generated code must satisfy*.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConstraintContext {
    /// Constraints that directly apply to the task's target file.
    pub direct_constraints: Vec<Constraint>,

    /// Constraints inherited from files that the target depends on.
    pub inherited_constraints: Vec<Constraint>,

    /// Concrete code examples that demonstrate the expected patterns.
    pub pattern_examples: Vec<PatternExample>,

    /// Type definitions and signatures discovered from the graph that
    /// the generated code should be compatible with.
    pub related_types: Vec<RelatedType>,

    /// Import/dependency relationships the generated code should respect.
    pub dependency_context: Vec<DependencyInfo>,

    /// Human-readable summary of the constraint context, suitable for
    /// inclusion in an LLM prompt.
    pub summary: String,

    /// Rendered recipe template (if the SCROLL specifies a recipe).
    /// Populated by Gandalf when a recipe_used is set on the SCROLL metadata.
    pub recipe_context: Option<String>,
}

/// A type definition discovered from the knowledge graph that is relevant
/// to the current task.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelatedType {
    pub file_path: String,
    pub type_name: String,
    pub kind: String,
    pub fields: Vec<String>,
}

/// Dependency edge relevant to the current task.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DependencyInfo {
    pub from_file: String,
    pub to_file: String,
    pub imported_items: Vec<String>,
}

// ---------------------------------------------------------------------------
// Gandalf
// ---------------------------------------------------------------------------

/// The Gandalf agent.
///
/// Gandalf reads the Kuzu knowledge graph and the SQLite SCROLL store to
/// produce rich constraint context that downstream agents (primarily Gimli)
/// consume when generating or validating code.
pub struct Gandalf {
    kuzu: Arc<KuzuGraph>,
    sqlite: Arc<SqliteDb>,
    search: Option<Arc<SearchOrchestrator>>,
    kb_manager: Option<Arc<KnowledgeManager>>,
    recipe_manager: Option<Arc<RecipeManager>>,
    llm: Option<Arc<dyn LlmProvider>>,
}

impl Gandalf {
    /// Create a new Gandalf instance.
    pub fn new(kuzu: Arc<KuzuGraph>, sqlite: Arc<SqliteDb>) -> Self {
        info!("Gandalf initialized");
        Self {
            kuzu,
            sqlite,
            search: None,
            kb_manager: None,
            recipe_manager: None,
            llm: None,
        }
    }

    /// Attach a SearchOrchestrator for tiered search when resolving examples.
    pub fn with_search(mut self, search: Arc<SearchOrchestrator>) -> Self {
        self.search = Some(search);
        self
    }

    /// Attach a KnowledgeManager for KB-sourced constraint enrichment.
    pub fn with_kb(mut self, kb_manager: Arc<KnowledgeManager>) -> Self {
        self.kb_manager = Some(kb_manager);
        self
    }

    /// Attach a RecipeManager for rendering recipe templates in constraint context.
    pub fn with_recipes(mut self, rm: Arc<RecipeManager>) -> Self {
        self.recipe_manager = Some(rm);
        self
    }

    /// Attach an LLM provider for task verification and augmentation.
    pub fn with_llm(mut self, llm: Arc<dyn LlmProvider>) -> Self {
        self.llm = Some(llm);
        self
    }

    // -----------------------------------------------------------------------
    // Public API
    // -----------------------------------------------------------------------

    /// Extract all constraints that apply to a given SCROLL.
    ///
    /// This walks the SCROLL's tasks, examines each target file in the graph,
    /// and collects both user-specified constraints (stored on the SCROLL
    /// itself) and discovered/inferred constraints from the dependency graph.
    pub async fn extract_constraints(&self, scroll_id: &str) -> Result<Vec<Constraint>> {
        info!(scroll_id = %scroll_id, "Extracting constraints for SCROLL");

        // 1. Load the SCROLL and its user-specified constraints.
        let scroll = self
            .sqlite
            .get_scroll(scroll_id)
            .await
            .map_err(|e| FurnaceError::Agent {
                agent: "Gandalf".into(),
                message: format!("Failed to load SCROLL {}: {}", scroll_id, e),
            })?;

        let mut constraints: Vec<Constraint> = scroll.constraints.clone();
        debug!(
            count = constraints.len(),
            "Loaded user-specified constraints"
        );

        // 2. Load tasks for this SCROLL.
        let tasks = self
            .sqlite
            .get_tasks_for_scroll(scroll_id)
            .await
            .map_err(|e| FurnaceError::Agent {
                agent: "Gandalf".into(),
                message: format!("Failed to load tasks for SCROLL {}: {}", scroll_id, e),
            })?;

        // 3. For each task, discover constraints from the graph.
        for task in &tasks {
            let discovered = self.discover_constraints_for_file(&task.target_file).await?;
            for c in discovered {
                if !constraints.iter().any(|existing| existing.id == c.id) {
                    constraints.push(c);
                }
            }
        }

        info!(
            scroll_id = %scroll_id,
            total = constraints.len(),
            "Constraint extraction complete"
        );

        Ok(constraints)
    }

    /// Build a full [`ConstraintContext`] for a single task.
    ///
    /// This is the main entry-point that Gimli calls before code generation.
    /// The returned context carries everything the LLM needs to respect
    /// existing patterns and architectural invariants.
    pub async fn build_constraint_context(&self, task: &Task) -> Result<ConstraintContext> {
        info!(
            task_id = %task.id,
            target = %task.target_file,
            "Building constraint context"
        );

        // 1. Extract direct constraints from the SCROLL.
        let scroll = self.sqlite.get_scroll(&task.scroll_id).await?;
        let direct_constraints = self.filter_constraints_for_file(
            &scroll.constraints,
            &task.target_file,
        );

        // 2. Discover inherited constraints from dependencies.
        let inherited_constraints = self
            .discover_inherited_constraints(&task.target_file)
            .await
            .unwrap_or_else(|e| {
                warn!(error = %e, "Failed to discover inherited constraints, continuing with empty set");
                Vec::new()
            });

        // 3. Find related types from the graph.
        let related_types = self
            .discover_related_types(&task.target_file)
            .await
            .unwrap_or_else(|e| {
                warn!(error = %e, "Failed to discover related types, continuing with empty set");
                Vec::new()
            });

        // 4. Build dependency context.
        let dependency_context = self
            .build_dependency_context(&task.target_file)
            .await
            .unwrap_or_else(|e| {
                warn!(error = %e, "Failed to build dependency context, continuing with empty set");
                Vec::new()
            });

        // 4b. Enrich with KB-sourced constraints.
        let kb_constraints = self.enrich_from_kb(task).await;
        let mut direct_constraints = direct_constraints;
        for c in kb_constraints {
            if !direct_constraints.iter().any(|existing| existing.id == c.id) {
                direct_constraints.push(c);
            }
        }

        // 5. Resolve pattern examples that match any constraint's pattern.
        let mut pattern_examples = Vec::new();
        // Cap pattern resolution to first 20 constraints to prevent
        // excessive async calls in large constraint sets.
        for constraint in direct_constraints
            .iter()
            .chain(inherited_constraints.iter())
            .take(20)
        {
            if !constraint.pattern.is_empty() {
                let mut examples = self
                    .resolve_pattern_examples(&constraint.pattern)
                    .await
                    .unwrap_or_else(|e| {
                        warn!(pattern = %constraint.pattern, error = %e, "Failed to resolve pattern examples");
                        Vec::new()
                    });
                pattern_examples.append(&mut examples);
            }
        }

        // Deduplicate examples by file path.
        pattern_examples.sort_by(|a, b| b.relevance.partial_cmp(&a.relevance).unwrap_or(std::cmp::Ordering::Equal));
        pattern_examples.dedup_by(|a, b| a.file_path == b.file_path);

        // Cap examples at a reasonable count to avoid prompt bloat.
        pattern_examples.truncate(10);

        // 6. Build human-readable summary.
        let summary = self.build_summary(
            task,
            &direct_constraints,
            &inherited_constraints,
            &related_types,
            &dependency_context,
        );

        // 7. Enrich with recipe context if the SCROLL specifies a recipe.
        let recipe_context = if let Some(ref recipe_name) = scroll.metadata.recipe_used {
            if let Some(ref rm) = self.recipe_manager {
                let mut vars = std::collections::HashMap::new();
                vars.insert("code".to_string(), task.description.clone());
                // Detect language from target file extension.
                let language = task
                    .target_file
                    .rsplit('.')
                    .next()
                    .unwrap_or("auto-detect")
                    .to_string();
                vars.insert("language".to_string(), language);
                match tokio::time::timeout(
                    std::time::Duration::from_secs(10),
                    rm.render_recipe(recipe_name, &vars),
                ).await {
                    Ok(Ok(rendered)) => Some(rendered),
                    Ok(Err(e)) => {
                        warn!(error = %e, "Failed to render recipe '{}'", recipe_name);
                        None
                    }
                    Err(_) => {
                        warn!("Recipe rendering timed out after 10s for '{}'", recipe_name);
                        None
                    }
                }
            } else {
                None
            }
        } else {
            None
        };

        let context = ConstraintContext {
            direct_constraints,
            inherited_constraints,
            pattern_examples,
            related_types,
            dependency_context,
            summary,
            recipe_context,
        };

        debug!(
            task_id = %task.id,
            direct = context.direct_constraints.len(),
            inherited = context.inherited_constraints.len(),
            examples = context.pattern_examples.len(),
            types = context.related_types.len(),
            deps = context.dependency_context.len(),
            "Constraint context built"
        );

        Ok(context)
    }

    /// Validate a set of constraints for internal consistency.
    ///
    /// Checks for:
    /// - Contradictory Must / MustNot pairs on the same pattern.
    /// - Empty descriptions or patterns.
    /// - Constraints referencing unknown files.
    pub async fn validate_constraints(
        &self,
        constraints: &[Constraint],
    ) -> Result<Vec<ValidationResult>> {
        info!(count = constraints.len(), "Validating constraints");

        let mut results = Vec::new();

        // Check for empty descriptions.
        for c in constraints {
            if c.description.trim().is_empty() {
                results.push(ValidationResult {
                    constraint_id: c.id.clone(),
                    passed: false,
                    message: "Constraint has an empty description".into(),
                    severity: Severity::Error,
                });
            }

            if c.pattern.trim().is_empty() {
                results.push(ValidationResult {
                    constraint_id: c.id.clone(),
                    passed: false,
                    message: "Constraint has an empty pattern".into(),
                    severity: Severity::Warning,
                });
            }
        }

        // Check for contradictory pairs.
        self.check_contradictions(constraints, &mut results);

        // Check that patterns reference known graph entities when applicable.
        self.check_pattern_references(constraints, &mut results)
            .await;

        // If no issues were found, emit one "all clear" result.
        if results.is_empty() {
            results.push(ValidationResult {
                constraint_id: "all".into(),
                passed: true,
                message: format!("All {} constraints are valid", constraints.len()),
                severity: Severity::Info,
            });
        }

        info!(
            total = constraints.len(),
            issues = results.iter().filter(|r| !r.passed).count(),
            "Constraint validation complete"
        );

        Ok(results)
    }

    /// Resolve concrete code examples from the knowledge graph for a given
    /// pattern name or description.
    ///
    /// Searches files in the graph whose functions, types, or imports match
    /// the pattern string and returns them ranked by relevance.
    pub async fn resolve_pattern_examples(
        &self,
        pattern_name: &str,
    ) -> Result<Vec<PatternExample>> {
        debug!(pattern = %pattern_name, "Resolving pattern examples");

        let pattern_lower = pattern_name.to_lowercase();
        let mut examples = Vec::new();

        // Search for types matching the pattern.
        let type_rows = self
            .kuzu
            .execute_query("DEFINES")
            .await
            .unwrap_or_default();

        for row in &type_rows {
            let type_name = row.get("type_name").cloned().unwrap_or_default();
            let file_path = row.get("file_path").cloned().unwrap_or_default();

            if type_name.to_lowercase().contains(&pattern_lower)
                || pattern_lower.contains(&type_name.to_lowercase())
            {
                examples.push(PatternExample {
                    file_path: file_path.clone(),
                    content: format!("type {} (defined in {})", type_name, file_path),
                    relevance: 0.8,
                });
            }
        }

        // Search for functions matching the pattern.
        let func_rows = self
            .kuzu
            .execute_query("CONTAINS")
            .await
            .unwrap_or_default();

        for row in &func_rows {
            let func_name = row.get("function_name").cloned().unwrap_or_default();
            let signature = row.get("signature").cloned().unwrap_or_default();
            let file_path = row.get("file_path").cloned().unwrap_or_default();

            if func_name.to_lowercase().contains(&pattern_lower)
                || signature.to_lowercase().contains(&pattern_lower)
            {
                examples.push(PatternExample {
                    file_path: file_path.clone(),
                    content: format!("fn {} in {}", signature, file_path),
                    relevance: 0.7,
                });
            }
        }

        // Enrich with SearchOrchestrator results if available.
        if let Some(ref search) = self.search {
            let request = SearchRequest {
                query: pattern_name.to_string(),
                max_results: 5,
                ..Default::default()
            };
            match tokio::time::timeout(
                std::time::Duration::from_secs(30),
                search.search(&request),
            ).await {
                Ok(Ok(response)) => {
                    for result in response.results {
                        examples.push(PatternExample {
                            file_path: result.file_path,
                            content: result.content,
                            relevance: result.relevance_score * 0.9,
                        });
                    }
                }
                Ok(Err(e)) => {
                    warn!(error = %e, "SearchOrchestrator enrichment failed, using graph-only results");
                }
                Err(_) => {
                    warn!("SearchOrchestrator enrichment timed out after 30s in Gandalf");
                }
            }
        }

        // Sort by relevance descending and deduplicate.
        examples.sort_by(|a, b| {
            b.relevance
                .partial_cmp(&a.relevance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        examples.dedup_by(|a, b| a.file_path == b.file_path && a.content == b.content);

        debug!(
            pattern = %pattern_name,
            count = examples.len(),
            "Pattern examples resolved"
        );

        Ok(examples)
    }

    /// Verify and augment tasks using LLM analysis before execution.
    ///
    /// Given a scroll's intent, constraints, and proposed tasks, asks the LLM to:
    /// 1. Verify the tasks are sufficient to fulfill the intent.
    /// 2. Identify missing files or dependencies.
    /// 3. Flag conflicts with constraints.
    ///
    /// Returns a `TaskAugmentation` with verification result, suggestions, and warnings.
    /// If no LLM is available, returns a default "verified" result.
    pub async fn verify_and_augment_tasks(
        &self,
        scroll: &Scroll,
        tasks: &[Task],
    ) -> Result<TaskAugmentation> {
        use crate::types::{AgentActivity, AgentActivityStatus};

        let llm = match &self.llm {
            Some(llm) if llm.is_available() => llm,
            _ => {
                debug!("No LLM available for task augmentation, skipping");
                return Ok(TaskAugmentation {
                    verified: true,
                    suggestions: vec![],
                    warnings: vec![],
                    groupings: vec![],
                });
            }
        };

        let activity_id = format!("act-gandalf-aug-{}-{}", scroll.id, chrono::Utc::now().timestamp_millis());
        let _ = self.sqlite.record_agent_activity(&AgentActivity {
            id: activity_id.clone(),
            scroll_id: scroll.id.clone(),
            task_id: String::new(),
            agent_name: "Gandalf".into(),
            action: "verify_and_augment".into(),
            detail: format!("{} tasks", tasks.len()),
            started_at: chrono::Utc::now(),
            completed_at: None,
            status: AgentActivityStatus::Running,
        }).await;

        let (system, prompt) = Self::build_augmentation_prompts(&scroll, &tasks);
        let system = system.as_str();

        match tokio::time::timeout(
            std::time::Duration::from_secs(120),
            llm.complete_with_system(system, &prompt),
        ).await {
            Err(_) => {
                warn!("Gandalf LLM augmentation timed out after 120s");
                let _ = self.sqlite.complete_agent_activity(
                    &activity_id,
                    &AgentActivityStatus::Failed,
                    "timed out after 120s",
                ).await;
                Ok(TaskAugmentation {
                    verified: true,
                    suggestions: vec![],
                    warnings: vec!["Gandalf augmentation timed out after 120s".to_string()],
                    groupings: vec![],
                })
            }
            Ok(Ok(response)) => {
                // Try to parse LLM response as JSON.
                let cleaned = response.trim().trim_start_matches("```json").trim_end_matches("```").trim();
                match serde_json::from_str::<TaskAugmentation>(cleaned) {
                    Ok(augmentation) => {
                        info!(
                            scroll_id = %scroll.id,
                            verified = augmentation.verified,
                            suggestions = augmentation.suggestions.len(),
                            warnings = augmentation.warnings.len(),
                            "Gandalf task augmentation complete"
                        );
                        let _ = self.sqlite.complete_agent_activity(
                            &activity_id,
                            &AgentActivityStatus::Completed,
                            &format!(
                                "verified={}, {} suggestions, {} warnings",
                                augmentation.verified,
                                augmentation.suggestions.len(),
                                augmentation.warnings.len(),
                            ),
                        ).await;
                        Ok(augmentation)
                    }
                    Err(e) => {
                        warn!(error = %e, "Failed to parse Gandalf augmentation response, treating as verified");
                        let _ = self.sqlite.complete_agent_activity(
                            &activity_id,
                            &AgentActivityStatus::Completed,
                            &format!("parse error (treated as verified): {}", e),
                        ).await;
                        Ok(TaskAugmentation {
                            verified: true,
                            suggestions: vec![],
                            warnings: vec![format!("Gandalf analysis returned unparseable response: {}", e)],
                            groupings: vec![],
                        })
                    }
                }
            }
            Ok(Err(e)) => {
                warn!(error = %e, "Gandalf LLM augmentation failed, continuing without");
                let _ = self.sqlite.complete_agent_activity(
                    &activity_id,
                    &AgentActivityStatus::Failed,
                    &format!("{}", e),
                ).await;
                Ok(TaskAugmentation {
                    verified: true,
                    suggestions: vec![],
                    warnings: vec![format!("Gandalf augmentation unavailable: {}", e)],
                    groupings: vec![],
                })
            }
        }
    }

    // -----------------------------------------------------------------------
    // Prefetch helpers
    // -----------------------------------------------------------------------

    /// Build the system and user prompts for Gandalf augmentation.
    ///
    /// Shared by both the direct-call (`verify_and_augment_tasks`) and
    /// prefetch (`prepare_augmentation_prompts`) paths to prevent drift.
    fn build_augmentation_prompts(scroll: &Scroll, tasks: &[Task]) -> (String, String) {
        let task_summary: Vec<String> = tasks
            .iter()
            .map(|t| {
                format!(
                    "- [{}] {} -> {} (has code: {})",
                    match t.task_type {
                        TaskType::Create => "Create",
                        TaskType::Modify => "Modify",
                        TaskType::Delete => "Delete",
                        TaskType::Test => "Test",
                    },
                    t.description,
                    t.target_file,
                    t.generated_code.is_some()
                )
            })
            .collect();

        let constraint_summary: Vec<String> = scroll
            .constraints
            .iter()
            .map(|c| format!("- [{:?}] {}", c.constraint_type, c.description))
            .collect();

        let user_prompt = format!(
            "## Scroll Intent\n{}\n\n## Constraints\n{}\n\n## Proposed Tasks\n{}\n\n\
             Analyze the above and respond with ONLY valid JSON (no markdown fences):\n\
             {{\"verified\": bool, \"suggestions\": [{{\"description\": \"...\", \"target_file\": \"...\", \"task_type\": \"Create|Modify|Test\", \"group\": \"optional-group-name\"}}], \"warnings\": [\"...\"], \"groupings\": [{{\"group_label\": \"feature-name\", \"task_ids\": [\"existing-task-id-1\"]}}]}}",
            scroll.intent,
            if constraint_summary.is_empty() {
                "None".to_string()
            } else {
                constraint_summary.join("\n")
            },
            task_summary.join("\n"),
        );

        let system_prompt = "You are Gandalf, a wise software architect and constraint guardian. \
                       Given a scroll's intent, constraints, and proposed tasks, verify: \
                       1. Are the tasks sufficient to fulfill the intent? \
                       2. Are any files or dependencies missing? \
                       3. Do any tasks conflict with the constraints? \
                       4. Group related tasks by returning a \"groupings\" array. Group files that form a single logical unit \
                       (e.g., a module + its tests, an interface + its implementation). Use descriptive group labels. \
                       IMPORTANT: Respond with your JSON analysis directly. Do not ask questions. \
                       Respond as JSON only.".to_string();

        (system_prompt, user_prompt)
    }

    /// Build augmentation prompts without calling the LLM.
    ///
    /// Returns a single `PrefetchPrompt` containing the same system + user
    /// prompt that `verify_and_augment_tasks` would send to the LLM.
    pub fn prepare_augmentation_prompts(
        &self,
        scroll: &Scroll,
        tasks: &[Task],
    ) -> Vec<crate::mcp::prefetch::PrefetchPrompt> {
        let (system_prompt, user_prompt) = Self::build_augmentation_prompts(scroll, tasks);

        vec![crate::mcp::prefetch::PrefetchPrompt {
            id: format!("gandalf-aug-{}", scroll.id),
            agent: "Gandalf".into(),
            task_id: None,
            system_prompt,
            user_prompt,
            description: format!("Verify and augment tasks for scroll {}", scroll.id),
        }]
    }

    /// Parse a raw LLM response into a `TaskAugmentation`.
    ///
    /// Shared by both the sampling path (`verify_and_augment_tasks`) and
    /// the prefetch path (process_round1).
    pub fn parse_augmentation_response(&self, response: &str) -> TaskAugmentation {
        let cleaned = response.trim().trim_start_matches("```json").trim_end_matches("```").trim();
        match serde_json::from_str::<TaskAugmentation>(cleaned) {
            Ok(augmentation) => {
                info!(
                    verified = augmentation.verified,
                    suggestions = augmentation.suggestions.len(),
                    warnings = augmentation.warnings.len(),
                    "Gandalf: parsed augmentation response"
                );
                augmentation
            }
            Err(e) => {
                warn!(error = %e, "Failed to parse Gandalf augmentation response, treating as verified");
                TaskAugmentation {
                    verified: true,
                    suggestions: vec![],
                    warnings: vec![format!("Gandalf analysis returned unparseable response: {}", e)],
                    groupings: vec![],
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Internal helpers
    // -----------------------------------------------------------------------

    /// Discover constraints implied by the graph neighbourhood of `file_path`.
    async fn discover_constraints_for_file(
        &self,
        file_path: &str,
    ) -> Result<Vec<Constraint>> {
        let mut constraints = Vec::new();

        // Look at the file entity in the graph.
        if let Some(file_entity) = self.kuzu.get_file(file_path).await? {
            // Infer constraints from public API surface.
            for func in &file_entity.functions {
                if func.is_public {
                    constraints.push(Constraint {
                        id: format!("discovered-pub-{}-{}", file_path, func.name),
                        constraint_type: ConstraintType::Should,
                        description: format!(
                            "Public function '{}' exists in '{}'; new code should maintain API compatibility",
                            func.name, file_path
                        ),
                        pattern: func.signature.clone(),
                        source: ConstraintSource::Discovered,
                    });
                }
            }

            // Infer constraints from type definitions.
            for type_def in &file_entity.types {
                constraints.push(Constraint {
                    id: format!("discovered-type-{}-{}", file_path, type_def.name),
                    constraint_type: ConstraintType::Should,
                    description: format!(
                        "Type '{}' ({}) is defined in '{}'; changes must be compatible",
                        type_def.name, type_def.kind, file_path
                    ),
                    pattern: type_def.name.clone(),
                    source: ConstraintSource::Discovered,
                });
            }
        }

        // Look at dependents -- files that import this file.
        let dependents = self.kuzu.get_dependents(file_path).await?;
        if !dependents.is_empty() {
            constraints.push(Constraint {
                id: format!("discovered-dependents-{}", file_path),
                constraint_type: ConstraintType::Must,
                description: format!(
                    "'{}' is imported by {} other file(s) ({}); breaking changes must be avoided",
                    file_path,
                    dependents.len(),
                    dependents.join(", ")
                ),
                pattern: "backward-compatible".into(),
                source: ConstraintSource::Inferred,
            });
        }

        Ok(constraints)
    }

    /// Collect constraints that are inherited through the dependency chain
    /// of `file_path` (i.e., constraints on the files that `file_path`
    /// depends on).
    async fn discover_inherited_constraints(
        &self,
        file_path: &str,
    ) -> Result<Vec<Constraint>> {
        let mut inherited = Vec::new();

        let deps = self.kuzu.get_dependencies(file_path).await?;
        for dep_path in &deps {
            let dep_constraints = self.discover_constraints_for_file(dep_path).await?;
            for c in dep_constraints {
                // Only inherit Must and MustNot -- advisory constraints are
                // not propagated.
                if c.constraint_type == ConstraintType::Must
                    || c.constraint_type == ConstraintType::MustNot
                {
                    inherited.push(Constraint {
                        id: format!("inherited-{}", c.id),
                        description: format!("[inherited from {}] {}", dep_path, c.description),
                        ..c
                    });
                }
            }
        }

        Ok(inherited)
    }

    /// Discover types defined in files that are related to `file_path`
    /// through the dependency graph.
    async fn discover_related_types(&self, file_path: &str) -> Result<Vec<RelatedType>> {
        let mut related = Vec::new();

        // Types from the file itself.
        if let Some(entity) = self.kuzu.get_file(file_path).await? {
            for td in &entity.types {
                related.push(RelatedType {
                    file_path: file_path.to_string(),
                    type_name: td.name.clone(),
                    kind: td.kind.clone(),
                    fields: td.fields.clone(),
                });
            }
        }

        // Types from direct dependencies.
        let deps = self.kuzu.get_dependencies(file_path).await?;
        for dep_path in &deps {
            if let Some(entity) = self.kuzu.get_file(dep_path).await? {
                for td in &entity.types {
                    related.push(RelatedType {
                        file_path: dep_path.clone(),
                        type_name: td.name.clone(),
                        kind: td.kind.clone(),
                        fields: td.fields.clone(),
                    });
                }
            }
        }

        Ok(related)
    }

    /// Build dependency edges relevant to `file_path`.
    async fn build_dependency_context(&self, file_path: &str) -> Result<Vec<DependencyInfo>> {
        let mut deps_info = Vec::new();

        // Outgoing dependencies.
        if let Some(entity) = self.kuzu.get_file(file_path).await? {
            for import in &entity.imports {
                deps_info.push(DependencyInfo {
                    from_file: file_path.to_string(),
                    to_file: import.source.clone(),
                    imported_items: import.items.clone(),
                });
            }
        }

        // Incoming dependencies (files that depend on the target).
        let dependents = self.kuzu.get_dependents(file_path).await?;
        for dep in &dependents {
            if let Some(entity) = self.kuzu.get_file(dep).await? {
                for import in &entity.imports {
                    if import.source == file_path {
                        deps_info.push(DependencyInfo {
                            from_file: dep.clone(),
                            to_file: file_path.to_string(),
                            imported_items: import.items.clone(),
                        });
                    }
                }
            }
        }

        Ok(deps_info)
    }

    /// Filter SCROLL-level constraints to those relevant to a specific file.
    fn filter_constraints_for_file(
        &self,
        constraints: &[Constraint],
        file_path: &str,
    ) -> Vec<Constraint> {
        constraints
            .iter()
            .filter(|c| {
                // Keep constraints that are generic (no file-specific pattern)
                // or that mention this file's name or extension.
                let pattern_lower = c.pattern.to_lowercase();
                let file_lower = file_path.to_lowercase();

                // Universal constraints always apply.
                if pattern_lower.is_empty() {
                    return true;
                }

                // Pattern mentions the file directly.
                if file_lower.contains(&pattern_lower) || pattern_lower.contains(&file_lower) {
                    return true;
                }

                // Pattern matches the file extension.
                if let Some(ext) = std::path::Path::new(file_path)
                    .extension()
                    .and_then(|e| e.to_str())
                {
                    if pattern_lower.contains(ext) {
                        return true;
                    }
                }

                // Broad architectural constraints always apply.
                matches!(
                    c.constraint_type,
                    ConstraintType::Must | ConstraintType::MustNot
                )
            })
            .cloned()
            .collect()
    }

    /// Check for contradictory Must / MustNot pairs.
    fn check_contradictions(
        &self,
        constraints: &[Constraint],
        results: &mut Vec<ValidationResult>,
    ) {
        let must_patterns: Vec<&Constraint> = constraints
            .iter()
            .filter(|c| c.constraint_type == ConstraintType::Must)
            .collect();

        let must_not_patterns: Vec<&Constraint> = constraints
            .iter()
            .filter(|c| c.constraint_type == ConstraintType::MustNot)
            .collect();

        for must in &must_patterns {
            for must_not in &must_not_patterns {
                if !must.pattern.is_empty()
                    && must.pattern.to_lowercase() == must_not.pattern.to_lowercase()
                {
                    results.push(ValidationResult {
                        constraint_id: format!("{} vs {}", must.id, must_not.id),
                        passed: false,
                        message: format!(
                            "Contradictory constraints: '{}' is both Must and MustNot for pattern '{}'",
                            must.description, must.pattern
                        ),
                        severity: Severity::Error,
                    });
                }
            }
        }
    }

    /// Check that constraint patterns reference entities that exist in the
    /// knowledge graph.
    async fn check_pattern_references(
        &self,
        constraints: &[Constraint],
        results: &mut Vec<ValidationResult>,
    ) {
        for c in constraints {
            if c.pattern.trim().is_empty() {
                continue;
            }

            // If the pattern looks like a file path, verify it exists in the graph.
            if c.pattern.contains('/') || c.pattern.contains('.') {
                if let Ok(file) = self.kuzu.get_file(&c.pattern).await {
                    if file.is_none() {
                        results.push(ValidationResult {
                            constraint_id: c.id.clone(),
                            passed: true, // warning, not failure
                            message: format!(
                                "Pattern '{}' looks like a file path but is not in the knowledge graph (it may not have been indexed yet)",
                                c.pattern
                            ),
                            severity: Severity::Warning,
                        });
                    }
                }
            }
        }
    }

    /// Enrich constraints from knowledge base sources.
    ///
    /// Queries all registered KBs with the task description, converts
    /// high-relevance entries (>0.6) to `Should` constraints with
    /// `ConstraintSource::Discovered`.
    async fn enrich_from_kb(&self, task: &Task) -> Vec<Constraint> {
        let kb = match &self.kb_manager {
            Some(kb) => kb,
            None => return Vec::new(),
        };

        let has_sources = tokio::time::timeout(
            std::time::Duration::from_secs(5),
            kb.has_sources(),
        ).await.unwrap_or(false);
        if !has_sources {
            return Vec::new();
        }

        let results = match tokio::time::timeout(
            std::time::Duration::from_secs(15),
            kb.query_all_kbs(&task.description, 5),
        ).await {
            Ok(r) => r,
            Err(_) => {
                warn!("KB query timed out after 15s in Gandalf enrich_from_kb");
                return Vec::new();
            }
        };
        let mut constraints = Vec::new();

        for qr in &results {
            for entry in &qr.entries {
                if entry.relevance_score > 0.6 {
                    let constraint_id = format!("kb-{}-{}", entry.source, entry.id);
                    constraints.push(Constraint {
                        id: constraint_id,
                        constraint_type: ConstraintType::Should,
                        description: format!(
                            "[KB:{}] {}",
                            entry.source,
                            if entry.title.is_empty() {
                                entry.content.chars().take(120).collect::<String>()
                            } else {
                                entry.title.clone()
                            }
                        ),
                        pattern: String::new(),
                        source: ConstraintSource::Discovered,
                    });
                }
            }
        }

        if !constraints.is_empty() {
            debug!(
                task_id = %task.id,
                kb_constraints = constraints.len(),
                "Enriched constraints from KB"
            );
        }

        constraints
    }

    /// Build a human-readable summary suitable for inclusion in LLM prompts.
    fn build_summary(
        &self,
        task: &Task,
        direct: &[Constraint],
        inherited: &[Constraint],
        types: &[RelatedType],
        deps: &[DependencyInfo],
    ) -> String {
        let mut parts = Vec::new();

        parts.push(format!(
            "Task: {} ({})\nTarget: {}\nType: {:?}",
            task.description, task.id, task.target_file, task.task_type
        ));

        // Direct constraints.
        if !direct.is_empty() {
            let mut section = String::from("\n--- Direct Constraints ---");
            for c in direct {
                section.push_str(&format!(
                    "\n[{:?}] {}: {}",
                    c.constraint_type, c.id, c.description
                ));
            }
            parts.push(section);
        }

        // Inherited constraints.
        if !inherited.is_empty() {
            let mut section = String::from("\n--- Inherited Constraints ---");
            for c in inherited {
                section.push_str(&format!(
                    "\n[{:?}] {}: {}",
                    c.constraint_type, c.id, c.description
                ));
            }
            parts.push(section);
        }

        // Related types.
        if !types.is_empty() {
            let mut section = String::from("\n--- Related Types ---");
            for t in types {
                let fields_str = if t.fields.is_empty() {
                    String::new()
                } else {
                    format!(" {{ {} }}", t.fields.join(", "))
                };
                section.push_str(&format!(
                    "\n{} {} ({}){}  [{}]",
                    t.kind, t.type_name, t.file_path, fields_str, t.file_path
                ));
            }
            parts.push(section);
        }

        // Dependency context.
        if !deps.is_empty() {
            let mut section = String::from("\n--- Dependencies ---");
            for d in deps {
                let items = if d.imported_items.is_empty() {
                    "*".to_string()
                } else {
                    d.imported_items.join(", ")
                };
                section.push_str(&format!(
                    "\n{} -> {} (imports: {})",
                    d.from_file, d.to_file, items
                ));
            }
            parts.push(section);
        }

        parts.join("\n")
    }
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use crate::types::*;
    use tempfile::TempDir;

    /// Helper: stand up an in-memory SqliteDb + KuzuGraph.
    async fn setup() -> (Arc<SqliteDb>, Arc<KuzuGraph>, TempDir) {
        let tmp = TempDir::new().expect("failed to create temp dir");
        let db_path = tmp.path().join("test.db");
        let graph_path = tmp.path().join("graph");

        let sqlite = Arc::new(SqliteDb::new(&db_path).expect("sqlite init"));
        let kuzu = Arc::new(KuzuGraph::new(&graph_path).expect("kuzu init"));

        (sqlite, kuzu, tmp)
    }

    /// Helper: create a minimal SCROLL with one task in SQLite.
    async fn seed_scroll(sqlite: &SqliteDb) -> (Scroll, Task) {
        let now = Utc::now();
        let scroll = Scroll {
            id: "SCROLL-test-001".into(),
            title: "Test Scroll".into(),
            intent: "Add a new handler".into(),
            constraints: vec![
                Constraint {
                    id: "c-1".into(),
                    constraint_type: ConstraintType::Must,
                    description: "Use async/await".into(),
                    pattern: "async".into(),
                    source: ConstraintSource::UserSpecified,
                },
                Constraint {
                    id: "c-2".into(),
                    constraint_type: ConstraintType::MustNot,
                    description: "No unwrap in production code".into(),
                    pattern: "unwrap".into(),
                    source: ConstraintSource::UserSpecified,
                },
            ],
            tasks: vec![],
            metadata: ScrollMetadata {
                execution_mode: ExecutionMode::Forge,
                auto_forge_config: None,
                recipe_used: None,
                search_context: vec![],
            },
            status: ScrollStatus::Draft,
            created_at: now,
            updated_at: now,
            cleanup_after: None,
        };
        sqlite.create_scroll(&scroll).await.expect("create scroll");

        let task = Task {
            id: "task-001".into(),
            scroll_id: scroll.id.clone(),
            description: "Create handler module".into(),
            target_file: "src/handlers/auth.rs".into(),
            task_type: TaskType::Create,
            status: TaskStatus::Pending,
            dependencies: vec![],
            generated_code: None,
            diff: None,
            created_at: now,
            group_label: None,
        };
        sqlite.create_task(&task).await.expect("create task");

        (scroll, task)
    }

    /// Helper: populate the graph with a file entity.
    async fn seed_graph(kuzu: &KuzuGraph) {
        let entity = FileEntity {
            path: "src/handlers/auth.rs".into(),
            language: "rust".into(),
            functions: vec![
                FunctionDef {
                    name: "authenticate".into(),
                    signature: "pub async fn authenticate(req: Request) -> Result<Response>".into(),
                    line_start: 10,
                    line_end: 30,
                    is_public: true,
                },
            ],
            types: vec![
                TypeDef {
                    name: "AuthConfig".into(),
                    kind: "struct".into(),
                    fields: vec!["secret: String".into(), "ttl: u64".into()],
                },
            ],
            imports: vec![
                ImportDef {
                    source: "src/types.rs".into(),
                    items: vec!["Request".into(), "Response".into()],
                },
            ],
            exports: vec!["authenticate".into()],
        };
        kuzu.add_file_to_graph(&entity).await.expect("add file");

        // Also add the dependency target.
        let dep_entity = FileEntity {
            path: "src/types.rs".into(),
            language: "rust".into(),
            functions: vec![],
            types: vec![
                TypeDef {
                    name: "Request".into(),
                    kind: "struct".into(),
                    fields: vec!["headers: HashMap<String, String>".into()],
                },
                TypeDef {
                    name: "Response".into(),
                    kind: "struct".into(),
                    fields: vec!["status: u16".into(), "body: String".into()],
                },
            ],
            imports: vec![],
            exports: vec!["Request".into(), "Response".into()],
        };
        kuzu.add_file_to_graph(&dep_entity)
            .await
            .expect("add dep file");
    }

    #[tokio::test]
    async fn test_extract_constraints_includes_user_specified() {
        let (sqlite, kuzu, _tmp) = setup().await;
        let (scroll, _task) = seed_scroll(&sqlite).await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let constraints = gandalf.extract_constraints(&scroll.id).await.unwrap();

        // Should contain at least the two user-specified constraints.
        assert!(constraints.len() >= 2, "Expected >= 2, got {}", constraints.len());
        assert!(constraints.iter().any(|c| c.id == "c-1"));
        assert!(constraints.iter().any(|c| c.id == "c-2"));
    }

    #[tokio::test]
    async fn test_extract_constraints_discovers_from_graph() {
        let (sqlite, kuzu, _tmp) = setup().await;
        let (_scroll, _task) = seed_scroll(&sqlite).await;
        seed_graph(&kuzu).await;

        let gandalf = Gandalf::new(Arc::clone(&kuzu), Arc::clone(&sqlite));
        let constraints = gandalf.extract_constraints("SCROLL-test-001").await.unwrap();

        // Should discover the public function constraint.
        assert!(
            constraints.iter().any(|c| c.description.contains("authenticate")),
            "Expected discovered constraint for public fn 'authenticate'"
        );
    }

    #[tokio::test]
    async fn test_build_constraint_context() {
        let (sqlite, kuzu, _tmp) = setup().await;
        let (_scroll, task) = seed_scroll(&sqlite).await;
        seed_graph(&kuzu).await;

        let gandalf = Gandalf::new(Arc::clone(&kuzu), Arc::clone(&sqlite));
        let ctx = gandalf.build_constraint_context(&task).await.unwrap();

        // Direct constraints should be non-empty (from the SCROLL).
        assert!(!ctx.direct_constraints.is_empty());

        // Related types should include AuthConfig.
        assert!(
            ctx.related_types.iter().any(|t| t.type_name == "AuthConfig"),
            "Expected AuthConfig in related types"
        );

        // Dependency context should include the import of types.rs.
        assert!(
            ctx.dependency_context.iter().any(|d| d.to_file == "src/types.rs"),
            "Expected dependency on src/types.rs"
        );

        // Summary should be non-empty.
        assert!(!ctx.summary.is_empty());
    }

    #[tokio::test]
    async fn test_validate_constraints_clean() {
        let (sqlite, kuzu, _tmp) = setup().await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let constraints = vec![
            Constraint {
                id: "c-ok".into(),
                constraint_type: ConstraintType::Must,
                description: "Use error handling".into(),
                pattern: "Result".into(),
                source: ConstraintSource::UserSpecified,
            },
        ];

        let results = gandalf.validate_constraints(&constraints).await.unwrap();
        assert!(results.iter().all(|r| r.passed), "Expected all validations to pass");
    }

    #[tokio::test]
    async fn test_validate_constraints_empty_description() {
        let (sqlite, kuzu, _tmp) = setup().await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let constraints = vec![
            Constraint {
                id: "c-bad".into(),
                constraint_type: ConstraintType::Must,
                description: "".into(),
                pattern: "something".into(),
                source: ConstraintSource::UserSpecified,
            },
        ];

        let results = gandalf.validate_constraints(&constraints).await.unwrap();
        assert!(
            results.iter().any(|r| !r.passed && r.message.contains("empty description")),
            "Expected error for empty description"
        );
    }

    #[tokio::test]
    async fn test_validate_constraints_contradiction() {
        let (sqlite, kuzu, _tmp) = setup().await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let constraints = vec![
            Constraint {
                id: "c-must".into(),
                constraint_type: ConstraintType::Must,
                description: "Use async".into(),
                pattern: "async".into(),
                source: ConstraintSource::UserSpecified,
            },
            Constraint {
                id: "c-must-not".into(),
                constraint_type: ConstraintType::MustNot,
                description: "No async".into(),
                pattern: "async".into(),
                source: ConstraintSource::UserSpecified,
            },
        ];

        let results = gandalf.validate_constraints(&constraints).await.unwrap();
        assert!(
            results.iter().any(|r| !r.passed && r.message.contains("Contradictory")),
            "Expected contradiction detected"
        );
    }

    #[tokio::test]
    async fn test_resolve_pattern_examples() {
        let (sqlite, kuzu, _tmp) = setup().await;
        seed_graph(&kuzu).await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let examples = gandalf.resolve_pattern_examples("authenticate").await.unwrap();

        assert!(!examples.is_empty(), "Expected at least one example for 'authenticate'");
        assert!(
            examples.iter().any(|e| e.content.contains("authenticate")),
            "Expected example content to reference 'authenticate'"
        );
    }

    #[tokio::test]
    async fn test_resolve_pattern_examples_type() {
        let (sqlite, kuzu, _tmp) = setup().await;
        seed_graph(&kuzu).await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let examples = gandalf.resolve_pattern_examples("AuthConfig").await.unwrap();

        assert!(!examples.is_empty(), "Expected examples for type 'AuthConfig'");
    }

    #[tokio::test]
    async fn test_resolve_pattern_examples_no_match() {
        let (sqlite, kuzu, _tmp) = setup().await;

        let gandalf = Gandalf::new(kuzu, sqlite);
        let examples = gandalf
            .resolve_pattern_examples("nonexistent_xyz_pattern")
            .await
            .unwrap();

        assert!(examples.is_empty(), "Expected no examples for unknown pattern");
    }

    #[tokio::test]
    async fn test_inherited_constraints_propagation() {
        let (sqlite, kuzu, _tmp) = setup().await;
        seed_graph(&kuzu).await;

        let gandalf = Gandalf::new(Arc::clone(&kuzu), Arc::clone(&sqlite));

        // src/handlers/auth.rs depends on src/types.rs, so constraints from
        // types.rs should propagate.
        let inherited = gandalf
            .discover_inherited_constraints("src/handlers/auth.rs")
            .await
            .unwrap();

        // types.rs has public types (Request, Response) which generate
        // discovered constraints that are Must-level (due to dependents).
        // These may or may not exist depending on whether types.rs has
        // dependents. In this seed data, auth.rs imports from types.rs so
        // types.rs has the auth.rs dependent.
        //
        // The inherited constraints come from discover_constraints_for_file
        // on types.rs. types.rs has types but no public functions, and auth.rs
        // is a dependent, so there should be a "backward-compatible" constraint.
        // However, only Must/MustNot are inherited.
        // The dependents constraint is Must, so it should be inherited.
        // Note: we check that it doesn't panic and returns some reasonable result.
        // The exact count depends on internal logic.
        debug!("Inherited constraints: {:?}", inherited);
    }

    #[tokio::test]
    async fn test_prepare_augmentation_prompts_returns_valid_prompt() {
        let (sqlite, kuzu, _tmp) = setup().await;
        let (scroll, task) = seed_scroll(&sqlite).await;

        let gandalf = Gandalf::new(
            Arc::clone(&kuzu),
            Arc::clone(&sqlite),
        );

        let prompts = gandalf.prepare_augmentation_prompts(&scroll, &[task]);

        assert_eq!(prompts.len(), 1);
        let p = &prompts[0];
        assert_eq!(p.agent, "Gandalf");
        assert!(p.task_id.is_none());
        assert!(p.id.contains(&scroll.id));
        assert!(p.system_prompt.contains("Gandalf"));
        assert!(p.user_prompt.contains("Scroll Intent"));
        assert!(p.user_prompt.contains(&scroll.intent));
        assert!(p.description.contains(&scroll.id));
    }

    #[test]
    fn test_parse_augmentation_response_valid_json() {
        let tmp = TempDir::new().unwrap();
        let db_path = tmp.path().join("test.db");
        let sqlite = Arc::new(SqliteDb::new(&db_path).expect("sqlite"));
        let graph_path = tmp.path().join("graph");
        let kuzu = Arc::new(KuzuGraph::new(&graph_path).expect("kuzu"));

        let gandalf = Gandalf::new(kuzu, sqlite);
        let response = r#"{"verified": true, "suggestions": [{"description": "Add tests", "target_file": "src/test.rs", "task_type": "Test"}], "warnings": ["Consider error handling"]}"#;

        let aug = gandalf.parse_augmentation_response(response);
        assert!(aug.verified);
        assert_eq!(aug.suggestions.len(), 1);
        assert_eq!(aug.suggestions[0].target_file, "src/test.rs");
        assert_eq!(aug.warnings.len(), 1);
    }

    #[test]
    fn test_parse_augmentation_response_invalid_json() {
        let tmp = TempDir::new().unwrap();
        let db_path = tmp.path().join("test.db");
        let sqlite = Arc::new(SqliteDb::new(&db_path).expect("sqlite"));
        let graph_path = tmp.path().join("graph");
        let kuzu = Arc::new(KuzuGraph::new(&graph_path).expect("kuzu"));

        let gandalf = Gandalf::new(kuzu, sqlite);
        let response = "not valid json at all";

        let aug = gandalf.parse_augmentation_response(response);
        assert!(aug.verified); // Falls back to verified
        assert!(aug.suggestions.is_empty());
        assert!(!aug.warnings.is_empty()); // Should have a warning about parse failure
    }

    #[test]
    fn test_parse_augmentation_response_with_groupings() {
        let tmp = TempDir::new().unwrap();
        let db_path = tmp.path().join("test.db");
        let sqlite = Arc::new(SqliteDb::new(&db_path).expect("sqlite"));
        let graph_path = tmp.path().join("graph");
        let kuzu = Arc::new(KuzuGraph::new(&graph_path).expect("kuzu"));

        let gandalf = Gandalf::new(kuzu, sqlite);
        let response = r#"{"verified": true, "suggestions": [{"description": "Add tests", "target_file": "src/test.rs", "task_type": "Test", "group": "auth-module"}], "warnings": [], "groupings": [{"group_label": "auth-module", "task_ids": ["TASK-001", "TASK-002"]}]}"#;

        let aug = gandalf.parse_augmentation_response(response);
        assert!(aug.verified);
        assert_eq!(aug.suggestions.len(), 1);
        assert_eq!(aug.suggestions[0].group, Some("auth-module".to_string()));
        assert_eq!(aug.groupings.len(), 1);
        assert_eq!(aug.groupings[0].group_label, "auth-module");
        assert_eq!(aug.groupings[0].task_ids, vec!["TASK-001", "TASK-002"]);
    }

    #[test]
    fn test_parse_augmentation_response_backwards_compat_no_groupings() {
        let tmp = TempDir::new().unwrap();
        let db_path = tmp.path().join("test.db");
        let sqlite = Arc::new(SqliteDb::new(&db_path).expect("sqlite"));
        let graph_path = tmp.path().join("graph");
        let kuzu = Arc::new(KuzuGraph::new(&graph_path).expect("kuzu"));

        let gandalf = Gandalf::new(kuzu, sqlite);
        // Old-format response without groupings or group fields — must still parse.
        let response = r#"{"verified": true, "suggestions": [{"description": "Add tests", "target_file": "src/test.rs", "task_type": "Test"}], "warnings": ["check imports"]}"#;

        let aug = gandalf.parse_augmentation_response(response);
        assert!(aug.verified);
        assert_eq!(aug.suggestions.len(), 1);
        assert_eq!(aug.suggestions[0].group, None);
        assert!(aug.groupings.is_empty());
        assert_eq!(aug.warnings.len(), 1);
    }
}
