"""Shared proposal types and session-owned policy memory. No game or model calls."""
from __future__ import annotations

import json
from dataclasses import dataclass, field


@dataclass
class Decision:
    action: str
    params: dict = field(default_factory=dict)
    reason: str = ""
    source: str = "code"          # code | jev | fallback
    confidence: float | None = None

    def __str__(self) -> str:
        params = ", ".join(f"{k}={v!r}" for k, v in self.params.items())
        conf = f" conf={self.confidence:.2f}" if self.confidence is not None else ""
        return f"[{self.source}{conf}] {self.action}({params})  # {self.reason}"


# rewards and card_reward are two views of one flow: claiming a card reward
# opens the card screen, and skipping returns to the rewards screen. Treat them
# as ONE screen group so per-flow state is not cleared on every hop.
SCREEN_GROUPS = {"rewards": "rewards_flow", "card_reward": "rewards_flow"}


def _screen_group(state_type: str | None, menu_screen: str | None = None) -> str | None:
    """
    A stable key for "which screen am I on".

    The `menu` state_type covers the main menu, mode select, character select
    and the tutorial prompt. Those are different screens with different valid
    actions, so the menu_screen is part of the key.
    """
    if state_type == "menu":
        return f"menu:{menu_screen}"
    return SCREEN_GROUPS.get(state_type, state_type)


@dataclass
class PendingAction:
    """An accepted request, not proof that the game completed it."""

    decision: Decision
    screen: dict


@dataclass
class PolicyContext:
    """Session-owned memory. Proposals never record execution.

    Grid toggles and character selection lack observable selection fields in
    the mod. Their accepted requests are tracked explicitly, not called facts.
    Other guarded actions wait for relevant screen evidence before continuing.
    """

    screen_group: str | None = None
    accepted_card_indices: set[int] = field(default_factory=set)
    accepted_card_grid: list[dict] | None = None
    accepted_crystal_cells: set[tuple[int, int]] = field(default_factory=set)
    rewards_skipped_card: bool = False
    character_selected: bool = False
    pending: PendingAction | None = None

    def observe(self, obs: dict) -> None:
        st = obs.get("state_type")
        if st in ("unknown", "overlay"):
            return  # A transient overlay is not evidence of completion.
        group = _screen_group(st, obs.get("menu_screen"))
        if group != self.screen_group:
            self.screen_group = group
            self.accepted_card_indices.clear()
            self.accepted_card_grid = None
            self.accepted_crystal_cells.clear()
            self.rewards_skipped_card = False
            self.character_selected = False
            self.pending = None
            return
        pending = self.pending
        if pending is None:
            return
        action, params = pending.decision.action, pending.decision.params
        screen = obs.get(st) or {}
        if action == "select_card" and pending.screen.get("screen_type") != "choose":
            # The API exposes no selected indices. Reserve accepted toggles
            # after a fresh read, even when their effect is not visible yet.
            if screen.get("cards") != pending.screen.get("cards"):
                return  # Cannot map a toggle safely onto a changed grid.
            self.accepted_card_indices.add(params["index"])
            self.accepted_card_grid = pending.screen.get("cards")
        elif action == "menu_select" and params.get("option") != "embark":
            self.character_selected = True
        elif action == "skip_card_reward":
            if st != "rewards":
                return
            self.rewards_skipped_card = True
        elif action in ("confirm_selection", "cancel_selection",
                        "confirm_bundle_selection", "cancel_bundle_selection",
                        "combat_confirm_selection"):
            return  # Wait for the screen to close; do not cancel a slow confirm.
        elif action == "crystal_sphere_click_cell":
            cell = (params["x"], params["y"])
            clickable = {(c.get("x"), c.get("y"))
                         for c in screen.get("clickable_cells", [])}
            if cell in clickable and not screen.get("can_proceed"):
                return
            self.accepted_crystal_cells.add(cell)
        elif action == "shop_purchase":
            # Gold alone can change for unrelated reasons. Require this item
            # to change/disappear, or a screen transition (e.g. card removal).
            before = pending.screen.get("shop", pending.screen)
            after = screen.get("shop", screen)
            index = params["index"]
            old = next((i for i in before.get("items", []) if i.get("index") == index), None)
            new = next((i for i in after.get("items", []) if i.get("index") == index), None)
            def inventory_item(item):
                return {k: v for k, v in item.items() if k not in ("can_afford", "price")} if item else None
            if inventory_item(old) == inventory_item(new):
                return
        elif action == "combat_select_card":
            if len(screen.get("selected_cards") or []) <= len(pending.screen.get("selected_cards") or []):
                return
        elif action == "select_bundle":
            if not screen.get("preview_showing") or not screen.get("can_confirm"):
                return
        else:
            return  # Direct choices and embark require a screen transition.
        self.pending = None

    def record_result(self, obs: dict, decision: Decision, *, accepted: bool) -> None:
        """Called only after a real action result. Transport errors stop the runner."""
        if not accepted:
            self.pending = None
            if decision.action == "menu_select" and decision.params.get("option") == "embark":
                self.character_selected = False
            return
        st = obs.get("state_type")
        guarded = decision.action in {
            "select_card", "confirm_selection", "cancel_selection",
            "select_bundle", "confirm_bundle_selection", "cancel_bundle_selection",
            "shop_purchase", "skip_card_reward", "crystal_sphere_click_cell",
            "combat_select_card", "combat_confirm_selection",
        } or (st == "menu" and obs.get("menu_screen") == "character_select")
        if guarded:
            # Detach from mutable fixtures/callers. This is a small screen node,
            # not another copy of the entire combat observation.
            screen = json.loads(json.dumps(obs.get(st) or {}))
            self.pending = PendingAction(decision, screen)