"""
PlannerAgent - AI-powered workflow planning using Claude.

This agent uses Claude to intelligently decompose user requests into
executable workflow steps, identifying the right agents and capabilities
to accomplish complex data engineering tasks.
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Optional
import uuid

try:
    from anthropic import Anthropic
    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False

try:
    from src.agents.pydantic_agents import (
        PYDANTIC_AI_AVAILABLE,
        run_planner_agent,
        PlanOutput,
    )
except ImportError:
    PYDANTIC_AI_AVAILABLE = False
    run_planner_agent = None  # type: ignore[assignment]
    PlanOutput = None  # type: ignore[misc,assignment]

logger = logging.getLogger(__name__)


class PlannerCapability(str, Enum):
    """Planner agent capabilities."""

    PLAN_WORKFLOW = "plan_workflow"
    ANALYZE_REQUEST = "analyze_request"
    SUGGEST_AGENTS = "suggest_agents"
    OPTIMIZE_PLAN = "optimize_plan"
    EXPLAIN_PLAN = "explain_plan"
    BUILD_DIMENSIONS = "build_dimensions"


class PlannerState(str, Enum):
    """Planner execution states."""

    IDLE = "idle"
    PLANNING = "planning"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class AgentInfo:
    """Information about an available agent."""

    name: str
    agent_type: str
    description: str
    capabilities: list[str]
    input_schema: dict[str, Any] = field(default_factory=dict)
    output_schema: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "agent_type": self.agent_type,
            "description": self.description,
            "capabilities": self.capabilities,
            "input_schema": self.input_schema,
            "output_schema": self.output_schema,
        }


@dataclass
class PlannedStep:
    """A step in the planned workflow."""

    id: str
    name: str
    description: str
    agent_type: str
    capability: str
    input_mapping: dict[str, str] = field(default_factory=dict)
    output_key: str = ""
    depends_on: list[str] = field(default_factory=list)
    config: dict[str, Any] = field(default_factory=dict)
    reasoning: str = ""  # Why this step is needed

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "agent_type": self.agent_type,
            "capability": self.capability,
            "input_mapping": self.input_mapping,
            "output_key": self.output_key,
            "depends_on": self.depends_on,
            "config": self.config,
            "reasoning": self.reasoning,
        }


@dataclass
class WorkflowPlan:
    """A complete workflow plan."""

    id: str
    name: str
    description: str
    user_request: str
    steps: list[PlannedStep] = field(default_factory=list)
    estimated_duration: str = ""
    confidence_score: float = 0.0
    warnings: list[str] = field(default_factory=list)
    alternatives: list[str] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "user_request": self.user_request,
            "steps": [s.to_dict() for s in self.steps],
            "step_count": len(self.steps),
            "estimated_duration": self.estimated_duration,
            "confidence_score": self.confidence_score,
            "warnings": self.warnings,
            "alternatives": self.alternatives,
            "created_at": self.created_at.isoformat(),
        }


@dataclass
class PlannerConfig:
    """Configuration for the PlannerAgent."""

    model: str = "claude-sonnet-4-20250514"
    max_tokens: int = 4096
    temperature: float = 0.3  # Lower for more consistent planning
    max_steps: int = 10
    enable_parallel_steps: bool = True
    include_reasoning: bool = True
    api_key: Optional[str] = None
    fallback_on_llm_error: bool = True


# Default available agents in the DataBridge system
DEFAULT_AGENTS: list[AgentInfo] = [
    AgentInfo(
        name="schema_scanner",
        agent_type="SchemaScanner",
        description="Scans database schemas to discover tables, columns, keys, and data types. Profiles data quality and infers table types (fact/dimension/staging).",
        capabilities=["scan_schema", "extract_metadata", "detect_keys", "sample_profiles"],
        input_schema={"connection_string": "str", "schema_name": "str", "sample_size": "int"},
        output_schema={"tables": "list", "columns": "list", "keys": "list", "profiles": "dict"},
    ),
    AgentInfo(
        name="logic_extractor",
        agent_type="LogicExtractor",
        description="Parses SQL statements to extract CASE logic, calculations, aggregations, and business rules for hierarchy creation.",
        capabilities=["parse_sql", "extract_case", "identify_calcs", "detect_aggregations"],
        input_schema={"sql": "str", "sql_file_path": "str"},
        output_schema={"case_statements": "list", "calculations": "list", "aggregations": "list"},
    ),
    AgentInfo(
        name="warehouse_architect",
        agent_type="WarehouseArchitect",
        description="Designs star schema structures, generates dimension and fact table specifications, and creates dbt models.",
        capabilities=["design_star_schema", "generate_dims", "generate_facts", "dbt_models"],
        input_schema={"tables": "list", "business_rules": "dict"},
        output_schema={"star_schema": "dict", "dimensions": "list", "facts": "list", "dbt_models": "list"},
    ),
    AgentInfo(
        name="deploy_validator",
        agent_type="DeployValidator",
        description="Executes DDL statements, runs dbt commands, and validates deployments with row counts and aggregate comparisons.",
        capabilities=["execute_ddl", "run_dbt", "validate_counts", "compare_aggregates"],
        input_schema={"ddl_statements": "list", "dbt_project": "str"},
        output_schema={"execution_results": "list", "validation_results": "dict"},
    ),
    AgentInfo(
        name="hierarchy_builder",
        agent_type="HierarchyBuilder",
        description="Creates and manages hierarchical data structures using MCP tools. Supports flexible import, property management, and deployment.",
        capabilities=["create_hierarchy", "import_hierarchy", "add_properties", "deploy_hierarchy"],
        input_schema={"project_id": "str", "hierarchy_data": "dict", "properties": "dict"},
        output_schema={"hierarchy_id": "str", "mappings": "list", "deployment_status": "str"},
    ),
    AgentInfo(
        name="data_reconciler",
        agent_type="DataReconciler",
        description="Compares data between sources, identifies discrepancies, performs fuzzy matching, and generates reconciliation reports.",
        capabilities=["compare_sources", "fuzzy_match", "identify_orphans", "generate_report"],
        input_schema={"source_a": "dict", "source_b": "dict", "match_columns": "list"},
        output_schema={"matches": "list", "orphans": "list", "conflicts": "list", "report": "dict"},
    ),
    # Unified Agent agents for Book ↔ Librarian ↔ Researcher operations
    AgentInfo(
        name="book_manipulator",
        agent_type="BookManipulator",
        description="Manipulates Book hierarchies in-memory. Creates Books from CSV/JSON, adds/removes nodes, applies formulas, and exports to various formats.",
        capabilities=["create_book", "add_node", "remove_node", "apply_formula", "export_book"],
        input_schema={"book_name": "str", "node_data": "dict", "formula": "dict"},
        output_schema={"book": "dict", "node_count": "int", "export_path": "str"},
    ),
    AgentInfo(
        name="librarian_sync",
        agent_type="LibrarianSync",
        description="Synchronizes between Book (Python) and Librarian (NestJS). Checks out projects to Books, promotes Books to Librarian, and handles bidirectional sync.",
        capabilities=["checkout_project", "promote_book", "sync_changes", "diff_systems"],
        input_schema={"book_name": "str", "project_id": "str", "direction": "str"},
        output_schema={"sync_result": "dict", "created": "int", "updated": "int", "diff": "dict"},
    ),
    AgentInfo(
        name="researcher_analyst",
        agent_type="ResearcherAnalyst",
        description="Runs analytics on Book/Librarian data using Researcher API. Validates source mappings, compares hierarchy data with databases, profiles source columns.",
        capabilities=["validate_mappings", "compare_schemas", "profile_sources", "analyze_data"],
        input_schema={"book_name": "str", "connection_id": "str", "analysis_type": "str"},
        output_schema={"validation": "dict", "profile": "dict", "comparison": "dict"},
    ),
    # Cortex Agent for Snowflake Cortex AI integration
    AgentInfo(
        name="cortex_agent",
        agent_type="CortexAgent",
        description="Executes Snowflake Cortex AI functions with orchestrated reasoning. Handles text generation, summarization, sentiment analysis, and multi-step data cleaning tasks.",
        capabilities=["complete", "summarize", "sentiment", "translate", "extract_answer", "reason", "analyze_data", "clean_data"],
        input_schema={"goal": "str", "table_name": "str", "context": "dict"},
        output_schema={"result": "str", "thinking_steps": "list", "conversation_id": "str"},
    ),
]


class PlannerAgent:
    """
    AI-powered workflow planner using Claude.

    The PlannerAgent analyzes user requests and creates executable workflow
    plans by selecting appropriate agents and ordering steps based on
    dependencies.

    Example:
        planner = PlannerAgent()

        # Plan a workflow
        plan = planner.plan_workflow(
            "Extract hierarchies from SQL CASE statements and deploy to Snowflake"
        )

        # Get workflow definition for orchestrator
        workflow_def = planner.to_workflow_definition(plan)
    """

    def __init__(
        self,
        config: Optional[PlannerConfig] = None,
        available_agents: Optional[list[AgentInfo]] = None,
        hierarchy_service: Optional[Any] = None,
        tool_registry: Optional[Any] = None,
        purpose_path: Optional[str] = None,
    ):
        """
        Initialize the PlannerAgent.

        Args:
            config: Planner configuration
            available_agents: List of available agents (uses defaults if not provided)
            hierarchy_service: Optional HierarchyService for context injection
            tool_registry: Optional DynamicToolRegistry for domain awareness
            purpose_path: Optional path to purpose.md strategic vision document
        """
        self._id = str(uuid.uuid4())[:8]
        self._config = config or PlannerConfig()
        self._state = PlannerState.IDLE
        self._available_agents = list(available_agents or DEFAULT_AGENTS)
        self._plans: list[WorkflowPlan] = []
        self._created_at = datetime.now()
        self._hierarchy_service = hierarchy_service
        self._tool_registry = tool_registry
        self._purpose_text = self._load_purpose(purpose_path)

        # Auto-register extended agent types
        self._load_skill_agents()
        self._load_langgraph_agents()
        self._load_domain_agents()

        # Initialize Anthropic client
        self._client: Optional[Anthropic] = None
        if ANTHROPIC_AVAILABLE:
            api_key = self._config.api_key or os.getenv("ANTHROPIC_API_KEY")
            if api_key:
                self._client = Anthropic(api_key=api_key)

    @property
    def id(self) -> str:
        """Get planner ID."""
        return self._id

    @property
    def state(self) -> PlannerState:
        """Get current state."""
        return self._state

    @property
    def available_agents(self) -> list[AgentInfo]:
        """Get available agents."""
        return self._available_agents

    def _load_skill_agents(self) -> None:
        """Auto-register skill-based agents from skills/index.json."""
        # Auto-discover skills/index.json relative to project root
        candidates = [
            Path(__file__).parent.parent.parent / "skills" / "index.json",
            Path(__file__).parent.parent / "skills" / "index.json",
        ]
        index_path = None
        for c in candidates:
            if c.exists():
                index_path = c
                break
        if not index_path:
            return

        try:
            data = json.loads(index_path.read_text(encoding="utf-8"))
        except Exception:
            return

        skills_dir = index_path.parent
        for skill in data.get("skills", []):
            skill_id = skill.get("id", "")
            if not skill_id:
                continue

            agent_name = f"skill_{skill_id.replace('-', '_')}"

            # Check if prompt file exists (non-draft)
            prompt_file = skill.get("prompt_file", "")
            is_draft = not (skills_dir / prompt_file).exists() if prompt_file else True

            capabilities = skill.get("capabilities", [])
            description = skill.get("description", skill.get("name", skill_id))
            if is_draft:
                description = f"[DRAFT] {description}"

            self._available_agents.append(AgentInfo(
                name=agent_name,
                agent_type="SkillAgent",
                description=description,
                capabilities=capabilities,
                input_schema={"skill_id": "str", "context": "dict"},
                output_schema={"result": "Any", "artifacts": "list"},
            ))

    def _load_langgraph_agents(self) -> None:
        """Auto-register LangGraph specialist agents."""
        langgraph_specs = [
            {
                "name": "langgraph_data_modeling",
                "description": "Schema discovery, FK detection, column classification, dimension/fact detection. Covers load_metadata, relationship_discovery, ai_classify, detect_dimensions, artifact_bundle phases.",
                "capabilities": ["load_metadata", "relationship_discovery", "ai_classify", "detect_dimensions", "artifact_bundle"],
            },
            {
                "name": "langgraph_hierarchy",
                "description": "Hierarchy creation and table grouping from classified schema. Covers the hierarchy phase.",
                "capabilities": ["hierarchy_creation", "table_grouping", "level_assignment"],
            },
            {
                "name": "langgraph_dbt",
                "description": "dbt project scaffolding, model generation, and source YAML creation. Covers the dbt phase.",
                "capabilities": ["dbt_scaffold", "model_generation", "source_yaml"],
            },
            {
                "name": "langgraph_snowflake",
                "description": "DDL execution and Wright pipeline deployment to Snowflake. Covers the wright phase.",
                "capabilities": ["ddl_execution", "wright_pipeline", "dynamic_tables"],
            },
            {
                "name": "langgraph_qa",
                "description": "Expectation suites, quality validation, observability baselines. Covers quality and observability phases.",
                "capabilities": ["expectation_suites", "quality_validation", "obs_baselines", "data_quality"],
            },
            {
                "name": "langgraph_cortex_llm",
                "description": "Text generation, summarization, and sentiment analysis via Snowflake Cortex LLM functions.",
                "capabilities": ["text_generation", "summarization", "sentiment_analysis", "translation"],
            },
            {
                "name": "langgraph_cortex_analyst",
                "description": "Semantic model queries and natural language SQL via Cortex Analyst.",
                "capabilities": ["semantic_model_query", "natural_language_sql", "metric_calculation"],
            },
            {
                "name": "langgraph_cortex_search",
                "description": "Cortex search service and RAG retrieval for document/knowledge search.",
                "capabilities": ["cortex_search", "rag_retrieval", "document_search"],
            },
        ]

        for spec in langgraph_specs:
            self._available_agents.append(AgentInfo(
                name=spec["name"],
                agent_type="LangGraphAgent",
                description=spec["description"],
                capabilities=spec["capabilities"],
                input_schema={"goal": "str", "context": "dict"},
                output_schema={"result": "Any", "artifacts": "list"},
            ))

    def _load_domain_agents(self) -> None:
        """Auto-register domain mini-agents from the tool registry."""
        if not self._tool_registry:
            return
        try:
            domains = self._tool_registry.get_all_domains()
        except Exception:
            return

        # Domains already covered by explicit agents — skip to avoid circularity
        skip_domains = {"planner", "gateway", "core"}
        # Map domains to existing agent types to avoid redundancy
        covered_by_langgraph = {"langgraph"}

        for d in domains:
            domain_name = d.get("domain", "")
            if not domain_name or domain_name in skip_domains or domain_name in covered_by_langgraph:
                continue

            # Get top 5 tool names for this domain
            try:
                tools = self._tool_registry.get_tools_by_domain(domain_name)
                top_tools = [t["name"] for t in tools[:5]]
            except Exception:
                top_tools = []

            description = d.get("description", f"Tools for {domain_name}")
            tool_count = d.get("tool_count", 0)

            self._available_agents.append(AgentInfo(
                name=f"domain_{domain_name}",
                agent_type="DomainToolAgent",
                description=f"{description} ({tool_count} tools)",
                capabilities=top_tools,
                input_schema={"tool_name": "str", "arguments": "dict"},
                output_schema={"result": "Any"},
            ))

    def register_agent(self, agent: AgentInfo) -> None:
        """
        Register an additional agent.

        Args:
            agent: Agent information
        """
        # Remove existing with same name
        self._available_agents = [a for a in self._available_agents if a.name != agent.name]
        self._available_agents.append(agent)

    def plan_workflow(
        self,
        user_request: str,
        context: Optional[dict[str, Any]] = None,
    ) -> WorkflowPlan:
        """
        Create a workflow plan from a user request.

        Tries PydanticAI (multi-step) first, then Claude SDK, then
        deterministic fallback.

        Args:
            user_request: Natural language description of what the user wants
            context: Optional additional context (e.g., available data, constraints)

        Returns:
            WorkflowPlan with steps to execute
        """
        self._state = PlannerState.PLANNING

        try:
            if PYDANTIC_AI_AVAILABLE and run_planner_agent is not None and self._client:
                plan = self._plan_with_pydantic_ai(user_request, context)
            elif self._client:
                plan = self._plan_with_claude(user_request, context)
            else:
                plan = self._plan_fallback(user_request, context)

            self._plans.append(plan)
            self._state = PlannerState.COMPLETED
            return plan

        except Exception as e:
            logger.error(f"Planning failed: {e}")
            if self._config.fallback_on_llm_error:
                plan = self._plan_fallback(user_request, context)
                plan.warnings.append("LLM path failed; returned deterministic fallback")
                self._plans.append(plan)
                self._state = PlannerState.COMPLETED
                return plan

            self._state = PlannerState.FAILED
            raise

    def _plan_with_pydantic_ai(
        self,
        user_request: str,
        context: Optional[dict[str, Any]] = None,
    ) -> WorkflowPlan:
        """Use PydanticAI multi-step agent for planning.

        Falls back to Claude SDK on failure.
        """
        import asyncio

        agents_dicts = [a.to_dict() for a in self._available_agents]

        try:
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None

            if loop is not None and loop.is_running():
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor() as pool:
                    plan_output = pool.submit(
                        asyncio.run,
                        run_planner_agent(user_request, agents_dicts, context),
                    ).result(timeout=60)
            else:
                plan_output = asyncio.run(
                    run_planner_agent(user_request, agents_dicts, context)
                )
        except Exception as exc:
            logger.warning("PydanticAI planning failed (%s), falling back to Claude SDK", exc)
            return self._plan_with_claude(user_request, context)

        if plan_output is None:
            return self._plan_with_claude(user_request, context)

        return self._convert_plan_output(plan_output, user_request)

    def _convert_plan_output(self, plan_output: Any, user_request: str) -> WorkflowPlan:
        """Convert PydanticAI ``PlanOutput`` to a ``WorkflowPlan`` dataclass."""
        steps = []
        for s in plan_output.steps:
            steps.append(PlannedStep(
                id=s.id or f"step_{len(steps)+1}",
                name=s.name,
                description=s.description,
                agent_type=s.agent_type,
                capability=s.capability,
                input_mapping=s.input_mapping,
                output_key=s.output_key,
                depends_on=s.depends_on,
                reasoning=s.reasoning,
            ))

        return WorkflowPlan(
            id=str(uuid.uuid4())[:8],
            name=plan_output.name or "PydanticAI Workflow",
            description=plan_output.description,
            user_request=user_request,
            steps=steps,
            estimated_duration=plan_output.estimated_duration,
            confidence_score=plan_output.confidence_score,
            warnings=plan_output.warnings + ["Planned via PydanticAI multi-step agent"],
            alternatives=plan_output.alternatives,
        )

    def _plan_with_claude(
        self,
        user_request: str,
        context: Optional[dict[str, Any]] = None,
    ) -> WorkflowPlan:
        """Use Claude to create an intelligent plan."""

        # Build the system prompt
        system_prompt = self._build_system_prompt()

        # Build the user message
        user_message = self._build_user_message(user_request, context)

        # Call Claude
        response = self._client.messages.create(
            model=self._config.model,
            max_tokens=self._config.max_tokens,
            temperature=self._config.temperature,
            system=system_prompt,
            messages=[{"role": "user", "content": user_message}],
        )

        # Parse the response
        plan = self._parse_claude_response(response.content[0].text, user_request)

        return plan

    def _build_system_prompt(self) -> str:
        """Build the system prompt for Claude."""

        # Group agents by type for structured display
        core_agents = [a for a in self._available_agents if a.agent_type not in ("SkillAgent", "LangGraphAgent", "DomainToolAgent")]
        skill_agents = [a for a in self._available_agents if a.agent_type == "SkillAgent"]
        langgraph_agents = [a for a in self._available_agents if a.agent_type == "LangGraphAgent"]
        domain_agents = [a for a in self._available_agents if a.agent_type == "DomainToolAgent"]

        def _format_agents(agents: list[AgentInfo]) -> str:
            return "\n".join(
                f"- **{a.name}** ({a.agent_type}): {a.description}\n"
                f"  Capabilities: {', '.join(a.capabilities[:8])}"
                for a in agents
            )

        sections = [f"## Available Agents ({len(self._available_agents)} total)\n"]

        if core_agents:
            sections.append(f"### Core Agents\n{_format_agents(core_agents)}")
        if skill_agents:
            sections.append(f"### Skill Agents\n{_format_agents(skill_agents)}")
        if langgraph_agents:
            sections.append(f"### LangGraph Specialist Agents\n{_format_agents(langgraph_agents)}")
        if domain_agents:
            sections.append(f"### Domain Tool Agents\n{_format_agents(domain_agents)}")

        agents_description = "\n\n".join(sections)

        return f"""You are a workflow planning agent for DataBridge AI, a data engineering platform.
