"""Action-dispatched handler functions for LangGraph MCP tools.

Follows the same pattern as src/agents/unified.py — each MCP tool
maps to a dispatch function that routes by ``action`` parameter.
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


# ============================================================================
# Lazy singletons
# ============================================================================

_memory_store: Dict[str, Any] = {}  # run_id -> ConversationMemory
_streaming_store: Dict[str, Any] = {}  # run_id -> StreamingProgressHandler


def _get_memory(run_id: str):
    """Get or create a ConversationMemory for a run."""
    if run_id not in _memory_store:
        from .conversation_memory import ConversationMemory
        _memory_store[run_id] = ConversationMemory(run_id=run_id)
    return _memory_store[run_id]


def _get_streaming(run_id: str):
    """Get or create a StreamingProgressHandler for a run."""
    if run_id not in _streaming_store:
        from .streaming_handler import StreamingProgressHandler
        _streaming_store[run_id] = StreamingProgressHandler(run_id=run_id)
    return _streaming_store[run_id]


def _run_async(coro):
    """Run an async coroutine from sync context."""
    try:
        loop = asyncio.get_running_loop()
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor() as pool:
            return loop.run_in_executor(pool, lambda: asyncio.run(coro))
    except RuntimeError:
        return asyncio.run(coro)


# ============================================================================
# langgraph_workflow action handlers
# ============================================================================

def _workflow_run(**kwargs) -> Dict[str, Any]:
    """Run a full multi-agent workflow."""
    config_json = kwargs.get("config", "{}")
    try:
        config = json.loads(config_json) if isinstance(config_json, str) else config_json
    except json.JSONDecodeError:
        return {"error": "Invalid JSON in config"}

    human_gates_json = kwargs.get("human_gates", "[]")
    try:
        human_gates = json.loads(human_gates_json) if isinstance(human_gates_json, str) else human_gates_json
    except json.JSONDecodeError:
        human_gates = []

    try:
        from .orchestrator_graph import run_graph

        result = asyncio.run(run_graph(
            config=config,
            human_gates=human_gates if human_gates else None,
        ))
        # Summarise for MCP response
        return {
            "run_id": result.get("run_id", "unknown"),
            "status": result.get("status", "unknown"),
            "completed_phases": result.get("completed_phases", []),
            "total_duration_seconds": result.get("total_duration_seconds", 0),
            "errors": result.get("errors", []),
            "phase_count": len(result.get("phase_results", {})),
        }
    except ImportError:
        return {"error": "langgraph is not installed — use PlatformOrchestrator instead"}
    except Exception as e:
        return {"error": f"Workflow execution failed: {e}"}


def _workflow_resume(**kwargs) -> Dict[str, Any]:
    """Resume a workflow from checkpoint."""
    run_id = kwargs.get("run_id")
    if not run_id:
        return {"error": "run_id is required for 'resume' action"}

    config_json = kwargs.get("config", "{}")
    try:
        config = json.loads(config_json) if isinstance(config_json, str) else config_json
    except json.JSONDecodeError:
        config = {}

    try:
        from .orchestrator_graph import run_graph

        result = asyncio.run(run_graph(
            config=config,
            checkpoint_id=run_id,
        ))
        return {
            "run_id": result.get("run_id", "unknown"),
            "status": result.get("status", "unknown"),
            "completed_phases": result.get("completed_phases", []),
            "resumed_from": run_id,
        }
    except ImportError:
        return {"error": "langgraph is not installed"}
    except Exception as e:
        return {"error": f"Resume failed: {e}"}


def _workflow_status(**kwargs) -> Dict[str, Any]:
    """Get status of a workflow run."""
    run_id = kwargs.get("run_id")
    if not run_id:
        return {"error": "run_id is required for 'status' action"}

    try:
        from .state_persistence import DataBridgeCheckpointSaver
        saver = DataBridgeCheckpointSaver()
        state = saver.get_latest_state(run_id)
        if state:
            return {
                "run_id": run_id,
                "status": state.get("status", "unknown"),
                "completed_phases": state.get("completed_phases", []),
                "errors": state.get("errors", []),
            }
        return {"run_id": run_id, "status": "not_found"}
    except Exception as e:
        return {"error": str(e)}


def _workflow_list_runs(**kwargs) -> Dict[str, Any]:
    """List past workflow runs."""
    limit = int(kwargs.get("limit", 20))
    try:
        from .state_persistence import DataBridgeCheckpointSaver
        saver = DataBridgeCheckpointSaver()
        runs = list(saver.list(limit=limit))
        return {
            "runs": [
                {
                    "thread_id": r.config.get("configurable", {}).get("thread_id", ""),
                    "checkpoint_id": r.config.get("configurable", {}).get("checkpoint_id", ""),
                }
                for r in runs
            ],
            "count": len(runs),
        }
    except Exception as e:
        return {"error": str(e)}


# ============================================================================
# langgraph_agent action handlers
# ============================================================================

def _agent_invoke(**kwargs) -> Dict[str, Any]:
    """Invoke a single specialist agent."""
    agent_name = kwargs.get("agent_name")
    if not agent_name:
        return {"error": "agent_name is required for 'invoke' action"}

    config_json = kwargs.get("config", "{}")
    try:
        config = json.loads(config_json) if isinstance(config_json, str) else config_json
    except json.JSONDecodeError:
        config = {}

    try:
        from .orchestrator_graph import _get_agents
        from .state_schema import create_initial_state

        agents = _get_agents()
        if agent_name not in agents:
            return {
                "error": f"Unknown agent: {agent_name}",
                "available": list(agents.keys()),
            }

        agent = agents[agent_name]()
        state = create_initial_state(config)
        state["current_phase"] = config.get("phase", agent.phase_name)

        result = asyncio.run(agent(state))
        return {
            "agent": agent_name,
            "phase": result.get("current_phase", ""),
            "status": "completed",
            "result": result.get("context", {}),
        }
    except ImportError:
        return {"error": "langgraph agents not available"}
    except Exception as e:
        return {"error": f"Agent invocation failed: {e}"}


def _agent_list(**kwargs) -> Dict[str, Any]:
    """List available specialist agents."""
    try:
        from .orchestrator_graph import list_available_agents
        agents = list_available_agents()
        return {"agents": agents, "count": len(agents)}
    except Exception as e:
        return {"error": str(e)}


def _agent_capabilities(**kwargs) -> Dict[str, Any]:
    """Get capabilities of a specific agent."""
    agent_name = kwargs.get("agent_name")
    if not agent_name:
        return {"error": "agent_name is required"}

    try:
        from .orchestrator_graph import _get_agents
        agents = _get_agents()
        if agent_name not in agents:
            return {"error": f"Unknown agent: {agent_name}"}
        instance = agents[agent_name]()
        return instance.get_capabilities()
    except Exception as e:
        return {"error": str(e)}


# ============================================================================
# langgraph_manage action handlers
# ============================================================================

def _manage_checkpoint(**kwargs) -> Dict[str, Any]:
    """Save or load a checkpoint."""
    run_id = kwargs.get("run_id")
    if not run_id:
        return {"error": "run_id is required"}

    operation = kwargs.get("operation", "load")
    try:
        from .state_persistence import DataBridgeCheckpointSaver
        saver = DataBridgeCheckpointSaver()

        if operation == "load":
            state = saver.get_latest_state(run_id)
            return {"run_id": run_id, "state": state, "found": state is not None}
        elif operation == "save":
            state_json = kwargs.get("state", "{}")
            state = json.loads(state_json) if isinstance(state_json, str) else state_json
            config = {"configurable": {"thread_id": run_id}}
            result = saver.put(config, state, {})
            return {"status": "saved", **result}
        else:
            return {"error": f"Unknown operation: {operation}"}
    except Exception as e:
        return {"error": str(e)}


def _manage_clear(**kwargs) -> Dict[str, Any]:
    """Clear run state."""
    run_id = kwargs.get("run_id")
    if not run_id:
        return {"error": "run_id is required"}

    try:
        from .state_persistence import DataBridgeCheckpointSaver
        saver = DataBridgeCheckpointSaver()
        cleared = saver.clear(run_id)
        return {"run_id": run_id, "cleared": cleared}
    except Exception as e:
        return {"error": str(e)}


def _manage_configure(**kwargs) -> Dict[str, Any]:
    """Update runtime configuration."""
    return {
        "status": "ok",
        "note": "Runtime configuration is set per-run via config parameter",
        "example_config": {
            "source_type": "csv",
            "csv_folder": "data/test",
            "erp": "enertia",
            "skip_phases": ["wright", "dbt"],
        },
    }


# ============================================================================
# langgraph_cortex action handlers
# ============================================================================

def _cortex_complete(**kwargs) -> Dict[str, Any]:
    """Run Cortex LLM completion via agent."""
    prompt = kwargs.get("prompt")
    if not prompt:
        return {"error": "prompt is required"}

    config = {
        "cortex_action": "complete",
        "cortex_prompt": prompt,
        "cortex_model": kwargs.get("model", "mistral-large"),
    }
    return _agent_invoke(agent_name="cortex_llm_agent", config=json.dumps(config))


def _cortex_analyst(**kwargs) -> Dict[str, Any]:
    """Run Cortex Analyst query via agent."""
    question = kwargs.get("question")
    model_name = kwargs.get("model_name")
    if not question or not model_name:
        return {"error": "question and model_name are required"}

    config = {
        "analyst_action": "ask",
        "analyst_question": question,
        "analyst_model": model_name,
    }
    return _agent_invoke(agent_name="cortex_analyst_agent", config=json.dumps(config))


def _cortex_search(**kwargs) -> Dict[str, Any]:
    """Run Cortex Search via agent."""
    query = kwargs.get("query")
    if not query:
        return {"error": "query is required"}

    config = {
        "search_action": kwargs.get("action", "search"),
        "search_query": query,
    }
    return _agent_invoke(agent_name="cortex_search_agent", config=json.dumps(config))


# ============================================================================
# langgraph_streaming action handlers
# ============================================================================

def _streaming_subscribe(**kwargs) -> Dict[str, Any]:
    """Subscribe to progress events for a run."""
    run_id = kwargs.get("run_id")
    if not run_id:
        return {"error": "run_id is required"}

    handler = _get_streaming(run_id)
    return {
        "run_id": run_id,
        "status": "subscribed",
        "event_count": handler.event_count,
    }


def _streaming_history(**kwargs) -> Dict[str, Any]:
    """Get event history for a run."""
    run_id = kwargs.get("run_id")
    if not run_id:
        return {"error": "run_id is required"}

    limit = int(kwargs.get("limit", 100))
    handler = _get_streaming(run_id)
    events = handler.get_history(limit=limit)
    return {
        "run_id": run_id,
        "events": events,
        "count": len(events),
    }


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

_WORKFLOW_ACTIONS = {
    "run": _workflow_run,
    "resume": _workflow_resume,
    "status": _workflow_status,
    "list_runs": _workflow_list_runs,
}

_AGENT_ACTIONS = {
    "invoke": _agent_invoke,
    "list": _agent_list,
    "capabilities": _agent_capabilities,
}

_MANAGE_ACTIONS = {
    "checkpoint": _manage_checkpoint,
    "clear": _manage_clear,
    "configure": _manage_configure,
}

_CORTEX_ACTIONS = {
    "complete": _cortex_complete,
    "analyst": _cortex_analyst,
    "search": _cortex_search,
}

_STREAMING_ACTIONS = {
    "subscribe": _streaming_subscribe,
    "history": _streaming_history,
}


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

def dispatch_langgraph_workflow(action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a langgraph_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("langgraph_workflow(%s) failed: %s", action, e)
        return {"error": f"langgraph_workflow({action}) failed: {e}"}


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


def dispatch_langgraph_manage(action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a langgraph_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("langgraph_manage(%s) failed: %s", action, e)
        return {"error": f"langgraph_manage({action}) failed: {e}"}


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


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