server/ba_agent_environment.py

"""
BA Agent RL Environment — full BA Toolkit pipeline (5 stages) over one feature.

Pipeline order:
    EXTRACT -> INTERVIEW -> GRAPH -> STORY_GEN -> PRD -> FINISH

Reward shape (dense):
    every step  -> shaped reward in [-0.10, +0.17]
    on FINISH   -> composite terminal reward in [-0.20, +1.05]
"""
from __future__ import annotations

import json
import random
from typing import Any, Dict, List, Optional
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment

from ba_agent_env.models import (
    ActionType,
    BAAgentAction,
    BAAgentObservation,
    BAAgentState,
    GeneratedStory,
    PipelineStage,
    TaskSample,
    TaskStatus,
)

from .dataset import get_task_by_id, sample_task
from .reward import step_reward, terminal_reward


_STAGE_FOR_ACTION = {
    ActionType.EXTRACT: PipelineStage.AWAITING_EXTRACT,
    ActionType.INTERVIEW: PipelineStage.AWAITING_INTERVIEW,
    ActionType.GRAPH: PipelineStage.AWAITING_GRAPH,
    ActionType.STORY_GEN: PipelineStage.AWAITING_STORY_GEN,
}

_NEXT_STAGE = {
    PipelineStage.AWAITING_EXTRACT: PipelineStage.AWAITING_INTERVIEW,
    PipelineStage.AWAITING_INTERVIEW: PipelineStage.AWAITING_GRAPH,
    PipelineStage.AWAITING_GRAPH: PipelineStage.AWAITING_STORY_GEN,
    PipelineStage.AWAITING_STORY_GEN: PipelineStage.READY_TO_FINISH,
}

_STAGE_TO_KEY = {
    ActionType.EXTRACT: "extract",
    ActionType.INTERVIEW: "interview",
    ActionType.GRAPH: "graph",
    ActionType.STORY_GEN: "story_gen",
}


def _parse_stories(payload: str) -> List[GeneratedStory]:
    text = (payload or "").strip()
    if not text:
        return []
    try:
        data = json.loads(text)
    except Exception:
        return []
    if isinstance(data, dict):
        data = data.get("stories") or data.get("generated_stories") or data.get("user_stories") or []
    if not isinstance(data, list):
        return []
    out: List[GeneratedStory] = []
    for i, s in enumerate(data):
        if not isinstance(s, dict):
            continue
        out.append(
            GeneratedStory(
                story_id=str(s.get("story_id") or s.get("id") or f"S-{i+1:03d}"),
                title=str(s.get("title") or s.get("name") or ""),
                description=str(s.get("description") or s.get("narrative") or ""),
                acceptance_criteria=s.get("acceptance_criteria") or s.get("ac") or "",
                story_points=s.get("story_points"),
                state=s.get("state"),
            )
        )
    return out


def _allowed_actions(stage: PipelineStage) -> List[ActionType]:
    if stage == PipelineStage.AWAITING_EXTRACT:
        return [ActionType.EXTRACT]
    if stage == PipelineStage.AWAITING_INTERVIEW:
        return [ActionType.INTERVIEW]
    if stage == PipelineStage.AWAITING_GRAPH:
        return [ActionType.GRAPH]
    if stage == PipelineStage.AWAITING_STORY_GEN:
        return [ActionType.STORY_GEN]
    if stage == PipelineStage.READY_TO_FINISH:
        return [ActionType.FINISH]
    return []


