from __future__ import annotations

from dataclasses import dataclass, field
from typing import Annotated, Any, Literal, Mapping

from typing_extensions import TypedDict

from agent_search.runtime.state import (
    CitationRowsByIndexChannel,
    DecompositionSubQuestionsChannel,
    StageSnapshotsChannel,
    SubQAChannel,
    SubQuestionArtifactsChannel,
    to_rag_state,
)
from schemas import GraphRunMetadata, RuntimeAgentRunRequest

RuntimeHitlMode = Literal["disabled", "strategic", "custom"]
RuntimeHitlCheckpointStage = Literal["subquestions", "query_expansion"]

HITL_MODE_DISABLED: RuntimeHitlMode = "disabled"
HITL_MODE_STRATEGIC: RuntimeHitlMode = "strategic"
HITL_MODE_CUSTOM: RuntimeHitlMode = "custom"

HITL_STAGE_SUBQUESTIONS: RuntimeHitlCheckpointStage = "subquestions"
HITL_STAGE_QUERY_EXPANSION: RuntimeHitlCheckpointStage = "query_expansion"
STRATEGIC_HITL_CHECKPOINT_STAGES: tuple[RuntimeHitlCheckpointStage, ...] = (
    HITL_STAGE_SUBQUESTIONS,
    HITL_STAGE_QUERY_EXPANSION,
)


def _merge_stable_main_question(current: str, update: str) -> str:
    if current and update and current != update:
        raise ValueError(f"main_question must remain stable across graph lanes: {current!r} != {update!r}")
    return update or current


def _merge_stable_optional_text(current: str, update: str) -> str:
    if current and update and current != update:
        raise ValueError(f"text channel must remain stable across graph lanes: {current!r} != {update!r}")
    return update or current


def _merge_stable_hitl_mode(current: RuntimeHitlMode, update: RuntimeHitlMode) -> RuntimeHitlMode:
    if current and update and current != update:
        raise ValueError(f"hitl_mode must remain stable across graph lanes: {current!r} != {update!r}")
    return update or current


def _merge_stable_run_metadata(current: GraphRunMetadata, update: GraphRunMetadata) -> GraphRunMetadata:
    if current.model_dump(mode="json") != update.model_dump(mode="json"):
        raise ValueError("run_metadata must remain stable across graph lanes.")
    return update.model_copy(deep=True)


def _merge_stable_initial_search_context(
    current: list[dict[str, Any]],
    update: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    if current and update and current != update:
        raise ValueError("initial_search_context must remain stable across graph lanes.")
    return [dict(item) for item in (update or current)]


def _merge_stable_hitl_checkpoint_stages(
    current: tuple[RuntimeHitlCheckpointStage, ...],
    update: tuple[RuntimeHitlCheckpointStage, ...],
) -> tuple[RuntimeHitlCheckpointStage, ...]:
    if current and update and current != update:
        raise ValueError("hitl_checkpoint_stages must remain stable across graph lanes.")
    return update or current


def _resolve_hitl_policy(
    payload: RuntimeAgentRunRequest,
) -> tuple[RuntimeHitlMode, tuple[RuntimeHitlCheckpointStage, ...]]:
    hitl = payload.controls.hitl if payload.controls else None
    if hitl is None or not hitl.enabled:
        return HITL_MODE_DISABLED, ()

    explicit_stages: list[RuntimeHitlCheckpointStage] = []
    if hitl.subquestions and hitl.subquestions.enabled:
        explicit_stages.append(HITL_STAGE_SUBQUESTIONS)
    if hitl.query_expansion and hitl.query_expansion.enabled:
        explicit_stages.append(HITL_STAGE_QUERY_EXPANSION)
    if explicit_stages:
        return HITL_MODE_CUSTOM, tuple(explicit_stages)
    return HITL_MODE_STRATEGIC, STRATEGIC_HITL_CHECKPOINT_STAGES


@dataclass(slots=True)
class RuntimeGraphContext:
    payload: RuntimeAgentRunRequest
    model: Any | None = None
    vector_store: Any | None = None
    callbacks: list[Any] = field(default_factory=list)
    langfuse_callback: Any | None = None
    initial_search_context: list[dict[str, Any]] = field(default_factory=list)


def to_runtime_graph_state(
    payload: RuntimeAgentRunRequest,
    *,
    run_metadata: GraphRunMetadata,
    initial_search_context: list[dict[str, Any]] | None = None,
) -> "RuntimeGraphState":
    hitl_mode, hitl_checkpoint_stages = _resolve_hitl_policy(payload)
    base_state = to_rag_state(
        {
            "main_question": payload.query,
            "decomposition_sub_questions": [],
            "sub_question_artifacts": [],
            "final_answer": "",
            "citation_rows_by_index": {},
            "run_metadata": run_metadata,
            "sub_qa": [],
            "output": "",
            "stage_snapshots": [],
        }
    )
    return RuntimeGraphState(
        **base_state,
        initial_search_context=list(initial_search_context or []),
        hitl_mode=hitl_mode,
        hitl_checkpoint_stages=hitl_checkpoint_stages,
    )


class RuntimeGraphState(TypedDict):
    main_question: Annotated[str, _merge_stable_main_question]
    decomposition_sub_questions: DecompositionSubQuestionsChannel
    sub_question_artifacts: SubQuestionArtifactsChannel
    final_answer: Annotated[str, _merge_stable_optional_text]
    citation_rows_by_index: CitationRowsByIndexChannel
    run_metadata: Annotated[GraphRunMetadata, _merge_stable_run_metadata]
    sub_qa: SubQAChannel
    output: Annotated[str, _merge_stable_optional_text]
    stage_snapshots: StageSnapshotsChannel
    initial_search_context: Annotated[list[dict[str, Any]], _merge_stable_initial_search_context]
    hitl_mode: Annotated[RuntimeHitlMode, _merge_stable_hitl_mode]
    hitl_checkpoint_stages: Annotated[
        tuple[RuntimeHitlCheckpointStage, ...], _merge_stable_hitl_checkpoint_stages
    ]


RuntimeGraphStateMapping = Mapping[str, Any]


__all__ = [
    "RuntimeGraphContext",
    "RuntimeHitlCheckpointStage",
    "RuntimeHitlMode",
    "RuntimeGraphState",
    "RuntimeGraphStateMapping",
    "HITL_MODE_CUSTOM",
    "HITL_MODE_DISABLED",
    "HITL_MODE_STRATEGIC",
    "HITL_STAGE_QUERY_EXPANSION",
    "HITL_STAGE_SUBQUESTIONS",
    "STRATEGIC_HITL_CHECKPOINT_STAGES",
    "to_runtime_graph_state",
]
