"""MCP tool registration for LangGraph multi-agent framework.

Registers 15 MCP tools organized into 5 action-dispatched groups:
- langgraph_workflow (4 actions): run, resume, status, list_runs
- langgraph_agent (3 actions): invoke, list, capabilities
- langgraph_manage (3 actions): checkpoint, clear, configure
- langgraph_cortex (3 actions): complete, analyst, search
- langgraph_streaming (2 actions): subscribe, history
"""

import json
import logging
from typing import Any, Dict, Optional

from .unified import (
    dispatch_langgraph_workflow,
    dispatch_langgraph_agent,
    dispatch_langgraph_manage,
    dispatch_langgraph_cortex,
    dispatch_langgraph_streaming,
)

logger = logging.getLogger(__name__)


def register_langgraph_tools(mcp: Any) -> Dict[str, Any]:
    """Register all LangGraph MCP tools.

    Args:
        mcp: FastMCP server instance.

    Returns:
        Registration summary dict.
    """

    # =================================================================
    # Tool 1: langgraph_workflow
    # =================================================================

    @mcp.tool()
    def langgraph_workflow(
        action: str,
        config: Optional[str] = None,
        run_id: Optional[str] = None,
        human_gates: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> str:
        """
        LangGraph multi-agent workflow execution.

        Actions:
        - run: Execute a full multi-agent workflow (requires config JSON)
        - resume: Resume a workflow from checkpoint (requires run_id)
        - status: Get status of a workflow run (requires run_id)
        - list_runs: List past workflow runs (optional limit)

        Args:
            action: The action to perform (run, resume, status, list_runs)
            config: JSON config dict for the workflow (source_type, database, etc.)
            run_id: Workflow run ID (for resume/status)
            human_gates: JSON list of phase names requiring human approval
            limit: Max results for list_runs

        Returns:
            JSON result with action-specific data
        """
        kwargs = {k: v for k, v in {
            "config": config, "run_id": run_id,
            "human_gates": human_gates, "limit": limit,
        }.items() if v is not None}
        result = dispatch_langgraph_workflow(action, **kwargs)
        return json.dumps(result, indent=2, default=str)

    # =================================================================
    # Tool 2: langgraph_agent
    # =================================================================

    @mcp.tool()
    def langgraph_agent(
        action: str,
        agent_name: Optional[str] = None,
        config: Optional[str] = None,
    ) -> str:
        """
        Invoke or inspect LangGraph specialist agents.

        Actions:
        - invoke: Run a single specialist agent (requires agent_name)
        - list: List all available agents
        - capabilities: Get agent capabilities (requires agent_name)

        Available agents: data_modeling_agent, hierarchy_agent, dbt_agent,
        snowflake_agent, qa_agent, cortex_llm_agent, cortex_analyst_agent,
        cortex_search_agent

        Args:
            action: The action to perform (invoke, list, capabilities)
            agent_name: Name of the agent to invoke/inspect
            config: JSON config dict for agent invocation

        Returns:
            JSON result with agent data
        """
        kwargs = {k: v for k, v in {
            "agent_name": agent_name, "config": config,
        }.items() if v is not None}
        result = dispatch_langgraph_agent(action, **kwargs)
        return json.dumps(result, indent=2, default=str)

    # =================================================================
    # Tool 3: langgraph_manage
    # =================================================================

    @mcp.tool()
    def langgraph_manage(
        action: str,
        run_id: Optional[str] = None,
        operation: Optional[str] = None,
        state: Optional[str] = None,
    ) -> str:
        """
        Manage LangGraph workflow checkpoints and configuration.

        Actions:
        - checkpoint: Save or load a checkpoint (requires run_id, optional operation: save/load)
        - clear: Clear all checkpoints for a run (requires run_id)
        - configure: Show configuration options

        Args:
            action: The action to perform (checkpoint, clear, configure)
            run_id: Workflow run ID
            operation: For checkpoint action: 'save' or 'load' (default: load)
            state: JSON state dict for checkpoint save

        Returns:
            JSON result with management data
        """
        kwargs = {k: v for k, v in {
            "run_id": run_id, "operation": operation, "state": state,
        }.items() if v is not None}
        result = dispatch_langgraph_manage(action, **kwargs)
        return json.dumps(result, indent=2, default=str)

    # =================================================================
    # Tool 4: langgraph_cortex
    # =================================================================

    @mcp.tool()
    def langgraph_cortex(
        action: str,
        prompt: Optional[str] = None,
        question: Optional[str] = None,
        query: Optional[str] = None,
        model: Optional[str] = None,
        model_name: Optional[str] = None,
    ) -> str:
        """
        Cortex AI operations via LangGraph agents.

        Actions:
        - complete: Run Cortex LLM completion (requires prompt)
        - analyst: Query Cortex Analyst semantic model (requires question + model_name)
        - search: Run Cortex Search query (requires query)

        Args:
            action: The action to perform (complete, analyst, search)
            prompt: Text prompt for completion
            question: Question for Cortex Analyst
            query: Search query for Cortex Search
            model: Cortex model name (default: mistral-large)
            model_name: Semantic model name for Analyst

        Returns:
            JSON result with Cortex data
        """
        kwargs = {k: v for k, v in {
            "prompt": prompt, "question": question, "query": query,
            "model": model, "model_name": model_name,
        }.items() if v is not None}
        result = dispatch_langgraph_cortex(action, **kwargs)
        return json.dumps(result, indent=2, default=str)

    # =================================================================
    # Tool 5: langgraph_streaming
    # =================================================================

    @mcp.tool()
    def langgraph_streaming(
        action: str,
        run_id: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> str:
        """
        Subscribe to and retrieve LangGraph workflow progress events.

        Actions:
        - subscribe: Subscribe to progress events for a run (requires run_id)
        - history: Get event history for a run (requires run_id, optional limit)

        Args:
            action: The action to perform (subscribe, history)
            run_id: Workflow run ID to track
            limit: Max events to return (default: 100)

        Returns:
            JSON result with streaming data
        """
        kwargs = {k: v for k, v in {
            "run_id": run_id, "limit": limit,
        }.items() if v is not None}
        result = dispatch_langgraph_streaming(action, **kwargs)
        return json.dumps(result, indent=2, default=str)

    # =================================================================
    # Summary
    # =================================================================

    logger.info("Registered 5 LangGraph MCP tools (15 actions)")
    return {
        "tools_registered": 5,
        "total_actions": 15,
        "tools": [
            "langgraph_workflow",
            "langgraph_agent",
            "langgraph_manage",
            "langgraph_cortex",
            "langgraph_streaming",
        ],
    }