Your job is to analyze user requests and create optimal workflow plans using available agents.

{agents_description}

## Planning Guidelines

1. **Identify the Goal**: Understand what the user wants to achieve
2. **Select Agents**: Choose the minimum set of agents needed
3. **Order Steps**: Determine dependencies and optimal execution order
4. **Map Data Flow**: Define how outputs from one step feed into the next
5. **Consider Parallelism**: Steps without dependencies can run in parallel

## Output Format

Respond with a JSON object containing the workflow plan:

```json
{{
    "name": "Descriptive workflow name",
    "description": "What this workflow accomplishes",
    "confidence_score": 0.95,
    "estimated_duration": "5-10 minutes",
    "steps": [
        {{
            "id": "step_1",
            "name": "Step name",
            "description": "What this step does",
            "agent_type": "agent_name",
            "capability": "capability_name",
            "input_mapping": {{"param": "source.path"}},
            "output_key": "step_1_output",
            "depends_on": [],
            "reasoning": "Why this step is needed"
        }}
    ],
    "warnings": ["Any concerns or limitations"],
    "alternatives": ["Other approaches considered"]
}}
```

## DataBridge Workflow Knowledge — E2E Assessment + BLCE Pipeline

When a user asks about creating a data model, data warehouse, or data mapping, follow this two-phase approach:

