models.py

from __future__ import annotations

import re
from enum import Enum
from typing import Any, Dict, List, Optional

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



class ActionType(str, Enum):
    EXECUTE = "execute"
    DEFER = "defer"
    DELEGATE = "delegate"
    OPTIMIZE = "optimize"


class TaskStatus(str, Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    DEFERRED = "deferred"
    FAILED = "failed"


class DifficultyLevel(str, Enum):
    EASY = "easy"
    MEDIUM = "medium"
    HARD = "hard"


def _coerce_action_type_from_text(text: str) -> ActionType:
    """
    Convert loose text from the web UI into the closest valid action type.

    The default web interface validates form fields against the Action model
    before our environment can run `message_to_action()`, so we accept
    friendly free-form input here and normalize it.
    """
    normalized = (text or "").strip().lower()
    if not normalized:
        return ActionType.EXECUTE

    keyword_map = (
        (ActionType.OPTIMIZE, ("optimize", "optimise", "tune", "analyze", "analyse")),
        (ActionType.DELEGATE, ("delegate", "assign", "handoff", "hand off", "offload")),
        (ActionType.DEFER, ("defer", "later", "wait", "skip", "postpone")),
        (ActionType.EXECUTE, ("execute", "run", "do", "complete", "process", "start")),
    )
    for action_type, keywords in keyword_map:
        if any(keyword in normalized for keyword in keywords):
            return action_type

    return ActionType.EXECUTE


def _extract_task_id_from_text(text: str) -> Optional[int]:
    if not text:
        return None

    explicit_match = re.search(r"(?:task|id|#)\s*(\d+)", text, flags=re.IGNORECASE)
    if explicit_match:
        return int(explicit_match.group(1))

    loose_match = re.search(r"\b(\d+)\b", text)
    if loose_match:
        return int(loose_match.group(1))

    return None


class TaskInfo(object):
    """Lightweight task descriptor (not a BaseModel to avoid nesting issues)."""

    def __init__(
        self,
        task_id: int,
        name: str,
        priority: float,
        deadline: int,
        uncertainty: float,
        value: float,
        required_energy: float,
        required_budget: float,
        category: str,
        status: str = "pending",
    ):
        self.task_id = task_id
        self.name = name
        self.priority = priority
        self.deadline = deadline
        self.uncertainty = uncertainty
        self.value = value
        self.required_energy = required_energy
        self.required_budget = required_budget
        self.category = category
        self.status = status

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "name": self.name,
            "priority": round(self.priority, 3),
            "deadline": self.deadline,
            "uncertainty": round(self.uncertainty, 3),
            "value": round(self.value, 3),
            "required_energy": round(self.required_energy, 3),
            "required_budget": round(self.required_budget, 3),
            "category": self.category,
            "status": self.status,
        }


class AetherTaskFlowAction(Action):
    """
    Action for the AETHER-TaskFlow environment.

    The agent selects a task by ID and decides how to act on it.
    """

    @model_validator(mode="before")
    @classmethod
    def normalize_web_input(cls, data: Any) -> Any:
        """
        Make the default web form resilient to casual text input.

        Examples that should validate cleanly:
          {"action_type": "hi"}
          {"action_type": "execute task 3"}
          {"message": "delegate 2"}
        """
        if isinstance(data, str):
            data = {"action_type": data}

        if not isinstance(data, dict):
            return data

        payload = dict(data)
        raw_message = payload.get("message")
        raw_action_type = payload.get("action_type")

        # When the UI sends a free-form message, treat it as action text.
        if isinstance(raw_message, str) and not raw_action_type:
            raw_action_type = raw_message
        payload.pop("message", None)

        if isinstance(raw_action_type, str):
            parsed_task_id = _extract_task_id_from_text(raw_action_type)
            payload["action_type"] = _coerce_action_type_from_text(raw_action_type).value

            if payload.get("task_id") in (None, "") and parsed_task_id is not None:
                payload["task_id"] = parsed_task_id

            if payload.get("reasoning") in (None, "") and raw_action_type.strip():
                payload["reasoning"] = f"parsed from '{raw_action_type.strip()[:80]}'"

        if payload.get("task_id") in (None, ""):
            payload["task_id"] = 0

        return payload

    action_type: ActionType = Field(
        ...,
        description=(
            "How to act on the selected task. "
            "'execute': consume resources and complete the task; "
            "'defer': postpone to a later step (low penalty); "
            "'delegate': offload at reduced reward but no resource cost; "
            "'optimize': reduce task uncertainty before execution."
        ),
    )
    task_id: int = Field(
        ...,
        ge=0,
        description="ID of the task to act on (from the current task list).",
    )
    reasoning: Optional[str] = Field(
        default=None,
        max_length=500,
        description="Optional agent reasoning for this action (logged but not scored).",
    )


