# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# BSD-style license

"""
Data models for the OrgSim Environment.

Defines typed Action, Observation, and State for the organization simulation.
"""

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


class OrgAction(Action):
    """Action for organization simulation.

    Contains agent action to interact with the organization structure.
    """

    action_type: Literal[
        "REQUEST_TASK",
        "ACCEPT_TASK",
        "COMPLETE_TASK",
        "REQUEST_HELP",
        "PROVIDE_HELP",
        "ESCALATE",
        "REQUEST_RESOURCE",
        "REPORT_STATUS",
    ] = Field(..., description="Type of action to perform")

    target_id: str = Field(default="", description="Target agent/team ID")

    payload: dict[str, Any] = Field(
        default_factory=dict, description="Additional action parameters"
    )


class OrgState(State):
    """State of the OrgSimEnvironment containing organization context."""

    episode_id: str = Field(default="", description="Unique episode identifier")

    step_count: int = Field(default=0, description="Current step number")

    tasks: dict[str, dict[str, Any]] = Field(
        default_factory=dict, description="All tasks in the organization"
    )

    team_members: dict[str, list[str]] = Field(
        default_factory=dict, description="Team membership: team_id -> [agent_ids]"
    )

    resource_status: dict[str, bool] = Field(
        default_factory=dict, description="Resource availability status"
    )


class OrgObservation(Observation):
    """Observation returned by OrgSimEnvironment.

    Contains current organization state from agent's perspective.
    """

    my_agent_id: str = Field(..., description="Current agent ID")
    my_team: str = Field(..., description="Current agent's team")
    my_role: str = Field(default="member", description="Role: lead or member")

    available_tasks: list[dict[str, Any]] = Field(
        default_factory=list, description="Tasks available for assignment"
    )

    active_task: dict[str, Any] | None = Field(
        default=None, description="Currently assigned task"
    )

    inbox: list[dict[str, Any]] = Field(
        default_factory=list, description="Messages/messages for this agent"
    )

    team_status: dict[str, Any] = Field(
        default_factory=dict, description="Team workload and availability"
    )

    resources: dict[str, bool] = Field(
        default_factory=dict, description="Available resources"
    )

    metrics: dict[str, Any] = Field(default_factory=dict, description="Episode metrics")

    done: bool = Field(default=False, description="Whether the episode has ended")

    reward: float = Field(default=0.0, description="Reward for the last action")

    metadata: dict[str, Any] = Field(default_factory=dict, description="Extra metadata")