### Phase 1: E2E Assessment (Initial Data Model & Warehouse)
The **E2E Assessment Pipeline** is what we use to create the initial data model and data warehouse
from raw source system tables. It runs 22 phases (schema discovery → AI classification →
star schema design → dbt generation → DDL deployment → quality validation).
The output is a working data warehouse with DDL scripts built directly from source tables.

**After Phase 1 completes**, proactively suggest Phase 2 to the user:
> "Now that the initial data warehouse is built from your raw source tables, we can refine it
> further. If you share your existing reports, tables, formulas, SQL queries, or any system
> requirements from your knowledge base, I can create a final version of the data model that
> caters to your actual business needs."

### Phase 2: BLCE Refinement (Business Logic → Data Marts)
The **BLCE Engine** (Business Logic Comprehension Engine, 72 tools, 21 phases) takes the
E2E output warehouse and enriches it using business context:
1. **Gather context**: Existing reports, SQL queries, formulas, system requirements, knowledge base docs
2. **Extract business logic**: Parse SQL CASE statements, calculations, aggregations, and business rules
3. **Identify configurations**: Core configurations, cross-references, and hierarchies the user can build
4. **Generate data marts**: Create additional data marts on top of the BLCE output warehouse
5. **Iterate**: The user can continue refining — building more hierarchies, adding cross-references,
   and generating additional mart layers even after the initial BLCE output is produced

