from __future__ import annotations
"""完整多角色工作流主图：编排从需求解析→资源查询→候选筛选→地图研究→路线评估→行程规划→校验→评审→最终定稿的全部节点与条件边。"""

from langchain_core.tools import BaseTool
from langgraph.graph import (
    END,
    START,
    StateGraph,
)

from app.agents.itinerary_agent import (
    ItineraryPlannerAgent,
)
from app.agents.map_agent import (
    MapResearchAgent,
)
from app.agents.requirement_agent import (
    RequirementAgent,
)
from app.agents.resource_agent import (
    ResourceSearchAgent,
)
from app.graph.map_nodes import (
    make_map_research_node,
)
from app.graph.nodes import (
    clarification_node,
    error_node,
    make_parse_request_node,
    route_after_parse,
)
from app.graph.planner_nodes import (
    make_itinerary_planner_node,
    route_after_route_evaluation,
)
from app.graph.resource_nodes import (
    make_resource_search_node,
)
from app.graph.route_nodes import (
    make_route_evaluation_node,
    route_after_map_research,
    route_after_selection,
)
from app.graph.selection_nodes import (
    make_candidate_selection_node,
    route_after_resources_to_selection,
)
from app.graph.state import TravelState
from app.mcp_client import load_mcp_tools
from app.graph.validation_nodes import (
    make_plan_validation_node,
    route_after_planning,
)
from app.agents.reviewer_agent import (
    TripReviewerAgent,
)
from app.graph.review_nodes import (
    make_trip_reviewer_node,
    route_after_validation,
)
from app.graph.final_nodes import (
    make_finalization_node,
)

from app.graph.supervisor_nodes import (
    make_supervisor_node,
    route_after_supervisor,
)


def find_tool(
    tools: list[BaseTool],
    *names: str,
) -> BaseTool | None:
    """按名称查找已经加载的MCP工具。"""

    for name in names:
        for tool in tools:
            if tool.name == name:
                return tool

    return None


async def build_planner_graph():
    """构建直到行程规划完成的多角色工作流。"""

    tool_bundle = await load_mcp_tools()

    return_flight_tool = find_tool(
        tool_bundle.travel_tools,
        "search_return_flights",
    )

    distance_tool = find_tool(
        tool_bundle.amap_tools,
        "maps_distance",
    )

    around_search_tool = find_tool(
        tool_bundle.amap_tools,
        "maps_around_search",
        "maps_search_around",
    )

    if return_flight_tool is None:
        raise RuntimeError(
            "没有加载search_return_flights工具。"
        )

    if distance_tool is None:
        raise RuntimeError(
            "没有加载maps_distance工具。"
        )

    requirement_agent = RequirementAgent()

    resource_agent = ResourceSearchAgent(
        tools=tool_bundle.travel_tools,
    )

    map_agent = MapResearchAgent(
        tools=tool_bundle.amap_tools,
    )

    planner_agent = ItineraryPlannerAgent()
    reviewer_agent = TripReviewerAgent()

    builder = StateGraph(TravelState)

    builder.add_node(
        "parse_request",
        make_parse_request_node(
            requirement_agent
        ),
    )

    builder.add_node(
        "clarify",
        clarification_node,
    )

    builder.add_node(
        "error",
        error_node,
    )

    builder.add_node(
        "search_resources",
        make_resource_search_node(
            resource_agent
        ),
    )

    builder.add_node(
        "select_candidates",
        make_candidate_selection_node(
            return_flight_tool,
            max_outbound_queries=2,
            max_returns_per_outbound=5,
        ),
    )

    builder.add_node(
        "research_destination",
        make_map_research_node(
            map_agent
        ),
    )

    builder.add_node(
        "evaluate_routes",
        make_route_evaluation_node(
            distance_tool=distance_tool,
            around_search_tool=(
                around_search_tool
            ),
            max_hotels=3,
            max_attractions=4,
        ),
    )

    builder.add_node(
        "plan_itinerary",
        make_itinerary_planner_node(
            planner_agent
        ),
    )

    builder.add_node(
        "validate_plan",
        make_plan_validation_node(),
    )

    builder.add_node(
    "review_plan",
    make_trip_reviewer_node(
        reviewer_agent
    ),
    )

    builder.add_node(
    "supervise",
    make_supervisor_node(
        max_planning_attempts=3,
    ),
)

    builder.add_node(
        "finalize_response",
        make_finalization_node(),
    )

    builder.add_edge(
        START,
        "parse_request",
    )

    builder.add_conditional_edges(
        "parse_request",
        route_after_parse,
        {
            "clarify": "clarify",
            "error": "error",
            "ready": "search_resources",
        },
    )

    builder.add_edge("clarify", END)
    builder.add_edge("error", END)

    builder.add_conditional_edges(
        "search_resources",
        route_after_resources_to_selection,
        {
            "select_candidates": (
                "select_candidates"
            ),
            "end": END,
        },
    )

    builder.add_conditional_edges(
        "select_candidates",
        route_after_selection,
        {
            "research_destination": (
                "research_destination"
            ),
            "end": END,
        },
    )

    builder.add_conditional_edges(
        "research_destination",
        route_after_map_research,
        {
            "evaluate_routes": (
                "evaluate_routes"
            ),
            "end": END,
        },
    )

    builder.add_conditional_edges(
        "evaluate_routes",
        route_after_route_evaluation,
        {
            "plan_itinerary": (
                "plan_itinerary"
            ),
            "end": END,
        },
    )

    builder.add_conditional_edges(
        "plan_itinerary",
        route_after_planning,
        {
            "validate_plan": "validate_plan",
            "end": END,
        },
    )

    builder.add_conditional_edges(
    "validate_plan",
    route_after_validation,
    {
        "review_plan": "review_plan",
        "end": END,
    },
    )

    builder.add_edge(
    "review_plan",
    "supervise",
    )

    builder.add_conditional_edges(
        "supervise",
        route_after_supervisor,
        {
            "replan": "plan_itinerary",
            "finalize": "finalize_response",
            "finalize_with_risks": (
                "finalize_response"
            ),
            "terminate": END,
        },
        )

    builder.add_edge(
        "finalize_response",
        END,
    )

    return builder.compile()