class BAAgentEnvironment(Environment[BAAgentAction, BAAgentObservation, BAAgentState]):
    """OpenEnv environment wrapping the full BA Toolkit pipeline over one feature."""

    def __init__(
        self,
        task_id: Optional[str] = None,
        max_steps: int = 12,
        seed: Optional[int] = None,
    ):
        self._task_id_override = task_id
        self._max_steps = max_steps
        self._rng = random.Random(seed)
        self._state: BAAgentState = BAAgentState()
        self.reset()

    def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs: Any) -> BAAgentObservation:  # type: ignore[override]
        if seed is not None:
            self._rng = random.Random(seed)
        if self._task_id_override:
            task = get_task_by_id(self._task_id_override) or sample_task(self._rng)
        else:
            task = sample_task(self._rng)
        self._state = BAAgentState(
            task=task,
            current_stage=PipelineStage.AWAITING_EXTRACT,
            pipeline_outputs={},
            generated_stories=[],
            step_count=0,
            max_steps=self._max_steps,
            status=TaskStatus.RUNNING,
            last_reward=0.0,
            episode_id=episode_id or str(uuid4()),
        )
        return self._build_obs(feedback="Episode reset. Run EXTRACT to begin.")

    def step(self, action: BAAgentAction, timeout_s: Optional[float] = None, **kwargs: Any) -> BAAgentObservation:  # type: ignore[override]
        state = self._state
        state.step_count += 1
        task = state.task
        assert task is not None

        if state.step_count > state.max_steps:
            state.status = TaskStatus.TASK_LIMIT_REACHED
            return self._build_obs(feedback="Step limit exceeded.", reward=-0.10, done=True, metadata={"status": state.status.value, "step_breakdown": {"step_limit": -0.10}})

        if action.action_type == ActionType.FINISH:
            if state.current_stage != PipelineStage.READY_TO_FINISH:
                return self._build_obs(
                    feedback=f"FINISH called too early. Current stage: {state.current_stage.value}.",
                    reward=-0.10, done=False,
                    metadata={"status": "early_finish", "step_breakdown": {"early_finish": -0.10}},
                )
            story_gen_present = bool(state.generated_stories)
            reward, details = terminal_reward(state.generated_stories, task, story_gen_present=story_gen_present)
            state.last_reward = reward
            state.status = TaskStatus.COMPLETED
            state.current_stage = PipelineStage.DONE
            return self._build_obs(
                feedback=f"Episode complete. Composite {details['composite_0_to_100']:.1f}/100. Terminal {details['terminal_total']:.3f}.",
                reward=reward, done=True, metadata=details,
            )

        out_of_order = False
        expected_stage = _STAGE_FOR_ACTION.get(action.action_type)
        if expected_stage is None or state.current_stage != expected_stage:
            out_of_order = True

        prior_outputs = dict(state.pipeline_outputs)
        parsed = _parse_stories(action.payload) if action.action_type == ActionType.STORY_GEN else None

        r, bd = step_reward(
            action_type=action.action_type.value,
            payload=action.payload or "",
            task=task,
            prior_outputs=prior_outputs,
            out_of_order=out_of_order,
            parsed_stories=parsed,
        )

        if out_of_order:
            state.status = TaskStatus.INVALID_ACTION_ORDER
            return self._build_obs(
                feedback=(
                    f"Invalid action {action.action_type.value} for stage {state.current_stage.value}. "
                    f"Allowed: {[a.value for a in _allowed_actions(state.current_stage)]}"
                ),
                reward=r, done=False, metadata={"status": state.status.value, "step_breakdown": bd},
            )

        if action.action_type == ActionType.STORY_GEN and not parsed:
            state.status = TaskStatus.INVALID_PAYLOAD
            return self._build_obs(
                feedback="STORY_GEN payload must be a JSON list of stories.",
                reward=r, done=False, metadata={"status": state.status.value, "step_breakdown": bd},
            )

        key = _STAGE_TO_KEY[action.action_type]
        state.pipeline_outputs[key] = action.payload or ""
        if action.action_type == ActionType.STORY_GEN:
            state.generated_stories = parsed or []
        state.status = TaskStatus.RUNNING
        state.current_stage = _NEXT_STAGE[state.current_stage]
        return self._build_obs(
            feedback=f"Stage {action.action_type.value} accepted. Next: {state.current_stage.value}.",
            reward=r, done=False, metadata={"status": "ok", "step_breakdown": bd},
        )

    @property
    def state(self) -> BAAgentState:  # type: ignore[override]
        return self._state

    def _build_obs(
        self, feedback: str = "", reward: float = 0.0, done: bool = False,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> BAAgentObservation:
        s = self._state
        task = s.task
        assert task is not None
        return BAAgentObservation(
            task_id=task.task_id,
            title=task.title,
            description=task.description,
            input_documents=task.input_documents,
            current_stage=s.current_stage,
            available_actions=_allowed_actions(s.current_stage),
            pipeline_outputs=dict(s.pipeline_outputs),
            step_count=s.step_count,
            max_steps=s.max_steps,
            feedback=feedback,
            status=s.status,
            reward=reward,
            done=done,
            metadata=metadata or {},
        )