### GraphRAG Context Persistence
After the E2E Assessment builds the initial warehouse, its DDLs, schema, and data model are
stored in **GraphRAG** (vector store + lineage graph). This is critical because:
- When suggesting edits, data mart additions, or hierarchy changes later, the AI knows exactly
  what exists in the current warehouse and what can be safely modified
- Wright pipeline generation for new reporting hierarchies references the GraphRAG-stored DDLs
  to determine what tables/columns exist and what needs to be added or altered
- The Hierarchy-Graph Bridge auto-syncs any hierarchy changes into GraphRAG
- BLCE cross-references and configurations are also indexed in GraphRAG for retrieval

### Planning Data Warehouse Requests
- For "create a data model/warehouse" → Plan the E2E Assessment workflow
- Always mention Phase 2 (BLCE refinement) as a suggested next step in alternatives
- If user provides reports/SQL/formulas → Plan a BLCE refinement workflow
- If user asks about data marts → Plan BLCE mart generation from existing warehouse
- After any warehouse change, mention that DDLs/schema are persisted in GraphRAG for future reference

## Important Rules

- Only use agents and capabilities from the available list
- Ensure step dependencies form a valid DAG (no cycles)
- Use clear, descriptive names for steps and output keys
- Include reasoning for each step when beneficial
- Flag any ambiguities or missing information as warnings
- Keep plans focused and avoid unnecessary steps
- Prefer SkillAgents for domain-specific financial/operational tasks
- Use LangGraph agents for multi-phase pipeline orchestration
- Use DomainToolAgents when specific tool capabilities are needed

