from typing import Any, Dict, List, Literal, Optional
from pydantic import Field
from openenv.core.env_server.types import Action, Observation, State


class PMOpsAction(Action):
    action_type: Literal[
        "ticketing.create_ticket",
        "ticketing.update_ticket",
        "ticketing.get_ticket",
        "ticketing.list_tickets",
        "ticketing.assign_ticket",
        "ticketing.comment_ticket",
        "ticketing.transition_ticket",
        "codebase.list_commits",
        "codebase.get_commit",
        "codebase.list_prs",
        "chat.post_message",
        "chat.read_channel",
        "chat.list_channels",
        "chat.search",
        "meta.read_runbook",
        "meta.finish",
        "meta.noop",
    ] = Field(default="meta.noop")
    args: Dict[str, Any] = Field(default_factory=dict)


class PMOpsObservation(Observation):
    step: int = Field(default=0)
    max_steps: int = Field(default=40)
    task_brief: str = Field(default="")
    last_action_result: Dict[str, Any] = Field(default_factory=dict)
    app_state_deltas: Dict[str, List[Any]] = Field(
        default_factory=lambda: {"ticketing": [], "chat": [], "codebase": []}
    )
    steps_remaining: int = Field(default=40)
    token_budget_remaining: int = Field(default=12000)


class PMOpsState(State):
    org_config: Dict[str, Any] = Field(default_factory=dict)
    task_config: Dict[str, Any] = Field(default_factory=dict)
    ticketing: Dict[str, Any] = Field(default_factory=dict)
    chat: Dict[str, Any] = Field(default_factory=dict)
    codebase: Dict[str, Any] = Field(default_factory=dict)
    step_count: int = Field(default=0)
    finished: bool = Field(default=False)
    episode_id: str | None = Field(default="")