"""A9 · Agent Planning & Reasoning v1.1 demo.

最小 STRIPS 风格动作模型 + BFS 规划 + 执行前权限检查 +
反馈驱动重规划。标准库 only；不访问真实网页 / LLM。
运行：python3 A9-Agent-Planning-and-Reasoning-v1.1-demo.py
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from typing import Dict, FrozenSet, Iterable, List, Optional, Set, Tuple


State = FrozenSet[str]


@dataclass(frozen=True)
class Action:
    name: str
    requires: FrozenSet[str]
    adds: FrozenSet[str]
    cost: int = 1
    removes: FrozenSet[str] = frozenset()
    permission: Optional[str] = None  # required grant name, if any


ACTIONS: Tuple[Action, ...] = (
    Action("search", frozenset(), frozenset({"sources"})),
    Action("read", frozenset({"sources"}), frozenset({"evidence"})),
    Action(
        "draft",
        frozenset({"evidence"}),
        frozenset({"draft"}),
        removes=frozenset(),  # draft does not remove evidence by itself
    ),
    Action("verify", frozenset({"draft", "evidence"}), frozenset({"verified"})),
    Action(
        "publish",
        frozenset({"verified", "approved"}),
        frozenset({"published"}),
        permission="publish",
    ),
)


def apply(state: State, action: Action) -> State:
    if not action.requires.issubset(state):
        raise ValueError(f"precondition failed: {action.name}")
    return frozenset((state - action.removes) | action.adds)


def applicable(state: State, action: Action) -> bool:
    return action.requires.issubset(state)


def bfs_plan(
    start: State,
    goal: State,
    actions: Iterable[Action] = ACTIONS,
    *,
    max_expansions: int = 200,
) -> Optional[List[Action]]:
    """单位成本 BFS：最少步数计划；max_expansions 限制搜索扩展。"""
    if goal.issubset(start):
        return []
    queue: deque[Tuple[State, List[Action]]] = deque([(start, [])])
    seen: Set[State] = {start}
    expansions = 0
    action_list = list(actions)
    while queue:
        if expansions >= max_expansions:
            return None
        state, path = queue.popleft()
        expansions += 1
        for action in action_list:
            if not applicable(state, action):
                continue
            nxt = apply(state, action)
            if nxt in seen:
                continue
            new_path = path + [action]
            if goal.issubset(nxt):
                return new_path
            seen.add(nxt)
            queue.append((nxt, new_path))
    return None


@dataclass
class Runtime:
    grants: FrozenSet[str] = frozenset({"search", "read", "draft", "verify"})
    max_steps: int = 20
    invalidate_after_first_draft: bool = False
    _draft_count: int = 0
    audit: List[str] = field(default_factory=list)

    def authorized(self, action: Action) -> bool:
        # action permission key defaults to action name
        need = action.permission or action.name
        return need in self.grants

    def execute(self, state: State, action: Action) -> State:
        if not self.authorized(action):
            raise PermissionError(action.name)
        if not applicable(state, action):
            raise ValueError(f"precondition failed at execute: {action.name}")
        nxt = apply(state, action)
        self.audit.append(action.name)
        if action.name == "draft":
            self._draft_count += 1
            if self.invalidate_after_first_draft and self._draft_count == 1:
                # 模拟证据失效：去掉 evidence 与 draft，保留 sources
                nxt = frozenset(s for s in nxt if s not in {"evidence", "draft"})
                self.audit.append("OBSERVATION:evidence_invalidated")
        return nxt


def run_loop(
    start: State,
    goal: State,
    runtime: Runtime,
    *,
    max_expansions: int = 200,
) -> Tuple[State, List[str]]:
    state = start
    steps = 0
    while not goal.issubset(state):
        if steps >= runtime.max_steps:
            runtime.audit.append("BUDGET_EXHAUSTED")
            break
        plan = bfs_plan(state, goal, max_expansions=max_expansions)
        if plan is None:
            runtime.audit.append("NO_PLAN")
            break
        if not plan:
            break
        action = plan[0]
        state = runtime.execute(state, action)
        steps += 1
    return state, runtime.audit


if __name__ == "__main__":
    # 1) 正常规划顺序：search -> read -> draft -> verify
    plan = bfs_plan(frozenset(), frozenset({"verified"}))
    assert plan is not None
    assert [a.name for a in plan] == ["search", "read", "draft", "verify"]

    # 2) 未审批时发布不可达（目标含 published，但状态永远没有 approved）
    no_pub = bfs_plan(frozenset(), frozenset({"published"}))
    assert no_pub is None

    # 3) 搜索扩展预算限制：极小 expansions 无法找到计划
    limited = bfs_plan(frozenset(), frozenset({"verified"}), max_expansions=1)
    assert limited is None

    # 4) 证据失效后的重新规划
    rt = Runtime(invalidate_after_first_draft=True, max_steps=20)
    final, trace = run_loop(frozenset(), frozenset({"verified"}), rt)
    assert "verified" in final
    assert trace == [
        "search",
        "read",
        "draft",
        "OBSERVATION:evidence_invalidated",
        "read",
        "draft",
        "verify",
    ]

    # 5) 未授权动作被拒
    denied_rt = Runtime(grants=frozenset({"search", "read", "draft"}))  # no verify
    try:
        denied_rt.execute(frozenset({"draft", "evidence"}), ACTIONS[3])  # verify
        raise AssertionError("expected PermissionError")
    except PermissionError as err:
        assert str(err) == "verify"

    # 6) 审批条件满足后的发布
    with_approval = bfs_plan(
        frozenset({"verified", "approved"}),
        frozenset({"published"}),
    )
    assert with_approval is not None
    assert [a.name for a in with_approval] == ["publish"]
    pub_rt = Runtime(grants=frozenset({"publish"}))
    published = pub_rt.execute(frozenset({"verified", "approved"}), ACTIONS[4])
    assert "published" in published

    # 7) 执行步数预算停止
    budget_rt = Runtime(max_steps=2, invalidate_after_first_draft=False)
    st, budget_trace = run_loop(frozenset(), frozenset({"verified"}), budget_rt)
    assert "BUDGET_EXHAUSTED" in budget_trace
    assert "verified" not in st
    assert budget_trace[:2] == ["search", "read"]

    print("A9 planning lab: 7 checks passed")
    print("Replanning trace: " + " -> ".join(trace))