## Connection Handling
If the user hasn't specified connection details, help them configure the connection step by
step like a normal LLM assistant would — ask about database type, host, credentials, schema.
Don't just warn about missing connections; actively guide them through setup.

## LLM Fallback for Complex Questions
If a user's question is beyond the scope of available agents (e.g., explaining concepts,
debugging issues, interpreting results), provide direct helpful answers like a knowledgeable
LLM assistant rather than only returning workflow plans. The planner can act as a general
data engineering advisor when no specific agent action is needed.

{self._get_purpose_context()}{self._get_tool_registry_context()}{self._get_hierarchy_context()}"""

    def _get_hierarchy_context(self) -> str:
        """Get active hierarchy project context for prompt injection."""
        if not self._hierarchy_service:
            return ""

        try:
            projects = self._hierarchy_service.list_projects()
            if not projects:
                return ""

            lines = ["## Active Hierarchy Context\n"]
            for proj in projects[:5]:
                pid = proj.get("id", "")
                name = proj.get("name", "")
                count = proj.get("hierarchy_count", 0)
                hierarchies = self._hierarchy_service.list_hierarchies(pid)
                mapping_count = sum(len(h.get("mapping", [])) for h in hierarchies)
                lines.append(
                    f"- **{name}** (id: {pid}): {count} hierarchies, {mapping_count} source mappings"
                )

            return "\n".join(lines) + "\n"
        except Exception:
            return ""

    @staticmethod
    def _load_purpose(purpose_path: Optional[str] = None) -> Optional[str]:
        """Load purpose.md strategic vision document."""
        candidates = []
        if purpose_path:
            candidates.append(Path(purpose_path))
        # Auto-discover relative to project root
        candidates.append(Path(__file__).parent.parent.parent / "purpose.md")
        for p in candidates:
            try:
                if p.exists():
                    return p.read_text(encoding="utf-8")
            except Exception:
                continue
        return None

    def _get_purpose_context(self) -> str:
        """Return curated strategic summary from purpose.md for prompt injection."""
        if not self._purpose_text:
            return ""
        return """## DataBridge AI Strategic Context

Core strategy: **Intelligence Arbitrage** — converting consulting expertise into scalable autonomous agents.

Key principles the planner must embody:
- **Automated Trust**: Continuous reconciliation → root-cause diagnosis → automated correction → certified closure
- **GraphRAG Memory**: Preserve edge-case knowledge across projects; later projects run faster than earlier ones
- **Sprinkler Behavior**: Don't just flag mismatches — detect, diagnose, fix, verify (drive closure)
- **CEO/IT Bridge**: Map technical structures into business semantics (EBITDA, margin, churn)
- **Delivery Compression**: Automate discovery, embed industry logic, self-heal quality loops

Planning directives:
- Prioritize reusable knowledge capture (GraphRAG ingestion steps)
- Always include validation/reconciliation steps in plans
- Map technical outputs to business-semantic naming
- Prefer skill agents for domain-specific financial/operational tasks
- Include quality closure loops, not just one-pass execution

"""

    def _get_tool_registry_context(self) -> str:
        """Return domain-level tool summary from the dynamic registry."""
        if not self._tool_registry:
            return ""
        try:
            domains = self._tool_registry.get_all_domains()
            if not domains:
                return ""
            total_tools = sum(d.get("tool_count", 0) for d in domains)
            lines = [f"## Available Tool Domains ({total_tools} tools across {len(domains)} domains)\n"]
            for d in domains:
                name = d.get("domain", "unknown")
                count = d.get("tool_count", 0)
                desc = d.get("description", "")
                lines.append(f"- **{name}** ({count} tools): {desc}")
            return "\n".join(lines) + "\n\n"
        except Exception:
            return ""

    def suggest_tool_domains(self, user_request: str) -> list[dict[str, Any]]:
        """Suggest relevant tool domains for a user request.

        Returns top 5 domains scored by keyword match against domain descriptions.
        """
        if not self._tool_registry:
            return []
        try:
            domains = self._tool_registry.get_all_domains()
        except Exception:
            return []

        request_lower = user_request.lower()
        request_words = set(request_lower.split())
        scored = []
        for d in domains:
            desc_words = set(d.get("description", "").lower().split())
            name = d.get("domain", "")
            # Score: keyword overlap + name match
            overlap = len(request_words & desc_words)
            score = min(overlap * 0.15, 0.8)
            if name.lower() in request_lower:
                score += 0.3
            if score > 0:
                scored.append({
                    "domain": name,
                    "tool_count": d.get("tool_count", 0),
                    "description": d.get("description", ""),
                    "relevance_score": round(min(score, 1.0), 2),
                })
        scored.sort(key=lambda x: x["relevance_score"], reverse=True)
        return scored[:5]

    def _build_user_message(
        self,
        user_request: str,
        context: Optional[dict[str, Any]] = None,
    ) -> str:
        """Build the user message for Claude."""

        message = f"""Please create a workflow plan for the following request:

## User Request
{user_request}
"""

        if context:
            message += f"""
## Additional Context
```json
{json.dumps(context, indent=2)}
```
"""

        message += """
