"""Main StateGraph wiring — nodes, edges, checkpoints.

Builds a LangGraph StateGraph that maps PHASE_ORDER phases to specialist
agent nodes with conditional edges, parallel branches, and human gates.

When ``langgraph`` is not installed, ``build_workflow_graph()`` raises
ImportError — callers should check ``LANGGRAPH_AVAILABLE`` first.
"""

from __future__ import annotations

import asyncio
import logging
import time
import uuid
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

from .state_schema import WorkflowState, create_initial_state

logger = logging.getLogger(__name__)

# Phase ordering (aligned with src/workflows/orchestrator.py:118-130)
PHASE_ORDER = [
    "load_metadata",
    "relationship_discovery",
    "group_tables",
    "hierarchy",
    "ai_classify",
    "detect_dimensions",
    "wright",
    "dbt",
    "quality",
    "observability",
    "artifact_bundle",
]

# Default phases that can run in parallel
PARALLEL_PHASES = {
    "dbt": "quality",  # dbt and quality can run concurrently after wright
}

# Default human gate points
DEFAULT_HUMAN_GATES = ["hierarchy"]


def _get_agents():
    """Lazy-import agent classes to avoid circular imports."""
    from .data_modeling_agent import DataModelingAgent
    from .hierarchy_agent import HierarchyAgent
    from .dbt_agent import DbtAgent
    from .snowflake_agent import SnowflakeAgent
    from .qa_agent import QAAgent
    from .cortex_llm_agent import CortexLLMAgent
    from .cortex_analyst_agent import CortexAnalystAgent
    from .cortex_search_agent import CortexSearchAgent

    return {
        "data_modeling_agent": DataModelingAgent,
        "hierarchy_agent": HierarchyAgent,
        "dbt_agent": DbtAgent,
        "snowflake_agent": SnowflakeAgent,
        "qa_agent": QAAgent,
        "cortex_llm_agent": CortexLLMAgent,
        "cortex_analyst_agent": CortexAnalystAgent,
        "cortex_search_agent": CortexSearchAgent,
    }


# Phase -> Agent class mapping
PHASE_AGENT_MAP = {
    "load_metadata": "data_modeling_agent",
    "relationship_discovery": "data_modeling_agent",
    "group_tables": "data_modeling_agent",
    "hierarchy": "hierarchy_agent",
    "ai_classify": "data_modeling_agent",
    "detect_dimensions": "data_modeling_agent",
    "wright": "snowflake_agent",
    "dbt": "dbt_agent",
    "quality": "qa_agent",
    "observability": "qa_agent",
    "artifact_bundle": "data_modeling_agent",
}


