from __future__ import annotations

from typing import Literal

# 会话级状态集（9 个）：8 个设计态 + cancelled（T3 架构审查整改）
SessionState = Literal[
    "uploading",
    "parsing",
    "awaiting_parse_confirm",
    "impact_running",
    "awaiting_impact_confirm",
    "writing",
    "qa",
    "done",
    "cancelled",
]

STATES: frozenset[str] = frozenset({
    "uploading",
    "parsing",
    "awaiting_parse_confirm",
    "impact_running",
    "awaiting_impact_confirm",
    "writing",
    "qa",
    "done",
    "cancelled",
})

# 合法转移白名单（不含 cancelled 的动态 resume 转移，见 SessionStateMachine.resume）
_TRANSITIONS: dict[str, frozenset[str]] = {
    "uploading": frozenset({"parsing"}),
    "parsing": frozenset({"awaiting_parse_confirm", "cancelled"}),
    "awaiting_parse_confirm": frozenset({"impact_running", "parsing", "cancelled"}),
    "impact_running": frozenset({"awaiting_impact_confirm", "cancelled"}),
    "awaiting_impact_confirm": frozenset({"writing", "impact_running", "awaiting_parse_confirm", "cancelled"}),
    "writing": frozenset({"qa", "awaiting_impact_confirm", "cancelled"}),
    "qa": frozenset({"done", "writing", "cancelled"}),
    "done": frozenset(),
    "cancelled": frozenset(),
}

# 可被取消的（非终态、非人工等待确认态之外全部执行中；done 不可取消）
_CANCELLABLE: frozenset[str] = frozenset({
    "parsing",
    "impact_running",
    "writing",
    "qa",
})


class StateTransitionError(Exception):
    """非法状态转移（对应 api-design §7 STATE_TRANSITION_INVALID 409）。"""


class SessionStateMachine:
    """会话级流程状态机：白名单转移 + cancelled/resume（T3）。

    取消（cancel）从任意执行中状态进入 cancelled 终态，并记录中断前状态
    （cancelled_from）；resume 从 cancelled 回到中断前状态，恢复后继续
    正常白名单流转。人工等待确认态（awaiting_*）与 done 不可取消。
    """

    def __init__(self, initial: str = "uploading") -> None:
        if initial not in STATES:
            raise StateTransitionError(f"未知初始状态: {initial}")
        self._state: str = initial
        self._cancelled_from: str | None = None

    @property
    def state(self) -> str:
        return self._state

    @property
    def cancelled_from(self) -> str | None:
        """取消前的中断状态（resume 目标）；仅 cancelled 态非 None。"""
        return self._cancelled_from

    def transition(self, target: str) -> str:
        """按白名单推进状态机；非法转移抛 StateTransitionError。"""
        if target not in STATES:
            raise StateTransitionError(f"未知目标状态: {target}")
        allowed = _TRANSITIONS[self._state]
        if target not in allowed:
            raise StateTransitionError(
                f"非法状态转移: {self._state} → {target}（白名单外）"
            )
        self._state = target
        return self._state

    def cancel(self) -> str:
        """取消当前执行：进入 cancelled 终态，记录中断前状态。"""
        if self._state not in _CANCELLABLE:
            raise StateTransitionError(
                f"当前状态不可取消: {self._state}（仅执行中状态可取消）"
            )
        self._cancelled_from = self._state
        self._state = "cancelled"
        return self._state

    def resume(self) -> str:
        """从 cancelled 恢复：回到中断前状态（cancelled_from）。"""
        if self._state != "cancelled":
            raise StateTransitionError(
                f"仅 cancelled 状态可 resume（当前: {self._state}）"
            )
        assert self._cancelled_from is not None  # cancelled 态必有记录
        self._state = self._cancelled_from
        self._cancelled_from = None
        return self._state