"""
Unified MCP tools for PlannerAgent (Phase 3D-2 consolidation).

Consolidates 11 individual planner tools into 3 action-dispatched tools:
- planner_workflow(action, ...) — 3 actions: plan, analyze, suggest_agents
- planner_transform(operation, ...) — 3 operations: explain, optimize, to_workflow_def
- planner_manage(action, ...) — 5 actions: list_agents, register_agent, status, history, configure

Original tool names are preserved as deprecation wrappers in mcp_tools.py.
Uses mcp_only pattern (no settings parameter).
"""

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

logger = logging.getLogger(__name__)

# ============================================================================
# Lazy singleton + context holders
# ============================================================================

_planner = None
_tool_registry = None
_purpose_path = None


def set_planner_context(tool_registry=None, purpose_path=None):
    """Set registry and purpose context for the planner singleton.

    Call after DynamicToolRegistry.finalize() in server.py.
    If the planner already exists, it is recreated with the new context.
    """
    global _tool_registry, _purpose_path, _planner
    _tool_registry = tool_registry
    _purpose_path = purpose_path
    if _planner is not None:
        # Recreate with new context
        from .planner_agent import PlannerAgent
        _planner = PlannerAgent(
            config=_planner._config,
            tool_registry=_tool_registry,
            purpose_path=_purpose_path,
        )


def _ensure_planner():
    global _planner
    if _planner is None:
        from .planner_agent import PlannerAgent
        _planner = PlannerAgent(
            tool_registry=_tool_registry,
            purpose_path=_purpose_path,
        )
    return _planner


def _reset_planner(config=None):
    """Reset the planner with new configuration."""
    global _planner
    from .planner_agent import PlannerAgent
    _planner = PlannerAgent(
        config=config,
        tool_registry=_tool_registry,
        purpose_path=_purpose_path,
    )
    return _planner


def _reconstruct_plan(data: dict):
    """Reconstruct a WorkflowPlan from a JSON dict."""
    from .planner_agent import PlannedStep, WorkflowPlan

    steps = [
        PlannedStep(
            id=s.get("id", f"step_{i}"),
            name=s.get("name", ""),
            description=s.get("description", ""),
            agent_type=s.get("agent_type", ""),
            capability=s.get("capability", ""),
            input_mapping=s.get("input_mapping", {}),
            output_key=s.get("output_key", ""),
            depends_on=s.get("depends_on", []),
            config=s.get("config", {}),
            reasoning=s.get("reasoning", ""),
        )
        for i, s in enumerate(data.get("steps", []))
    ]

    return WorkflowPlan(
        id=data.get("id", ""),
        name=data.get("name", ""),
        description=data.get("description", ""),
        user_request=data.get("user_request", ""),
        steps=steps,
        estimated_duration=data.get("estimated_duration", ""),
        confidence_score=data.get("confidence_score", 0.0),
        warnings=data.get("warnings", []),
        alternatives=data.get("alternatives", []),
    )


# ============================================================================
# planner_workflow action handlers
# ============================================================================