def build_workflow_graph(
    config: Optional[Dict[str, Any]] = None,
    human_gates: Optional[List[str]] = None,
    checkpoint_saver: Optional[Any] = None,
    progress_callback: Optional[Callable] = None,
):
    """Build a LangGraph StateGraph for the multi-agent workflow.

    Args:
        config: Orchestrator configuration dict.
        human_gates: List of phase names that require human approval.
        checkpoint_saver: LangGraph checkpoint saver for resume support.
        progress_callback: Callback for streaming progress events.

    Returns:
        Compiled LangGraph StateGraph.

    Raises:
        ImportError: If langgraph is not installed.
    """
    from langgraph.graph import StateGraph, END

    config = config or {}
    gates = set(human_gates or DEFAULT_HUMAN_GATES)
    skip_phases = set(config.get("skip_phases", []))

    # Instantiate agents
    agent_classes = _get_agents()
    agent_instances = {}
    for agent_name, agent_cls in agent_classes.items():
        agent_instances[agent_name] = agent_cls(progress_callback=progress_callback)

    # Build graph
    graph = StateGraph(dict)  # Use dict as state type for flexibility

    # --- Add nodes ---
    for phase in PHASE_ORDER:
        if phase in skip_phases:
            continue

        agent_key = PHASE_AGENT_MAP.get(phase)
        if not agent_key or agent_key not in agent_instances:
            continue

        agent = agent_instances[agent_key]

        # Create a phase-specific node function
        def make_node(p, a):
            async def node_fn(state):
                state = {**state, "current_phase": p}
                return await a(state)
            node_fn.__name__ = f"node_{p}"
            return node_fn

        graph.add_node(phase, make_node(phase, agent))

    # --- Add edges ---
    active_phases = [p for p in PHASE_ORDER if p not in skip_phases]

    if not active_phases:
        # Empty graph — just start->end
        graph.add_node("noop", lambda state: state)
        graph.set_entry_point("noop")
        graph.add_edge("noop", END)
    else:
        graph.set_entry_point(active_phases[0])

        for i, phase in enumerate(active_phases):
            if i == len(active_phases) - 1:
                # Last phase -> END
                graph.add_edge(phase, END)
            else:
                next_phase = active_phases[i + 1]

                # Conditional edge after load_metadata: if no tables, go to END
                if phase == "load_metadata":
                    def load_meta_router(state):
                        ctx = state.get("context", {})
                        meta = ctx.get("load_metadata", {})
                        tables = meta.get("tables", [])
                        if not tables and not meta.get("error"):
                            return "end"
                        return "continue"

                    graph.add_conditional_edges(
                        phase,
                        load_meta_router,
                        {"continue": next_phase, "end": END},
                    )
                # Human gate: conditional edge for approval
                elif phase in gates:
                    def make_gate_router(next_p):
                        def gate_router(state):
                            if state.get("human_approval_needed"):
                                return "wait"
                            return "continue"
                        return gate_router

                    graph.add_conditional_edges(
                        phase,
                        make_gate_router(next_phase),
                        {"continue": next_phase, "wait": END},
                    )
                else:
                    graph.add_edge(phase, next_phase)

    # Compile with optional checkpoint saver
    compile_kwargs = {}
    if checkpoint_saver:
        compile_kwargs["checkpointer"] = checkpoint_saver

    return graph.compile(**compile_kwargs)


async def run_graph(
    config: Dict[str, Any],
    checkpoint_id: Optional[str] = None,
    human_gates: Optional[List[str]] = None,
    progress_callback: Optional[Callable] = None,
) -> Dict[str, Any]:
    """Execute the multi-agent workflow graph.

    Args:
        config: Orchestrator configuration dict.
        checkpoint_id: Optional checkpoint ID to resume from.
        human_gates: Phases requiring human approval.
        progress_callback: Streaming progress callback.

    Returns:
        Final workflow state dict.
    """
    from .state_persistence import DataBridgeCheckpointSaver

    # Create checkpoint saver
    try:
        saver = DataBridgeCheckpointSaver()
    except Exception:
        saver = None

    # Build graph
    graph = build_workflow_graph(
        config=config,
        human_gates=human_gates,
        checkpoint_saver=saver,
        progress_callback=progress_callback,
    )

    # Create initial state
    initial_state = create_initial_state(config)
    initial_state["status"] = "running"

    # Execute
    run_config = {"configurable": {"thread_id": initial_state["run_id"]}}

    if checkpoint_id:
        run_config["configurable"]["checkpoint_id"] = checkpoint_id

    t0 = time.time()
    try:
        final_state = await graph.ainvoke(dict(initial_state), config=run_config)
        final_state["status"] = _compute_status(final_state)
    except Exception as e:
        logger.error("Graph execution failed: %s", e)
        final_state = dict(initial_state)
        final_state["status"] = "failed"
        final_state["errors"] = [{"phase": "graph", "error": str(e)}]

    final_state["total_duration_seconds"] = round(time.time() - t0, 2)
    final_state["ended_at"] = datetime.utcnow().isoformat() + "Z"

    return final_state


def _compute_status(state: Dict[str, Any]) -> str:
    """Compute overall run status from phase results."""
    results = state.get("phase_results", {})
    if not results:
        return "completed"

    statuses = [r.get("status", "unknown") for r in results.values()]
    if all(s in ("completed", "skipped") for s in statuses):
        return "completed"
    if any(s == "error" for s in statuses):
        return "partial"
    return "running"


def list_available_agents() -> List[Dict[str, Any]]:
    """List all available specialist agents and their capabilities."""
    agents = _get_agents()
    result = []
    for name, cls in agents.items():
        instance = cls()
        result.append(instance.get_capabilities())
    return result
