env/aether_env.py

from __future__ import annotations

import random
import re
import sys
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional

# Allow running from repo root or server/
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from openenv.core import Environment

from models import (
    AetherTaskFlowAction,
    AetherTaskFlowObservation,
    AetherTaskFlowState,
    ActionType,
)
from env.tasks import generate_tasks, get_profile, apply_dynamic_updates
from env.grader import grade



class AetherTaskFlowEnvironment(Environment):
    """
    AETHER-TaskFlow: Adaptive Workflow Management RL Environment.

    The agent manages a dynamic task queue under resource constraints
    and system uncertainty. Three task scenarios of increasing difficulty
    test baseline reasoning, adaptation, and robustness.
    """

    SUPPORTS_CONCURRENT_SESSIONS = True
    DEFAULT_SEED = 42

    def __init__(self, difficulty: str = "easy", default_seed: int = DEFAULT_SEED) -> None:
        super().__init__()
        if difficulty not in ("easy", "medium", "hard"):
            raise ValueError(f"difficulty must be easy/medium/hard, got '{difficulty}'")
        self._difficulty = difficulty
        self._profile = get_profile(difficulty)
        self._default_seed = int(default_seed)
        self._state: AetherTaskFlowState = AetherTaskFlowState()
        self._rng = random.Random()
        self._deferred_tasks: List[Dict[str, Any]] = []



    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> AetherTaskFlowObservation:
        self._reset_rubric()
        seed = self._default_seed if seed is None else int(seed)
        
        # === CRITICAL: Ensure determinism ===
        random.seed(seed)
        
        self._rng = random.Random(seed)
        ep_id = episode_id or str(uuid.uuid4())

        profile = self._profile
        tasks = generate_tasks(self._difficulty, seed=seed)
        task_dicts = [t.to_dict() for t in tasks]

        resources = {
            "time": float(profile["initial_time"]),
            "energy": float(profile["initial_energy"]),
            "budget": float(profile["initial_budget"]),
        }

        self._state = AetherTaskFlowState(
            episode_id=ep_id,
            step_count=0,
            difficulty=self._difficulty,
            tasks=task_dicts,
            completed_tasks=[],
            failed_tasks=[],
            deferred_tasks=[],
            resources=dict(resources),
            initial_resources=dict(resources),
            system_health=1.0,
            cumulative_value=0.0,
            cumulative_reward=0.0,
            tasks_completed=0,
            tasks_failed=0,
            episode_done=False,
            seed=seed,
        )
        self._deferred_tasks = []
        self._sync_state_queues()

        return self._build_obs(
            last_action_type=None,
            last_action_task_id=None,
            last_action_outcome="Episode started. Select a task to act on.",
            reward=0.0,
            done=False,
        )


    def step(
        self,
        action: AetherTaskFlowAction | Dict[str, Any] | str,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> AetherTaskFlowObservation:
        if self._state.episode_id is None:
            self.reset(seed=self._default_seed)

        try:
            parsed_action = self._coerce_action(action)
            return self._step_impl(parsed_action)
        except Exception as exc:
            return self._safe_step_failure(action, exc)

    def _step_impl(self, action: AetherTaskFlowAction) -> AetherTaskFlowObservation:
        s = self._state

        if s.episode_done:
            return self._build_obs(
                last_action_type=None,
                last_action_task_id=None,
                last_action_outcome="Episode already finished.",
                reward=0.0,
                done=True,
            )

        s.step_count += 1
        profile = self._profile
        max_steps: int = profile["max_steps"]

        # ---- Apply dynamic task updates (medium/hard) ----
        if self._difficulty in ("medium", "hard"):
            from env.tasks import apply_dynamic_updates as _upd
            task_objs_updated = _upd(
                [self._make_task_info(t) for t in s.tasks],
                s.step_count,
                self._difficulty,
                self._rng,
            )
            s.tasks = [t.to_dict() for t in task_objs_updated]

        # ---- Deadline expiry check (before acting) ----
        still_alive, newly_failed = [], []
        for t in s.tasks:
            if t.get("deadline", 1) <= 0 and t["status"] == "pending":
                t["status"] = "failed"
                newly_failed.append(t)
                s.system_health = max(0.0, s.system_health - 0.05)
            else:
                still_alive.append(t)
        s.tasks = still_alive
        s.failed_tasks.extend(newly_failed)
        s.tasks_failed += len(newly_failed)

        # ---- Find the target task ----
        task = self._find_task(action.task_id, s.tasks)
        if task is None:
            # Try deferred list
            task = self._find_task(action.task_id, self._deferred_tasks)

        raw_reward = 0.0
        outcome = ""

        if task is None:
            raw_reward = -0.5
            outcome = (
                f"Task {action.task_id} not found in active queue. "
                "Choose a valid task_id from the observation."
            )
            s.system_health = max(0.0, s.system_health - 0.02)
        else:
            raw_reward, outcome = self._execute_action(action.action_type, task, s, max_steps)

        s.cumulative_reward += raw_reward

        # ---- Decrement deadlines each step ----
        for t in s.tasks:
            if t["status"] == "pending":
                t["deadline"] = max(0, t["deadline"] - 1)

        # ---- Recycle deferred tasks if resources improve ----
        from env.algorithms import AWFROX
        recycler = AWFROX()
        resources_dict = {
            "energy": s.resources["energy"],
            "budget": s.resources["budget"],
        }
        active_updated, still_deferred = recycler.recycle_deferred(
            s.tasks, self._deferred_tasks, resources_dict, s.step_count
        )
        s.tasks = active_updated
        self._deferred_tasks = still_deferred
        self._sync_state_queues()

        # ---- Done condition ----
        no_more_tasks = len(s.tasks) == 0 and len(self._deferred_tasks) == 0
        out_of_time = s.step_count >= max_steps
        out_of_resources = (
            s.resources["energy"] <= 0 or s.resources["time"] <= 0
        )
        system_collapse = s.system_health <= 0.0

        done = no_more_tasks or out_of_time or out_of_resources or system_collapse
        s.episode_done = done
        self._sync_state_queues()

        return self._build_obs(
            last_action_type=action.action_type.value,
            last_action_task_id=action.task_id,
            last_action_outcome=outcome,
            reward=self._normalize_step_reward(raw_reward),
            done=done,
        )

    def message_to_action(self, message: str) -> AetherTaskFlowAction:
        """Convert free-form UI text into a valid environment action."""
        return self._coerce_action(message)

    def _coerce_action(
        self,
        action: AetherTaskFlowAction | Dict[str, Any] | str | None,
    ) -> AetherTaskFlowAction:
        if isinstance(action, AetherTaskFlowAction):
            return action

        if action is None:
            return self._recommended_action("No action provided; selected a safe default.")

        if isinstance(action, str):
            return self._parse_action_message(action)

        if isinstance(action, dict):
            if "message" in action and isinstance(action["message"], str):
                return self._parse_action_message(action["message"])
            if "input" in action and isinstance(action["input"], str):
                return self._parse_action_message(action["input"])
            if "action" in action:
                nested_action = action["action"]
                if isinstance(nested_action, (dict, str)) or nested_action is None:
                    return self._coerce_action(nested_action)

            recommended = self._recommended_action("Filled missing action fields from the current state.")
            normalized_payload = {
                "action_type": action.get("action_type", recommended.action_type.value),
                "task_id": action.get("task_id", recommended.task_id),
                "reasoning": action.get("reasoning", recommended.reasoning),
            }
            return AetherTaskFlowAction.model_validate(normalized_payload)

        raise TypeError(f"Unsupported action input: {type(action)!r}")

    def _parse_action_message(self, message: str) -> AetherTaskFlowAction:
        normalized = (message or "").strip().lower()
        recommended = self._recommended_action(
            "Selected the top-ranked task from the current observation."
        )

        if not normalized:
            return recommended

        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")),
        )

        chosen_action = recommended.action_type
        for action_type, keywords in keyword_map:
            if any(keyword in normalized for keyword in keywords):
                chosen_action = action_type
                break

        requested_task_id = self._extract_task_id(normalized)
        if requested_task_id is not None and self._task_exists(requested_task_id):
            task_id = requested_task_id
        else:
            task_id = recommended.task_id

        return AetherTaskFlowAction(
            action_type=chosen_action,
            task_id=task_id,
            reasoning=f"parsed from '{message.strip()[:80]}'",
        )

    def _recommended_action(self, reasoning: str) -> AetherTaskFlowAction:
        candidates = self._iter_candidate_tasks()
        if not candidates:
            return AetherTaskFlowAction(
                action_type=ActionType.DEFER,
                task_id=0,
                reasoning=reasoning,
            )

        from env.algorithms import AETHER, RAPTOR

        resources = {
            "energy": self._state.resources.get("energy", 0.0),
            "budget": self._state.resources.get("budget", 0.0),
            "time": self._state.resources.get("time", 0.0),
        }
        max_steps = self._profile["max_steps"]
        ranked = AETHER().rank_tasks(candidates, resources, self._state.step_count, max_steps)
        best_task_id, _ = ranked[0]
        best_task = next(task for task in candidates if task["task_id"] == best_task_id)
        action_type = ActionType(
            RAPTOR().decide(best_task, resources, self._state.step_count, max_steps)
        )
        return AetherTaskFlowAction(
            action_type=action_type,
            task_id=best_task_id,
            reasoning=reasoning,
        )

    def _iter_candidate_tasks(self) -> List[Dict[str, Any]]:
        active_tasks = [task for task in self._state.tasks if task.get("status") == "pending"]
        if active_tasks:
            return active_tasks
        deferred_tasks = [
            task for task in self._deferred_tasks if task.get("status") in ("pending", "deferred")
        ]
        return deferred_tasks

    def _task_exists(self, task_id: int) -> bool:
        return self._find_task(task_id, self._state.tasks) is not None or self._find_task(
            task_id, self._deferred_tasks
        ) is not None

    def _extract_task_id(self, text: str) -> Optional[int]:
        explicit_match = re.search(r"(?:task|id|#)\s*(\d+)", text)
        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

    def _safe_step_failure(
        self,
        action: AetherTaskFlowAction | Dict[str, Any] | str,
        exc: Exception,
    ) -> AetherTaskFlowObservation:
        self._state.episode_done = True
        self._state.system_health = max(0.0, self._state.system_health - 0.1)
        self._sync_state_queues()

        last_action_type = None
        last_action_task_id = None
        if isinstance(action, AetherTaskFlowAction):
            last_action_type = action.action_type.value
            last_action_task_id = action.task_id
        elif isinstance(action, dict):
            raw_action_type = action.get("action_type")
            if isinstance(raw_action_type, str):
                last_action_type = raw_action_type
            raw_task_id = action.get("task_id")
            if isinstance(raw_task_id, int):
                last_action_task_id = raw_task_id

        return self._build_obs(
            last_action_type=last_action_type,
            last_action_task_id=last_action_task_id,
            last_action_outcome=(
                f"Step failed safely: {type(exc).__name__}: {str(exc)[:160]}"
            ),
            reward=self._normalize_step_reward(-1.0),
            done=True,
        )


    def _reset_rubric(self) -> None:
        """Called at the start of every reset() — OpenEnv lifecycle hook."""
        # No persistent rubric state in this env; this hook satisfies the
        # openenv.core.Environment base-class interface.
        pass

    def get_metadata(self) -> dict:
        """Return environment metadata (used by WebInterfaceManager on startup)."""
        return {
            "name": "aether_taskflow",
            "description": (
                "AETHER-TaskFlow: Adaptive Workflow Management RL Environment. "
                "Agent manages a dynamic task queue under resource constraints, "
                "uncertainty, and time pressure. Real-world enterprise tasks."
            ),
            "difficulty": self._difficulty,
            "max_steps": self._profile["max_steps"],
            "action_types": ["execute", "defer", "delegate", "optimize"],
            "version": "1.0.0",
        }

    def close(self) -> None:
        """Clean up environment resources (no-op for this in-memory env)."""
        pass

    # ------------------------------------------------------------------
    # state property (OpenEnv required)
    # ------------------------------------------------------------------

    @property
    def state(self) -> AetherTaskFlowState:
        self._sync_state_queues()
        return self._state


    def _execute_action(
        self,
        action_type: ActionType,
        task: Dict[str, Any],
        s: AetherTaskFlowState,
        max_steps: int,
    ) -> tuple[float, str]:
        """Execute the chosen action on a task. Returns (reward, outcome_str)."""

        energy_cost = task["required_energy"]
        budget_cost = task["required_budget"]
        uncertainty = task["uncertainty"]
        value = task["value"]
        priority = task["priority"]
        deadline = task["deadline"]
        time_left = max_steps - s.step_count

        if action_type == ActionType.EXECUTE:
            # Check resource sufficiency
            if s.resources["energy"] < energy_cost or s.resources["budget"] < budget_cost:
                s.system_health = max(0.0, s.system_health - 0.08)
                return -1.0, (
                    f"Cannot execute '{task['name']}': insufficient resources "
                    f"(need E={energy_cost:.1f}/B={budget_cost:.1f}, "
                    f"have E={s.resources['energy']:.1f}/B={s.resources['budget']:.1f})."
                )

            # Uncertainty-based failure chance
            success_prob = 1.0 - uncertainty * 0.4
            if self._rng.random() > success_prob:
                # Partial failure — lose resources, get partial reward
                s.resources["energy"] = max(0.0, s.resources["energy"] - energy_cost * 0.5)
                s.resources["budget"] = max(0.0, s.resources["budget"] - budget_cost * 0.5)
                s.system_health = max(0.0, s.system_health - 0.06)
                self._remove_task(task["task_id"], s)
                partial_reward = value * priority * 0.25
                s.cumulative_value += partial_reward
                return partial_reward, (
                    f"Partial failure on '{task['name']}' (uncertainty={uncertainty:.2f}). "
                    f"Partial reward: {partial_reward:.2f}."
                )

            # Success
            s.resources["energy"] = max(0.0, s.resources["energy"] - energy_cost)
            s.resources["budget"] = max(0.0, s.resources["budget"] - budget_cost)
            s.resources["time"] = max(0.0, s.resources["time"] - 1.0)
            task["status"] = "completed"
            self._remove_task(task["task_id"], s)
            s.completed_tasks.append(task)
            s.tasks_completed += 1

            # Reward: base value × priority, bonus for early completion
            deadline_bonus = max(0.0, deadline / max(time_left, 1)) * 0.5
            reward = value * priority + deadline_bonus
            s.cumulative_value += reward
            return reward, (
                f"Successfully executed '{task['name']}'. "
                f"Reward: {reward:.2f} (value={value:.1f}, priority={priority:.2f})."
            )

        elif action_type == ActionType.DEFER:
            # Low penalty; task goes to deferred queue
            task["status"] = "deferred"
            self._remove_task(task["task_id"], s)
            self._deferred_tasks.append(task)
            self._sync_state_queues()
            s.resources["time"] = max(0.0, s.resources["time"] - 0.5)
            defer_penalty = -0.2 * priority  # higher priority = bigger penalty for deferring
            return defer_penalty, (
                f"Deferred '{task['name']}'. "
                f"Penalty: {defer_penalty:.2f}. Will retry when resources recover."
            )

        elif action_type == ActionType.DELEGATE:
            # Offload — no resource cost, reduced reward
            task["status"] = "completed"
            self._remove_task(task["task_id"], s)
            s.completed_tasks.append(task)
            s.tasks_completed += 1
            delegate_reward = value * priority * 0.35
            s.cumulative_value += delegate_reward
            return delegate_reward, (
                f"Delegated '{task['name']}'. "
                f"Reward: {delegate_reward:.2f} (35% of full value)."
            )

        elif action_type == ActionType.OPTIMIZE:
            # Spend a small energy/budget to reduce uncertainty
            opt_energy = max(0.3, energy_cost * 0.2)
            opt_budget = max(1.0, budget_cost * 0.15)
            if s.resources["energy"] < opt_energy:
                return -0.1, f"Cannot optimize '{task['name']}': not enough energy."

            s.resources["energy"] = max(0.0, s.resources["energy"] - opt_energy)
            s.resources["budget"] = max(0.0, s.resources["budget"] - opt_budget)
            s.resources["time"] = max(0.0, s.resources["time"] - 0.5)

            # Reduce uncertainty significantly
            reduction = self._rng.uniform(0.2, 0.45)
            old_unc = task["uncertainty"]
            task["uncertainty"] = max(0.02, task["uncertainty"] - reduction)
            return 0.1, (
                f"Optimized '{task['name']}': uncertainty {old_unc:.2f} → {task['uncertainty']:.2f}. "
                f"Small positive reward for risk reduction."
            )

        return 0.0, "Unknown action type."

    def _find_task(self, task_id: int, task_list: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        for t in task_list:
            if t["task_id"] == task_id and t["status"] in ("pending", "deferred"):
                return t
        return None

    def _remove_task(self, task_id: int, s: AetherTaskFlowState) -> None:
        s.tasks = [t for t in s.tasks if t["task_id"] != task_id]
        self._deferred_tasks = [t for t in self._deferred_tasks if t["task_id"] != task_id]
        self._sync_state_queues()

    def _sync_state_queues(self) -> None:
        self._state.deferred_tasks = [dict(task) for task in self._deferred_tasks]

    def _max_positive_step_reward(self) -> float:
        """Upper-bound the raw positive reward for the current difficulty profile."""
        priority_max = float(self._profile["priority_range"][1])
        value_max = float(self._profile["value_range"][1])
        deadline_max = float(self._profile["deadline_range"][1])
        return max(1.0, (value_max * priority_max) + (deadline_max * 0.5))

    def _normalize_step_reward(self, raw_reward: float) -> float:
        """
        Map raw action rewards into [0, 1] for OpenEnv-facing observations.

        Negative rewards occupy [0.0, 0.5), zero remains a neutral midpoint in
        the action-reward scale, and positive rewards occupy (0.5, 1.0].
        Non-action lifecycle observations (reset/already-done) still emit their
        explicit reward values directly via _build_obs.
        """
        min_reward = -1.0
        if raw_reward <= 0.0:
            normalized = ((raw_reward - min_reward) / (0.0 - min_reward)) * 0.5
        else:
            max_reward = self._max_positive_step_reward()
            normalized = 0.5 + 0.5 * min(raw_reward / max_reward, 1.0)
        return round(max(0.0, min(1.0, normalized)), 4)

    def _get_obs(self) -> Dict[str, Any]:
        """Return a compact state snapshot for manual debugging and simple UIs."""
        s = self._state
        self._sync_state_queues()
        return {
            "episode_id": s.episode_id,
            "difficulty": s.difficulty,
            "step_count": s.step_count,
            "num_tasks": len(s.tasks),
            "num_deferred_tasks": len(self._deferred_tasks),
            "tasks_completed": s.tasks_completed,
            "tasks_failed": s.tasks_failed,
            "resources": {
                "time": round(s.resources.get("time", 0.0), 2),
                "energy": round(s.resources.get("energy", 0.0), 2),
                "budget": round(s.resources.get("budget", 0.0), 2),
            },
            "system_health": round(s.system_health, 2),
            "done": s.episode_done,
        }

    def _make_task_info(self, t: Dict[str, Any]):
        from env.tasks import TaskInfo as _TI
        return _TI(
            task_id=t["task_id"],
            name=t["name"],
            priority=t["priority"],
            deadline=t["deadline"],
            uncertainty=t["uncertainty"],
            value=t["value"],
            required_energy=t["required_energy"],
            required_budget=t["required_budget"],
            category=t["category"],
            status=t["status"],
        )

    def _build_obs(
        self,
        last_action_type: Optional[str],
        last_action_task_id: Optional[int],
        last_action_outcome: Optional[str],
        reward: float,
        done: bool,
    ) -> AetherTaskFlowObservation:
        s = self._state
        return AetherTaskFlowObservation(
            done=done,
            reward=reward,
            metadata={
                "difficulty": s.difficulty,
                "episode_id": s.episode_id,
                "step_count": s.step_count,
                "summary": self._get_obs(),
            },
            tasks=list(s.tasks),
            time_remaining=int(s.resources.get("time", 0)),
            energy_remaining=round(s.resources.get("energy", 0.0), 2),
            budget_remaining=round(s.resources.get("budget", 0.0), 2),
            system_health=round(s.system_health, 4),
            step_number=s.step_count,
            tasks_completed=s.tasks_completed,
            tasks_failed=s.tasks_failed,
            cumulative_value=round(s.cumulative_value, 4),
            last_action_type=last_action_type,
            last_action_task_id=last_action_task_id,
            last_action_outcome=last_action_outcome,
            difficulty=s.difficulty,
            episode_id=s.episode_id,
        )

    def _build_grade_result(self) -> Dict[str, Any]:
        s = self._state
        profile = self._profile
        return {
            "difficulty": s.difficulty,
            "tasks_completed": s.tasks_completed,
            "tasks_failed": s.tasks_failed,
            "total_tasks": s.tasks_completed + s.tasks_failed + len(s.tasks) + len(self._deferred_tasks),
            "remaining_time": s.resources.get("time", 0),
            "remaining_energy": s.resources.get("energy", 0),
            "remaining_budget": s.resources.get("budget", 0),
            "initial_time": s.initial_resources.get("time", profile["initial_time"]),
            "initial_energy": s.initial_resources.get("energy", profile["initial_energy"]),
            "initial_budget": s.initial_resources.get("budget", profile["initial_budget"]),
            "system_health": s.system_health,
            "steps_used": s.step_count,
            "max_steps": profile["max_steps"],
            "cumulative_value": s.cumulative_value,
        }

    def compute_final_score(self) -> float:
        """Compute the final grade [0, 1] for the completed episode."""
        result = self._build_grade_result()
        score = grade(self._difficulty, result)
        return max(0.0, min(1.0, score))