class AetherTaskFlowObservation(Observation):
    """
    Observation returned after each step in the AETHER-TaskFlow environment.

    Contains all information the agent needs to make the next decision.
    """

    # Task queue
    tasks: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Current list of pending/deferred tasks as dicts.",
    )

    # Resource pool
    time_remaining: int = Field(
        default=10,
        ge=0,
        description="Time steps remaining in this episode.",
    )
    energy_remaining: float = Field(
        default=10.0,
        ge=0.0,
        description="Energy units remaining.",
    )
    budget_remaining: float = Field(
        default=50.0,
        ge=0.0,
        description="Budget units remaining.",
    )

    # System health
    system_health: float = Field(
        default=1.0,
        ge=0.0,
        le=1.0,
        description="Overall system health [0,1]. Drops on overload or missed deadlines.",
    )

    # Episode progress
    step_number: int = Field(
        default=0,
        ge=0,
        description="Current step number in this episode.",
    )
    tasks_completed: int = Field(
        default=0,
        ge=0,
        description="Total tasks completed so far.",
    )
    tasks_failed: int = Field(
        default=0,
        ge=0,
        description="Total tasks that missed their deadline.",
    )
    cumulative_value: float = Field(
        default=0.0,
        description="Total value accumulated so far.",
    )

    # Last action feedback
    last_action_type: Optional[str] = Field(
        default=None,
        description="Action type taken in the previous step.",
    )
    last_action_task_id: Optional[int] = Field(
        default=None,
        description="Task ID acted on in the previous step.",
    )
    last_action_outcome: Optional[str] = Field(
        default=None,
        description="Human-readable outcome of the last action.",
    )

    # Episode info
    difficulty: str = Field(
        default="easy",
        description="Current task difficulty level.",
    )
    episode_id: Optional[str] = Field(
        default=None,
        description="Unique identifier for this episode.",
    )



class AetherTaskFlowState(State):
    """
    Internal state of the AETHER-TaskFlow environment.

    This is the ground truth state used by the grader.
    """

    difficulty: str = Field(default="easy")
    tasks: List[Dict[str, Any]] = Field(default_factory=list)
    completed_tasks: List[Dict[str, Any]] = Field(default_factory=list)
    failed_tasks: List[Dict[str, Any]] = Field(default_factory=list)
    deferred_tasks: List[Dict[str, Any]] = Field(default_factory=list)

    resources: Dict[str, float] = Field(
        default_factory=lambda: {"time": 10.0, "energy": 10.0, "budget": 50.0}
    )
    initial_resources: Dict[str, float] = Field(
        default_factory=lambda: {"time": 10.0, "energy": 10.0, "budget": 50.0}
    )

    system_health: float = Field(default=1.0)
    cumulative_value: float = Field(default=0.0)
    cumulative_reward: float = Field(default=0.0)

    tasks_completed: int = Field(default=0)
    tasks_failed: int = Field(default=0)

    episode_done: bool = Field(default=False)
    seed: Optional[int] = Field(default=None)