app/environment.py

"""AP Commander — Core AP Clerk Environment (OpenEnv interface: reset / step / state)."""

from __future__ import annotations
from typing import Optional, Tuple, Dict, Any, List

from .models import APObservation, APAction, APReward, DecisionType
from .tasks import TASKS, grade_action

# OpenEnv base class — use real package if available, else local stub
try:
    from openenv.core.env_server import Environment as _OpenEnvBase
except ImportError:
    class _OpenEnvBase:  # type: ignore[no-redef]
        """Stub matching the OpenEnv Environment interface."""
        def reset(self, seed=None, episode_id=None, **kwargs):
            raise NotImplementedError
        def step(self, action, timeout_s=None, **kwargs):
            raise NotImplementedError
        @property
        def state(self):
            raise NotImplementedError

# Intermediate actions: do not end the episode, reveal context instead
_INTERMEDIATE = frozenset({
    DecisionType.QUERY_VENDOR,
    DecisionType.ESCALATE,
    DecisionType.HOLD,
    DecisionType.HYPOTHETICAL,
})


class APClerkEnvironment(_OpenEnvBase):
    """
    AP Clerk environment. Intermediate actions (QUERY_VENDOR / ESCALATE / HOLD / HYPOTHETICAL)
    reveal context without ending the episode. Terminal actions grade and close it.
    Fixed seed → reproducible episode; None → fresh random each call.
    """

    def __init__(self) -> None:
        self._task_id:       Optional[str]          = None
        self._observation:   Optional[APObservation] = None
        self._step_count:    int                     = 0
        self._done:          bool                    = False
        self._episode_score: float                   = 0.01
        self._max_steps:     int                     = 1
        self._context_store: List[str]               = []

    def reset(self, task_id: str, seed: Optional[int] = None) -> APObservation:
        if task_id not in TASKS:
            raise ValueError(
                f"Unknown task_id {task_id!r}. "
                f"Valid options: {list(TASKS.keys())}"
            )
        spec = TASKS[task_id]
        self._task_id       = task_id
        self._step_count    = 0
        self._done          = False
        self._episode_score = 0.01
        self._max_steps     = spec.max_steps

        obs = spec.generator(seed=seed)
        obs.step_count  = 0
        obs.max_steps   = self._max_steps

        # Context notes pre-generated by the generator are hidden until revealed
        # by an intermediate action (ESCALATE / QUERY_VENDOR / HOLD).
        self._context_store = list(obs.context_notes)
        obs.context_notes   = []
        obs.action_history  = []

        self._observation = obs
        return obs

    def step(self, action: APAction) -> Tuple[APObservation, APReward, bool, Dict[str, Any]]:
        if self._observation is None:
            raise RuntimeError("Call reset(task_id) before step().")
        if self._done:
            raise RuntimeError("Episode already finished. Call reset() to start a new one.")

        self._step_count += 1
        is_intermediate = action.decision in _INTERMEDIATE
        at_limit        = self._step_count >= self._max_steps

        # ── HYPOTHETICAL action (training self-play, no context reveal) ─────────
        if action.decision == DecisionType.HYPOTHETICAL and not at_limit:
            self._observation.action_history.append({
                "step":        self._step_count,
                "decision":    "HYPOTHETICAL",
                "reason_code": action.reason_code.value,
                "explanation": action.explanation,
            })
            self._observation.step_count = self._step_count
            # Simulate a brief outcome hint without revealing graded context
            hint = (
                f"[HYPOTHETICAL] If you chose {action.reason_code.value}: "
                f"outcome would depend on whether the invoice data supports it. "
                "Review the invoice total, PO amounts, and GRN quantities carefully before committing."
            )
            self._observation.context_notes.append(hint)
            reward = APReward(
                score=0.01,
                breakdown={"hypothetical_step": action.reason_code.value},
                feedback="Hypothetical path explored. Now commit to a real decision.",
                done=False,
            )
            info: Dict[str, Any] = {
                "task_id":    self._task_id,
                "step_count": self._step_count,
                "episode_score": 0.01,
                "hypothetical": True,
            }
            return self._observation, reward, False, info

        # ── Intermediate action (episode continues) ───────────────────────────
        if is_intermediate and not at_limit:
            if self._context_store:
                prefix_map = {
                    DecisionType.ESCALATE:     "[MANAGER]",
                    DecisionType.QUERY_VENDOR: "[VENDOR]",
                    DecisionType.HOLD:         "[COMPLIANCE]",
                }
                preferred_prefix = prefix_map.get(action.decision, "")
                # Prefer context note matching the action type; fall back to FIFO
                matched_idx = next(
                    (i for i, n in enumerate(self._context_store)
                     if preferred_prefix and n.startswith(preferred_prefix)),
                    0,
                )
                note = self._context_store.pop(matched_idx)
                self._observation.context_notes.append(note)

            # History lets graders award the process bonus for correct sequences
            self._observation.action_history.append({
                "step":        self._step_count,
                "decision":    action.decision.value,
                "reason_code": action.reason_code.value,
                "explanation": action.explanation,
            })
            self._observation.step_count = self._step_count

            steps_left = self._max_steps - self._step_count
            context_revealed = self._observation.context_notes[-1] \
                if self._observation.context_notes else None

            reward = APReward(
                score=0.01,
                breakdown={
                    "intermediate_step": action.decision.value,
                    "steps_remaining":   steps_left,
                },
                feedback=(
                    f"Intermediate step recorded: {action.decision.value}. "
                    + (f"Context revealed — {context_revealed}"
                       if context_revealed else
                       "No additional context available. Proceed to final decision.")
                ),
                done=False,
            )
            info = {
                "task_id":    self._task_id,
                "step_count": self._step_count,
                "episode_score": 0.01,
                "intermediate": True,
            }
            return self._observation, reward, False, info

        # ── Terminal action (or forced terminal at step limit) ────────────────
        self._observation.action_history.append({
            "step":        self._step_count,
            "decision":    action.decision.value,
            "reason_code": action.reason_code.value,
            "explanation": action.explanation,
        })

        reward               = grade_action(self._task_id, self._observation, action)
        # Clamp to open interval (0, 1) as required by the evaluator
        clamped_score        = max(0.01, min(0.99, reward.score))
        reward               = APReward(
            score=clamped_score,
            breakdown=reward.breakdown,
            feedback=reward.feedback,
            done=reward.done,
        )
        self._done           = True
        self._episode_score  = reward.score
        self._observation.step_count = self._step_count

        info = {
            "task_id":       self._task_id,
            "step_count":    self._step_count,
            "episode_score": self._episode_score,
        }
        return self._observation, reward, self._done, info

    @property
    def state(self) -> Dict[str, Any]:
        return {
            "task_id":             self._task_id,
            "step_count":          self._step_count,
            "done":                self._done,
            "episode_score":       self._episode_score,
            "current_observation": self._observation,
        }

    @staticmethod
    def list_tasks() -> Dict[str, Dict[str, str]]:
        return {
            tid: {
                "name":       spec.name,
                "difficulty": spec.difficulty,
                "description": spec.description,
            }
            for tid, spec in TASKS.items()
        }