def _workflow_plan(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    request = kwargs.get("request")
    if not request:
        return {"error": "request is required for 'plan' action"}

    context = kwargs.get("context")
    ctx = None
    if context:
        try:
            ctx = json.loads(context)
        except json.JSONDecodeError:
            ctx = None

    plan = planner.plan_workflow(request, ctx)
    return plan.to_dict()


def _workflow_analyze(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    request = kwargs.get("request")
    if not request:
        return {"error": "request is required for 'analyze' action"}

    analysis = planner.analyze_request(request)
    return analysis


def _workflow_suggest_agents(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    request = kwargs.get("request")
    if not request:
        return {"error": "request is required for 'suggest_agents' action"}

    result = planner.suggest_agents(request)
    # suggest_agents now returns an enriched dict with agents, by_type, tool_domains
    return result


# ============================================================================
# planner_transform operation handlers
# ============================================================================

def _transform_explain(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    plan_json = kwargs.get("plan_json")
    if not plan_json:
        return {"error": "plan_json is required for 'explain' operation"}

    try:
        data = json.loads(plan_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON in plan_json"}

    plan = _reconstruct_plan(data)
    explanation = planner.explain_plan(plan)
    return {"explanation": explanation}


def _transform_optimize(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    plan_json = kwargs.get("plan_json")
    if not plan_json:
        return {"error": "plan_json is required for 'optimize' operation"}

    try:
        data = json.loads(plan_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON in plan_json"}

    plan = _reconstruct_plan(data)
    optimized = planner.optimize_plan(plan)
    return optimized.to_dict()


def _transform_to_workflow_def(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    plan_json = kwargs.get("plan_json")
    if not plan_json:
        return {"error": "plan_json is required for 'to_workflow_def' operation"}

    try:
        data = json.loads(plan_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON in plan_json"}

    plan = _reconstruct_plan(data)
    workflow_def = planner.to_workflow_definition(plan)
    return workflow_def


# ============================================================================
# planner_manage action handlers
# ============================================================================

def _manage_list_agents(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    agents = [a.to_dict() for a in planner.available_agents]
    return {"agents": agents, "count": len(agents)}


def _manage_register_agent(**kwargs) -> Dict[str, Any]:
    from .planner_agent import AgentInfo

    planner = _ensure_planner()
    name = kwargs.get("name")
    agent_type = kwargs.get("agent_type")
    description = kwargs.get("description")
    capabilities = kwargs.get("capabilities")

    if not all([name, agent_type, description, capabilities]):
        return {"error": "name, agent_type, description, capabilities are required for 'register_agent' action"}

    try:
        caps = json.loads(capabilities)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON in capabilities"}

    input_schema = {}
    if kwargs.get("input_schema"):
        try:
            input_schema = json.loads(kwargs["input_schema"])
        except json.JSONDecodeError:
            return {"error": "Invalid JSON in input_schema"}

    output_schema = {}
    if kwargs.get("output_schema"):
        try:
            output_schema = json.loads(kwargs["output_schema"])
        except json.JSONDecodeError:
            return {"error": "Invalid JSON in output_schema"}

    agent = AgentInfo(
        name=name,
        agent_type=agent_type,
        description=description,
        capabilities=caps,
        input_schema=input_schema,
        output_schema=output_schema,
    )

    planner.register_agent(agent)

    return {
        "status": "success",
        "message": f"Agent '{name}' registered successfully",
        "agent": agent.to_dict(),
    }


def _manage_status(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    return planner.get_status()


def _manage_history(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    limit = kwargs.get("limit", 10)
    history = planner.get_history(limit)
    return {"history": history, "count": len(history)}


def _manage_configure(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()

    model = kwargs.get("model")
    temperature = kwargs.get("temperature")
    max_steps = kwargs.get("max_steps")
    enable_parallel = kwargs.get("enable_parallel")

    if model:
        planner._config.model = model
    if temperature is not None and temperature >= 0:
        planner._config.temperature = temperature
    if max_steps is not None and max_steps > 0:
        planner._config.max_steps = max_steps
    if enable_parallel:
        planner._config.enable_parallel_steps = enable_parallel.lower() == "true"

    return {
        "status": "success",
        "config": {
            "model": planner._config.model,
            "temperature": planner._config.temperature,
            "max_steps": planner._config.max_steps,
            "enable_parallel_steps": planner._config.enable_parallel_steps,
        },
    }


# ============================================================================
# Action dispatch dictionaries
# ============================================================================

_WORKFLOW_ACTIONS = {
    "plan": _workflow_plan,
    "analyze": _workflow_analyze,
    "suggest_agents": _workflow_suggest_agents,
}

_TRANSFORM_OPS = {
    "explain": _transform_explain,
    "optimize": _transform_optimize,
    "to_workflow_def": _transform_to_workflow_def,
}

def _manage_agent_graph(**kwargs) -> Dict[str, Any]:
    planner = _ensure_planner()
    return planner.get_agent_graph()


_MANAGE_ACTIONS = {
    "list_agents": _manage_list_agents,
    "register_agent": _manage_register_agent,
    "status": _manage_status,
    "history": _manage_history,
    "configure": _manage_configure,
    "agent_graph": _manage_agent_graph,
}


# ============================================================================
# Unified dispatch functions (mcp_only — no settings parameter)
# ============================================================================

def dispatch_planner_workflow(action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a planner_workflow action."""
    handler = _WORKFLOW_ACTIONS.get(action)
    if not handler:
        return {
            "error": f"Unknown action: '{action}'",
            "valid_actions": sorted(_WORKFLOW_ACTIONS.keys()),
        }
    try:
        return handler(**kwargs)
    except Exception as e:
        logger.error(f"planner_workflow({action}) failed: {e}")
        return {"error": f"planner_workflow({action}) failed: {e}"}


def dispatch_planner_transform(operation: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a planner_transform operation."""
    handler = _TRANSFORM_OPS.get(operation)
    if not handler:
        return {
            "error": f"Unknown operation: '{operation}'",
            "valid_operations": sorted(_TRANSFORM_OPS.keys()),
        }
    try:
        return handler(**kwargs)
    except Exception as e:
        logger.error(f"planner_transform({operation}) failed: {e}")
        return {"error": f"planner_transform({operation}) failed: {e}"}


def dispatch_planner_manage(action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a planner_manage action."""
    handler = _MANAGE_ACTIONS.get(action)
    if not handler:
        return {
            "error": f"Unknown action: '{action}'",
            "valid_actions": sorted(_MANAGE_ACTIONS.keys()),
        }
    try:
        return handler(**kwargs)
    except Exception as e:
        logger.error(f"planner_manage({action}) failed: {e}")
        return {"error": f"planner_manage({action}) failed: {e}"}


# ============================================================================
# Registration (mcp_only — no settings parameter)
# ============================================================================

def register_unified_planner_tools(mcp):
    """Register the 3 unified planner MCP tools."""

    @mcp.tool()
    def planner_workflow(
        action: str,
        request: Optional[str] = None,
        context: Optional[str] = None,
    ) -> str:
        """
        Unified planner workflow tool. Replaces 3 individual tools.

        Actions:
        - plan: Create a workflow plan from a natural language request (requires request)
        - analyze: Analyze a user request to understand intent and requirements (requires request)
        - suggest_agents: Suggest which agents could handle a request (requires request)

        Args:
            action: The action to perform (plan, analyze, suggest_agents)
            request: Natural language description of the task
            context: Optional JSON context with additional information (for plan action)

        Returns:
            JSON result with action-specific data
        """
        result = dispatch_planner_workflow(action, **{
            k: v for k, v in {
                "request": request, "context": context,
            }.items() if v is not None
        })
        return json.dumps(result, indent=2, default=str)

    @mcp.tool()
    def planner_transform(
        operation: str,
        plan_json: Optional[str] = None,
    ) -> str:
        """
        Unified planner plan transformation tool. Replaces 3 individual tools.

        Operations:
        - explain: Generate a human-readable explanation of a workflow plan (requires plan_json)
        - optimize: Optimize a plan for parallel execution (requires plan_json)
        - to_workflow_def: Convert plan to Orchestrator workflow definition (requires plan_json)

        Args:
            operation: The operation to perform (explain, optimize, to_workflow_def)
            plan_json: JSON string of a workflow plan (from planner_workflow action='plan')

        Returns:
            JSON result with operation-specific data
        """
        result = dispatch_planner_transform(operation, **{
            k: v for k, v in {
                "plan_json": plan_json,
            }.items() if v is not None
        })
        return json.dumps(result, indent=2, default=str)

    @mcp.tool()
    def planner_manage(
        action: str,
        name: Optional[str] = None,
        agent_type: Optional[str] = None,
        description: Optional[str] = None,
        capabilities: Optional[str] = None,
        input_schema: Optional[str] = None,
        output_schema: Optional[str] = None,
        limit: Optional[int] = None,
        model: Optional[str] = None,
        temperature: Optional[float] = None,
        max_steps: Optional[int] = None,
        enable_parallel: Optional[str] = None,
    ) -> str:
        """
        Unified planner management tool. Replaces 5 individual tools.

        Actions:
        - list_agents: List all available agents
        - register_agent: Register a custom agent (requires name, agent_type, description, capabilities)
        - status: Get planner status (includes agent_counts by type)
        - history: Get planning history (optional limit)
        - configure: Configure planner settings (optional model, temperature, max_steps, enable_parallel)
        - agent_graph: Get vis.js-compatible agent dependency graph (nodes + edges)

        Args:
            action: The action to perform (list_agents, register_agent, status, history, configure)
            name: Agent name (for register_agent)
            agent_type: Agent type (for register_agent)
            description: Agent description (for register_agent)
            capabilities: JSON array of capability names (for register_agent)
            input_schema: JSON schema for agent inputs (for register_agent)
            output_schema: JSON schema for agent outputs (for register_agent)
            limit: Maximum number of plans to return (for history, default 10)
            model: Claude model to use (for configure)
            temperature: Sampling temperature (for configure)
            max_steps: Maximum steps in a plan (for configure)
            enable_parallel: "true" or "false" for parallel optimization (for configure)

        Returns:
            JSON result with action-specific data
        """
        result = dispatch_planner_manage(action, **{
            k: v for k, v in {
                "name": name, "agent_type": agent_type,
                "description": description, "capabilities": capabilities,
                "input_schema": input_schema, "output_schema": output_schema,
                "limit": limit, "model": model, "temperature": temperature,
                "max_steps": max_steps, "enable_parallel": enable_parallel,
            }.items() if v is not None
        })
        return json.dumps(result, indent=2, default=str)

    logger.info("Registered 3 unified planner tools: planner_workflow, planner_transform, planner_manage")
    return ["planner_workflow", "planner_transform", "planner_manage"]