Please respond with a valid JSON workflow plan.
"""

        return message

    def _parse_claude_response(
        self,
        response_text: str,
        user_request: str,
    ) -> WorkflowPlan:
        """Parse Claude's response into a WorkflowPlan."""

        # Extract JSON from response
        json_start = response_text.find("{")
        json_end = response_text.rfind("}") + 1

        if json_start == -1 or json_end == 0:
            raise ValueError("No JSON found in Claude response")

        json_str = response_text[json_start:json_end]
        data = json.loads(json_str)

        # Build steps
        steps = []
        for step_data in data.get("steps", []):
            step = PlannedStep(
                id=step_data.get("id", f"step_{len(steps)+1}"),
                name=step_data.get("name", "Unnamed Step"),
                description=step_data.get("description", ""),
                agent_type=step_data.get("agent_type", ""),
                capability=step_data.get("capability", ""),
                input_mapping=step_data.get("input_mapping", {}),
                output_key=step_data.get("output_key", f"step_{len(steps)+1}_output"),
                depends_on=step_data.get("depends_on", []),
                config=step_data.get("config", {}),
                reasoning=step_data.get("reasoning", ""),
            )
            steps.append(step)

        # Create plan
        plan = WorkflowPlan(
            id=str(uuid.uuid4())[:8],
            name=data.get("name", "Generated Workflow"),
            description=data.get("description", ""),
            user_request=user_request,
            steps=steps,
            estimated_duration=data.get("estimated_duration", "Unknown"),
            confidence_score=data.get("confidence_score", 0.8),
            warnings=data.get("warnings", []),
            alternatives=data.get("alternatives", []),
        )

        return plan

    def _plan_fallback(
        self,
        user_request: str,
        context: Optional[dict[str, Any]] = None,
    ) -> WorkflowPlan:
        """
        Create a basic plan without Claude (fallback mode).

        Uses keyword matching to identify relevant agents.
        """
        request_lower = user_request.lower()
        steps = []

        # Keyword-based agent selection
        agent_keywords = {
            "schema_scanner": ["scan", "schema", "database", "tables", "columns", "metadata"],
            "logic_extractor": ["sql", "case", "extract", "parse", "logic", "calculation"],
            "warehouse_architect": ["design", "star schema", "dimension", "fact", "dbt", "model"],
            "deploy_validator": ["deploy", "execute", "validate", "ddl", "test"],
            "hierarchy_builder": ["hierarchy", "tree", "parent", "child", "level", "mapping"],
            "data_reconciler": ["reconcile", "compare", "match", "discrepancy", "orphan"],
        }

        selected_agents = []
        for agent_name, keywords in agent_keywords.items():
            if any(kw in request_lower for kw in keywords):
                agent_info = next(
                    (a for a in self._available_agents if a.name == agent_name),
                    None
                )
                if agent_info:
                    selected_agents.append(agent_info)

        # Create steps for selected agents
        for i, agent in enumerate(selected_agents):
            step = PlannedStep(
                id=f"step_{i+1}",
                name=f"{agent.agent_type} Execution",
                description=agent.description,
                agent_type=agent.name,
                capability=agent.capabilities[0] if agent.capabilities else "",
                output_key=f"step_{i+1}_output",
                depends_on=[f"step_{i}"] if i > 0 else [],
                reasoning=f"Selected based on keyword match in user request",
            )
            steps.append(step)

        # Default to a generic step if no agents matched
        if not steps:
            steps.append(PlannedStep(
                id="step_1",
                name="Manual Review Required",
                description="Could not automatically plan this request",
                agent_type="manual",
                capability="review",
                reasoning="No agents matched the user request",
            ))

        return WorkflowPlan(
            id=str(uuid.uuid4())[:8],
            name="Auto-Generated Workflow",
            description=f"Workflow for: {user_request[:100]}",
            user_request=user_request,
            steps=steps,
            confidence_score=0.5,  # Lower confidence for fallback
            warnings=["Generated without AI assistance - review recommended"],
        )

    def analyze_request(
        self,
        user_request: str,
    ) -> dict[str, Any]:
        """
        Analyze a user request without creating a full plan.

        Args:
            user_request: Natural language request

        Returns:
            Analysis including intent, entities, and suggested agents
        """
        if not self._client:
            return self._analyze_fallback(user_request)

        response = self._client.messages.create(
            model=self._config.model,
            max_tokens=1024,
            temperature=0.2,
            system="""Analyze the user's data engineering request and extract:
1. Primary intent (what they want to achieve)
2. Key entities (tables, schemas, files mentioned)
3. Constraints (time, quality, format requirements)
4. Ambiguities (unclear aspects that need clarification)

Respond in JSON format.""",
            messages=[{"role": "user", "content": user_request}],
        )

        try:
            text = response.content[0].text
            json_start = text.find("{")
            json_end = text.rfind("}") + 1
            return json.loads(text[json_start:json_end])
        except Exception:
            return self._analyze_fallback(user_request)

    def _analyze_fallback(self, user_request: str) -> dict[str, Any]:
        """Basic analysis without Claude."""
        return {
            "intent": "unknown",
            "entities": [],
            "constraints": [],
            "ambiguities": ["Full analysis requires Claude API"],
            "suggested_agents": [
                a.name for a in self._available_agents
                if any(word in user_request.lower() for word in a.description.lower().split()[:5])
            ],
        }

    def suggest_agents(
        self,
        user_request: str,
    ) -> dict[str, Any]:
        """
        Suggest which agents could handle a request with enriched scoring.

        Args:
            user_request: Natural language request

        Returns:
            Dict with overall suggestions, grouped by type, and tool domains
        """
        request_lower = user_request.lower()
        request_words = set(request_lower.split())

        # Purpose alignment keywords
        purpose_keywords = {
            "trust", "reconciliation", "reconcile", "hierarchy", "quality",
            "validation", "validate", "recon", "audit", "compliance",
        }
        has_purpose_alignment = bool(request_words & purpose_keywords)

        suggestions = []
        for agent in self._available_agents:
            relevance = 0.0
            matched_capabilities = []

            # Capability match (0.3 each)
            for cap in agent.capabilities:
                cap_lower = cap.lower().replace("-", " ").replace("_", " ")
                cap_words = set(cap_lower.split())
                if cap_lower in request_lower or cap_words & request_words:
                    relevance += 0.3
                    matched_capabilities.append(cap)

            # Description keyword match (0.1 each, max 0.5)
            desc_words = set(agent.description.lower().split())
            # Filter out noise words
            noise = {"the", "a", "an", "and", "or", "for", "to", "of", "in", "is", "with", "from", "by"}
            meaningful_desc = desc_words - noise
            word_matches = len(meaningful_desc & request_words)
            relevance += min(word_matches * 0.1, 0.5)

            # Agent type bonus
            if agent.agent_type == "SkillAgent":
                # Skill agents get a boost when domain keywords match
                if any(tag in request_lower for tag in ["financial", "accounting", "gl", "p&l",
                        "oil", "gas", "cost", "budget", "operations", "geographic",
                        "saas", "manufacturing", "transportation", "workflow"]):
                    relevance += 0.2
            elif agent.agent_type == "LangGraphAgent":
                if any(phase in request_lower for phase in ["schema", "metadata", "classify",
                        "dimension", "dbt", "wright", "deploy", "quality", "cortex"]):
                    relevance += 0.15
            elif agent.agent_type == "DomainToolAgent":
                relevance += 0.1 if relevance > 0 else 0

            # Purpose alignment bonus
            if has_purpose_alignment:
                relevance += 0.1

            if relevance > 0:
                suggestions.append({
                    "agent_name": agent.name,
                    "agent_type": agent.agent_type,
                    "relevance_score": round(min(relevance, 1.0), 2),
                    "matched_capabilities": matched_capabilities,
                    "description": agent.description,
                })

        # Sort by relevance
        suggestions.sort(key=lambda x: x["relevance_score"], reverse=True)

        # Group by type
        by_type = {
            "core": [s for s in suggestions if s["agent_type"] not in ("SkillAgent", "LangGraphAgent", "DomainToolAgent")],
            "skills": [s for s in suggestions if s["agent_type"] == "SkillAgent"],
            "langgraph": [s for s in suggestions if s["agent_type"] == "LangGraphAgent"],
            "domains": [s for s in suggestions if s["agent_type"] == "DomainToolAgent"],
        }

        # Tool domain suggestions
        tool_domains = self.suggest_tool_domains(user_request)

        return {
            "agents": suggestions[:15],
            "by_type": by_type,
            "tool_domains": tool_domains,
            "total_matched": len(suggestions),
        }

    def optimize_plan(
        self,
        plan: WorkflowPlan,
    ) -> WorkflowPlan:
        """
        Optimize an existing plan for better performance.

        Args:
            plan: Existing workflow plan

        Returns:
            Optimized plan
        """
        if not self._config.enable_parallel_steps:
            return plan

        # Identify steps that can run in parallel
        # (steps with no dependencies on each other)
        optimized_steps = list(plan.steps)

        # Group steps by their dependency level
        levels: dict[int, list[PlannedStep]] = {}
        step_levels: dict[str, int] = {}

        for step in optimized_steps:
            if not step.depends_on:
                step_levels[step.id] = 0
            else:
                max_dep_level = max(
                    step_levels.get(dep, 0) for dep in step.depends_on
                )
                step_levels[step.id] = max_dep_level + 1

            level = step_levels[step.id]
            if level not in levels:
                levels[level] = []
            levels[level].append(step)

        # Add parallelism hints
        for level, steps_at_level in levels.items():
            if len(steps_at_level) > 1:
                for step in steps_at_level:
                    step.config["parallel_group"] = f"level_{level}"

        return WorkflowPlan(
            id=plan.id,
            name=plan.name,
            description=plan.description,
            user_request=plan.user_request,
            steps=optimized_steps,
            estimated_duration=plan.estimated_duration,
            confidence_score=plan.confidence_score,
            warnings=plan.warnings + [f"Optimized: {len(levels)} parallel levels identified"],
            alternatives=plan.alternatives,
            created_at=plan.created_at,
        )

    def explain_plan(
        self,
        plan: WorkflowPlan,
    ) -> str:
        """
        Generate a human-readable explanation of a plan.

        Args:
            plan: Workflow plan to explain

        Returns:
            Markdown explanation
        """
        explanation = f"""# Workflow Plan: {plan.name}

## Overview
{plan.description}

**Confidence Score:** {plan.confidence_score:.0%}
**Estimated Duration:** {plan.estimated_duration}
**Total Steps:** {len(plan.steps)}

## User Request
> {plan.user_request}

## Execution Steps

"""

        for i, step in enumerate(plan.steps, 1):
            deps = ", ".join(step.depends_on) if step.depends_on else "None (can start immediately)"

            explanation += f"""### Step {i}: {step.name}

- **Agent:** {step.agent_type}
- **Capability:** {step.capability}
- **Dependencies:** {deps}
- **Output Key:** {step.output_key}

{step.description}

"""
            if step.reasoning:
                explanation += f"**Reasoning:** {step.reasoning}\n\n"

        if plan.warnings:
            explanation += "## Warnings\n\n"
            for warning in plan.warnings:
                explanation += f"- {warning}\n"
            explanation += "\n"

        if plan.alternatives:
            explanation += "## Alternative Approaches\n\n"
            for alt in plan.alternatives:
                explanation += f"- {alt}\n"

        return explanation

    def to_workflow_definition(
        self,
        plan: WorkflowPlan,
    ) -> dict[str, Any]:
        """
        Convert a plan to a workflow definition for the Orchestrator.

        Args:
            plan: Workflow plan

        Returns:
            Workflow definition dict compatible with Orchestrator.create_workflow()
        """
        return {
            "name": plan.name,
            "description": plan.description,
            "steps": [
                {
                    "name": step.name,
                    "agent": step.agent_type,
                    "capability": step.capability,
                    "input_mapping": step.input_mapping,
                    "output_key": step.output_key,
                    "depends_on": step.depends_on,
                    "config": step.config,
                }
                for step in plan.steps
            ],
        }

    def to_python_script(
        self,
        plan: dict,
        workflow_type: str = "auto",
        config_overrides: Optional[dict] = None,
        user_edits: Optional[str] = None,
    ) -> dict[str, Any]:
        """Generate an executable Python script from a workflow plan.

        Args:
            plan: Plan dict (from WorkflowPlan.to_dict())
            workflow_type: "auto", "e2e", "blce", or "custom"
            config_overrides: Override source/dest/limit configs
            user_edits: User-provided code edits (custom only)

        Returns:
            Dict with script_path, code_preview, code_full,
            workflow_type, phases, warnings
        """
        from src.agents.script_generator import ScriptGenerator

        gen = ScriptGenerator(
            anthropic_client=self._client,
            output_dir="scripts/generated",
        )
        result = gen.generate(plan, workflow_type, config_overrides, user_edits)
        return {
            "script_path": result.script_path,
            "code_preview": result.code[:2000],
            "code_full": result.code,
            "workflow_type": result.workflow_type,
            "phases": result.phases,
            "warnings": result.warnings,
            "plan_id": result.plan_id,
        }

    def get_status(self) -> dict[str, Any]:
        """Get planner status with agent counts by type."""
        core = [a for a in self._available_agents if a.agent_type not in ("SkillAgent", "LangGraphAgent", "DomainToolAgent")]
        skills = [a for a in self._available_agents if a.agent_type == "SkillAgent"]
        langgraph = [a for a in self._available_agents if a.agent_type == "LangGraphAgent"]
        domains = [a for a in self._available_agents if a.agent_type == "DomainToolAgent"]

        return {
            "planner_id": self._id,
            "state": self._state.value,
            "purpose_loaded": self._purpose_text is not None,
            "tool_registry_available": self._tool_registry is not None,
            "agent_counts": {
                "core": len(core),
                "skill": len(skills),
                "langgraph": len(langgraph),
                "domain": len(domains),
                "total": len(self._available_agents),
            },
            "plans_created": len(self._plans),
            "claude_available": self._client is not None,
            "model": self._config.model,
            "created_at": self._created_at.isoformat(),
        }

    def get_history(self, limit: int = 10) -> list[dict[str, Any]]:
        """Get planning history."""
        return [p.to_dict() for p in self._plans[-limit:]]

    def get_agent_graph(self) -> dict[str, Any]:
        """Return a vis.js-compatible graph of all agents with dependency edges.

        Returns:
            Dict with 'nodes' and 'edges' lists for graph rendering.
        """
        type_colors = {
            "core": "#38bdf8",       # blue
            "SkillAgent": "#4ade80",       # green
            "LangGraphAgent": "#818cf8",   # purple
            "DomainToolAgent": "#f59e0b",  # amber
        }
        type_groups = {
            "SkillAgent": "skills",
            "LangGraphAgent": "langgraph",
            "DomainToolAgent": "domains",
        }

        nodes = []
        agent_names = set()
        for agent in self._available_agents:
            agent_names.add(agent.name)
            agent_group = type_groups.get(agent.agent_type, "core")
            color = type_colors.get(agent.agent_type, type_colors["core"])
            # Friendly label: remove prefix and underscores
            label = agent.name.replace("_", " ").replace("domain ", "").replace("skill ", "").replace("langgraph ", "").title()
            nodes.append({
                "id": agent.name,
                "label": label,
                "type": agent_group,
                "color": color,
                "group": agent_group,
                "description": agent.description[:120],
            })

        # Filter edges to only include agents that exist
        edges = []
        for dep in AGENT_DEPENDENCIES:
            src = dep.get("source", "")
            tgt = dep.get("target", "")
            if src in agent_names and tgt in agent_names:
                edges.append({
                    "from": src,
                    "to": tgt,
                    "label": dep.get("label", ""),
                    "dashes": True,
                    "arrows": "to",
                })

        return {
            "nodes": nodes,
            "edges": edges,
            "node_count": len(nodes),
            "edge_count": len(edges),
        }


