from __future__ import annotations

from collections.abc import Sequence

from langgraph.types import Send

from agent_search.runtime.graph.state import (
    HITL_STAGE_QUERY_EXPANSION,
    HITL_STAGE_SUBQUESTIONS,
    RuntimeGraphState,
    RuntimeHitlCheckpointStage,
)


def _hitl_stage_enabled(state: RuntimeGraphState, stage: RuntimeHitlCheckpointStage) -> bool:
    return stage in state.get("hitl_checkpoint_stages", ())


def route_post_decompose(state: RuntimeGraphState) -> str | Sequence[Send]:
    if not state["decomposition_sub_questions"]:
        return "synthesize"
    if _hitl_stage_enabled(state, HITL_STAGE_SUBQUESTIONS):
        return "subquestion_checkpoint"
    return route_subquestion_lanes(state)


def route_post_expand(state: RuntimeGraphState) -> str:
    if _hitl_stage_enabled(state, HITL_STAGE_QUERY_EXPANSION):
        return "query_expansion_checkpoint"
    return "search"


def route_subquestion_lanes(state: RuntimeGraphState) -> Sequence[Send]:
    return [
        Send(
            "expand",
            {
                "main_question": state["main_question"],
                "decomposition_sub_questions": [sub_question],
                "sub_question_artifacts": [],
                "final_answer": "",
                "citation_rows_by_index": {},
                "run_metadata": state["run_metadata"],
                "sub_qa": [],
                "output": "",
                "stage_snapshots": [],
                "hitl_mode": state["hitl_mode"],
                "hitl_checkpoint_stages": state["hitl_checkpoint_stages"],
            },
        )
        for sub_question in state["decomposition_sub_questions"]
    ]


__all__ = ["route_post_decompose", "route_post_expand", "route_subquestion_lanes"]
