# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

"""
Data models for the PostMortem incident-triage environment.

An agent plays an on-call SRE responding to a live incident. It queries fake
telemetry (logs / metrics / traces), scopes blast radius, hypothesises a root
cause, applies a mitigation, and writes a status-page update. The reward is a
5-stage process-reward ladder in [0, 1].
"""

from typing import Any, Dict, List

from openenv.core.env_server.types import Action, Observation
from pydantic import Field


class PostmortemAction(Action):
    """
    Single action envelope. `tool` selects an investigation / response verb,
    `args` is a tool-specific dict.

    Supported tools:
      - "ack"              args: {}
      - "query_logs"       args: {"service": str}
      - "query_metrics"    args: {"service": str}
      - "query_traces"     args: {"trace_id": str}
      - "scope"            args: {"services": list[str]}
      - "hypothesize"      args: {"root_cause": str}
      - "mitigate"         args: {"action": str}
      - "write_status"     args: {"text": str}
    """

    tool: str = Field(..., description="Investigation/response verb")
    args: Dict[str, Any] = Field(default_factory=dict, description="Tool args")


class PostmortemObservation(Observation):
    """Observation returned after each step."""

    task_id: str = Field(default="", description="Current scenario id")
    task_description: str = Field(default="", description="Incident brief")
    available_services: List[str] = Field(default_factory=list)
    available_trace_ids: List[str] = Field(default_factory=list)
    tool_result: str = Field(default="", description="Result from the last tool call")
    subgoals: Dict[str, bool] = Field(
        default_factory=lambda: {
            "acked": False,
            "scoped": False,
            "hypothesized": False,
            "mitigated": False,
            "written": False,
        }
    )
    reward_so_far: float = Field(default=0.0)
    steps_remaining: int = Field(default=0)
    last_error: str = Field(default="")