# ============================================================================
# Agent dependency graph — static edges between agent types
# ============================================================================

AGENT_DEPENDENCIES = [
    # Core → LangGraph (core agents map to LangGraph phases)
    {"source": "schema_scanner",      "target": "langgraph_data_modeling",  "label": "feeds schema"},
    {"source": "logic_extractor",     "target": "langgraph_data_modeling",  "label": "feeds SQL logic"},
    {"source": "warehouse_architect", "target": "langgraph_dbt",            "label": "feeds star schema"},
    {"source": "hierarchy_builder",   "target": "langgraph_hierarchy",      "label": "feeds hierarchy"},
    {"source": "deploy_validator",    "target": "langgraph_snowflake",      "label": "feeds DDL"},
    {"source": "data_reconciler",     "target": "langgraph_qa",             "label": "feeds recon"},
    {"source": "cortex_agent",        "target": "langgraph_cortex_llm",     "label": "feeds AI"},

    # LangGraph phase ordering
    {"source": "langgraph_data_modeling", "target": "langgraph_hierarchy",      "label": "schema → hierarchy"},
    {"source": "langgraph_data_modeling", "target": "langgraph_qa",             "label": "classify → quality"},
    {"source": "langgraph_hierarchy",     "target": "langgraph_snowflake",      "label": "hierarchy → wright"},
    {"source": "langgraph_hierarchy",     "target": "langgraph_dbt",            "label": "hierarchy → dbt"},
    {"source": "langgraph_snowflake",     "target": "langgraph_qa",             "label": "wright → quality"},
    {"source": "langgraph_qa",            "target": "langgraph_data_modeling",  "label": "quality → artifact"},

    # Skills → Core agents they depend on
    {"source": "skill_financial_analyst",   "target": "data_reconciler",           "label": "uses reconciliation"},
    {"source": "skill_financial_analyst",   "target": "hierarchy_builder",         "label": "uses hierarchy"},
    {"source": "skill_fpa_oil_gas_analyst", "target": "hierarchy_builder",         "label": "uses hierarchy"},
    {"source": "skill_fpa_oil_gas_analyst", "target": "cortex_agent",              "label": "uses AI"},
    {"source": "skill_fpa_cost_analyst",    "target": "hierarchy_builder",         "label": "uses hierarchy"},
    {"source": "skill_operations_analyst",  "target": "hierarchy_builder",         "label": "uses hierarchy"},
    {"source": "skill_operations_analyst",  "target": "data_reconciler",           "label": "uses reconciliation"},
    {"source": "skill_platform_workflow",   "target": "langgraph_data_modeling",   "label": "orchestrates"},

    # Domain agents → Core agents they extend
    {"source": "domain_reconciliation", "target": "data_reconciler",      "label": "extends"},
    {"source": "domain_hierarchy",      "target": "hierarchy_builder",    "label": "extends"},
    {"source": "domain_cortex",         "target": "cortex_agent",         "label": "extends"},
    {"source": "domain_blce",           "target": "logic_extractor",      "label": "extends"},
    {"source": "domain_wright",         "target": "langgraph_snowflake",  "label": "extends"},
    {"source": "domain_dbt",            "target": "langgraph_dbt",        "label": "extends"},
    {"source": "domain_quality",        "target": "langgraph_qa",         "label": "extends"},
    {"source": "domain_graphrag",       "target": "researcher_analyst",   "label": "extends"},
]


# Convenience function for quick planning
def plan_workflow(
    request: str,
    context: Optional[dict[str, Any]] = None,
    api_key: Optional[str] = None,
) -> WorkflowPlan:
    """
    Quick function to plan a workflow.

    Args:
        request: User request
        context: Optional context
        api_key: Optional Anthropic API key

    Returns:
        WorkflowPlan
    """
    config = PlannerConfig(api_key=api_key) if api_key else None
    planner = PlannerAgent(config=config)
    return planner.plan_workflow(request, context)
