from dataclasses import dataclass, field
from typing import Literal, NotRequired, TypedDict

from langgraph.runtime import get_runtime

from app.agent.data_tools import DataContext
from app.schemas.actions import Action
from app.schemas.base import JsonObject
from app.schemas.clarification import Clarification
from app.schemas.clarification_reply import ClarificationReply
from app.schemas.decision import DecisionResponse, Route
from app.schemas.input_context import InputContext


class QueryState(TypedDict):
    conversation_id: NotRequired[str]  # 当前聊天会话 ID
    thread_id: NotRequired[str]  # 当前任务 ID
    message_id: NotRequired[str]  # 本次触发执行的用户消息 ID，澄清恢复时更新
    request: JsonObject  # 原始请求，保留提交时的内容
    input_context: NotRequired[InputContext]  # 校验后的上下文，包含最新大屏和澄清历史
    intent_decision: NotRequired[DecisionResponse | None]  # Jev 原始意图判断及分数
    intent: NotRequired[Route | None]  # 门控后的去向
    actions: NotRequired[
        list[Action] | None
    ]  # screen agent 各次成功提交的动作，按执行顺序拼接
    clarification: NotRequired[Clarification | None]  # 当前澄清问题及选项
    clarification_reply: NotRequired[
        ClarificationReply | None
    ]  # 待合并的用户回答或取消指令
    answer: NotRequired[str | None]  # 问答回答，或 agent 给用户的说明
    suggestions: NotRequired[list[str] | None]  # 可以一键发送的后续请求，前端显示成按钮
    # 轮次记录，字段同 TurnChange、TurnQuery；用普通 dict，检查点不用为它们扩类型白名单
    changes: NotRequired[list[JsonObject] | None]  # 交付的修改，按组件合并
    queries: NotRequired[list[JsonObject] | None]  # 问答里成功执行的查询
    status: NotRequired[
        Literal["running", "waiting_for_input", "completed", "cancelled", "failed"]
    ]  # 后端处理、等待补充、处理完成、用户取消或失败，完成不代表前端已执行


@dataclass(frozen=True)
class RunContext:
    """图的运行上下文：本次运行要用、但不进检查点的依赖，每次启动和恢复时由调用方传入"""

    # 本会话的数据文件；状态里的 input_context.data_files 只有摘要，列概况和路径在这里
    data: DataContext = field(default_factory=DataContext)


def run_data() -> DataContext:
    """当前运行可用的数据表；单独调用节点（不在图里执行）或调用方没传 context 时为空"""
    try:
        runtime = get_runtime(RunContext)
    except RuntimeError:
        return DataContext()
    context = runtime.context if runtime is not None else None
    return context.data if context is not None else DataContext()