server/adaptive_world_environment.py

"""
adaptive_world_environment.py — Core AdaptiveWorld environment.

Replaces api_debug_environment.py from Round 1.

Key v2 features:
  - Drift injected silently (NO drift_occurred field in observation)
  - query_history action: agent reviews its own step log
  - prior_world_model: agent's beliefs carry across episodes
  - DriftDifficultyController: auto-escalates drift complexity
  - Cross-episode persistence: beliefs saved if accuracy >= 0.70
"""
import os
import json
import random
import copy
import httpx

try:
    from openenv.core.env_server import Environment
except ImportError:
    from openenv.core.env_server.interfaces import Environment

from models import AdaptiveAction, AdaptiveObservation, AdaptiveState
from scenarios.registry import SCENARIO_REGISTRY
from graders.grader import AdaptiveGrader
from server.drift_injector import DriftInjector
from server.difficulty_controller import DriftDifficultyController
from server.mock_api import (
    DYNAMIC_CONFIG,
    _issued_tokens,
    _request_log,
    _items_db,
    _orders_db,
    _bookings_db,
    _claims_db,
)

GRADER = AdaptiveGrader()
MOCK_BASE = os.getenv("MOCK_BASE_URL", "http://localhost:7860")
MAX_RESPONSE_BODY_LENGTH = 2000


def _apply_config(config: dict) -> None:
    """
    Directly mutate the shared DYNAMIC_CONFIG in-process.
    Avoids the internal HTTP round-trip to /_admin/mutate which fails
    silently in multi-worker or containerised deployments (HF Spaces).
    """
    _issued_tokens.clear()
    _request_log.clear()
    _items_db.clear()
    _orders_db.clear()
    _bookings_db.clear()
    _claims_db.clear()
    DYNAMIC_CONFIG.update(config)


def _read_schema() -> str:
    """Return the current API schema directly from DYNAMIC_CONFIG (no HTTP)."""
    schema = {
        "order_field":          DYNAMIC_CONFIG.get("order_field", "qty"),
        "required_extra":       DYNAMIC_CONFIG.get("required_extra"),
        "order_status":         DYNAMIC_CONFIG.get("order_status", "confirmed"),
        "rooms_endpoint":       DYNAMIC_CONFIG.get("rooms_endpoint", "/mock_api/rooms/book"),
        "auth_scheme":          DYNAMIC_CONFIG.get("auth_scheme", "Bearer"),
        "claims_endpoint":      DYNAMIC_CONFIG.get("claims_endpoint", "/mock_api/claims"),
        "claims_id_field":      DYNAMIC_CONFIG.get("claims_id_field", "policy_id"),
        "claims_amount_field":  DYNAMIC_CONFIG.get("claims_amount_field", "amount"),
        "product_key":          DYNAMIC_CONFIG.get("product_key", "products"),
        "product_id_field":     DYNAMIC_CONFIG.get("product_id_field", "id"),
        "product_price_field":  DYNAMIC_CONFIG.get("product_price_field", "price"),
        "max_rooms_per_request": DYNAMIC_CONFIG.get("max_rooms_per_request", 0),
    }
    return json.dumps(schema)

# Module-level controller and belief store persist across all episodes
# (as long as the server process is running — simulates a training run)
_GLOBAL_CONTROLLER = DriftDifficultyController()
_GLOBAL_BELIEFS: dict = {}   # domain → last high-confidence belief dict


class AdaptiveWorldEnvironment(Environment):
    """
    OpenEnv environment where the world mutates mid-episode.

    Agent completes multi-step professional tasks across 4 domains
    (e-commerce, hotel, flight, insurance) while API schemas drift silently.

    Tracked metrics:
      task_reward    — did the agent complete the goal?
      belief_accuracy — did the agent understand WHY it succeeded/failed?
    """

    def __init__(self):
        self._state = AdaptiveState()
        self._current_scenario = None
        self._injector: DriftInjector | None = None
        self._controller = _GLOBAL_CONTROLLER
        self._persistent_beliefs = _GLOBAL_BELIEFS
        self._step_log: list = []

    # ── reset ──────────────────────────────────────────────────────────────────

    def reset(self, scenario_id: str = "auto", **kwargs) -> AdaptiveObservation:
        difficulty = kwargs.get("difficulty", "easy")

        # Pick base scenario
        if scenario_id == "auto":
            pool = SCENARIO_REGISTRY.get(difficulty, SCENARIO_REGISTRY["easy"])
            base = copy.deepcopy(random.choice(pool))
        else:
            base = copy.deepcopy(self._find_scenario(scenario_id))

        # Apply adaptive difficulty escalation (v2)
        scenario = self._controller.get_scenario_params(base)
        self._current_scenario = scenario
        self._injector = DriftInjector(scenario["id"])
        self._step_log = []

        # Reset mock API to initial world state (direct in-process mutation)
        initial_config = self._injector.get_initial()
        _apply_config(initial_config)

        # v2: retrieve prior beliefs for this domain
        domain = scenario.get("domain", "e-commerce")
        prior_beliefs = copy.deepcopy(self._persistent_beliefs.get(domain, {}))

        self._state = AdaptiveState(
            episode_id=f"ep_{random.randint(10000, 99999)}",
            step_count=0,
            scenario_id=scenario["id"],
            domain=domain,
            drift_step=scenario["drift_trigger_step"],
            world_truth=self._injector.get_world_truth(),
        )

        return AdaptiveObservation(
            task_description=scenario["description"],
            domain=domain,
            prior_world_model=prior_beliefs,   # v2: agent sees its past knowledge
            current_step=0,
            max_steps=scenario["max_steps"],
            last_status_code=0,
            step_feedback=(
                "Episode started. Complete the task described above. "
                "The world may change mid-episode — stay alert. "
                "Use probe_schema to inspect the current API contract "
                "or query_history to compare past responses."
            ),
            difficulty_level=self._controller.level,
            done=False,
            reward=0.001,
        )

    # ── step ───────────────────────────────────────────────────────────────────

    def step(self, action: AdaptiveAction, **kwargs) -> AdaptiveObservation:
        if self._current_scenario is None:
            self.reset()

        self._state.step_count += 1

        # Check for transient error (expert EX1 scenario)
        if (self._injector and
                self._state.step_count == self._injector.get_transient_error_step()):
            return self._make_obs(
                status=503,
                body='{"error": "Service temporarily unavailable", "retryable": true}',
                feedback="503 Service Unavailable — this may be transient. "
                         "Retry before assuming a structural change."
            )

        # Inject primary drift at the configured step (SILENT — no notification)
        if (not self._state.drift_injected and
                self._state.step_count >= self._state.drift_step):
            self._inject_drift()

        # Inject secondary drift for expert scenarios
        if (self._injector and not self._injector.secondary_drifted):
            secondary_step = self._injector.get_secondary_drift_step()
            if secondary_step > 0 and self._state.step_count >= secondary_step:
                self._inject_secondary_drift()

        # Route by action type
        if action.action_type == "probe_schema":
            return self._handle_probe()

        elif action.action_type == "query_history":
            return self._handle_history(action.history_steps)

        elif action.action_type == "declare_belief":
            self._state.agent_belief = action.belief_state or {}
            return self._make_obs(
                feedback="Belief recorded. Continue with the task."
            )

        elif action.action_type == "submit_result":
            return self._finalize(action)

        else:  # "call_api"
            return self._execute_api_call(action)

    # ── internal: drift injection ──────────────────────────────────────────────

    def _inject_drift(self):
        """Silently mutate the mock API. Agent is NOT notified."""
        mutation = self._injector.inject()
        # Apply noisy_errors override from difficulty escalation if set
        noisy = self._current_scenario.get("noisy_errors", False)
        if noisy:
            mutation = dict(mutation)
            mutation["noisy_errors"] = True
        _apply_config(mutation)
        self._state.drift_injected = True
        self._state.world_truth = self._injector.get_world_truth()

    def _inject_secondary_drift(self):
        """Silently inject secondary drift for expert scenarios."""
        mutation = self._injector.inject_secondary()
        if not mutation:
            return
        _apply_config(mutation)
        self._state.world_truth = self._injector.get_world_truth()

    # ── internal: API call execution ───────────────────────────────────────────

    def _execute_api_call(self, action: AdaptiveAction) -> AdaptiveObservation:
        try:
            with httpx.Client(base_url=MOCK_BASE, timeout=5.0) as http:
                resp = http.request(
                    method=action.method.upper(),
                    url=action.url,
                    headers=action.headers,
                    json=action.body if action.body else None,
                    params=action.query_params,
                )
            status = resp.status_code
            resp_headers = dict(resp.headers)
            resp_body = resp.text[:MAX_RESPONSE_BODY_LENGTH]
        except Exception as e:
            status = 0
            resp_headers = {}
            resp_body = f"Connection error: {str(e)}"

        # Log step for query_history
        self._step_log.append({
            "step": self._state.step_count,
            "method": action.method.upper(),
            "url": action.url,
            "body": action.body,
            "status": status,
            "response": resp_body[:500],
        })

        # Track visited endpoints
        if status > 0 and action.url not in self._state.visited_endpoints:
            self._state.visited_endpoints.add(action.url)

        feedback = self._get_feedback(status)

        # Check task completion
        if status == 200 and self._check_task_completion(resp_body):
            self._state.task_completed = True
            feedback = "Task completed successfully. Please review the response for any silent changes, and use 'submit_result' to end the episode."

        done = (
            self._state.step_count >= self._current_scenario["max_steps"]
        )

        if done:
            # Always run the grader when episode ends, even without explicit submit_result
            task_reward = GRADER.grade_task(
                task_completed=self._state.task_completed,
                steps_taken=self._state.step_count,
                max_steps=self._current_scenario["max_steps"],
                drift_detected=self._drift_was_detected(),
            )
            belief_accuracy = GRADER.grade_belief(
                agent_belief=self._state.agent_belief,
                world_truth=self._state.world_truth,
                drift_type=self._current_scenario["drift_type"],
            )
            # Fallback: infer belief from action patterns if no explicit declaration
            if belief_accuracy == 0.0 and not self._state.agent_belief:
                belief_accuracy = GRADER.infer_belief_from_actions(
                    step_log=self._step_log,
                    drift_type=self._current_scenario["drift_type"],
                )

            # v2: Update difficulty controller
            self._controller.record(
                drift_type=self._current_scenario["drift_type"],
                belief_accuracy=belief_accuracy,
            )
            # v2: Persist agent's beliefs if accuracy is high enough
            if belief_accuracy >= 0.70 and self._state.agent_belief:
                domain = self._state.domain
                self._persistent_beliefs[domain] = copy.deepcopy(self._state.agent_belief)

            combined = round(task_reward * 0.7 + belief_accuracy * 0.3, 4)

            return AdaptiveObservation(
                task_description=self._current_scenario["description"],
                domain=self._state.domain,
                current_step=self._state.step_count,
                max_steps=self._current_scenario["max_steps"],
                last_status_code=status,
                last_response_body=resp_body,
                last_response_headers=resp_headers,
                task_reward=round(task_reward, 4),
                belief_accuracy=round(belief_accuracy, 4),
                step_feedback=self._summary_message(task_reward, belief_accuracy),
                difficulty_level=self._controller.level,
                done=True,
                reward=combined,
            )

        return AdaptiveObservation(
            task_description=self._current_scenario["description"],
            domain=self._state.domain,
            current_step=self._state.step_count,
            max_steps=self._current_scenario["max_steps"],
            last_status_code=status,
            last_response_body=resp_body,
            last_response_headers=resp_headers,
            step_feedback=feedback,
            difficulty_level=self._controller.level,
            done=done,
            reward=0.001,
        )

    # ── internal: probe schema ─────────────────────────────────────────────────

    def _handle_probe(self) -> AdaptiveObservation:
        # Read DYNAMIC_CONFIG directly — no HTTP round-trip needed
        schema_body = _read_schema()[:MAX_RESPONSE_BODY_LENGTH]
        status = 200

        self._step_log.append({
            "step": self._state.step_count,
            "method": "GET",
            "url": "/openapi.json",
            "status": status,
            "response": schema_body[:300],
        })

        return self._make_obs(
            status=status,
            body=schema_body,
            feedback="Schema retrieved. Compare against your prior beliefs to detect drift."
        )

    # ── internal: query history ────────────────────────────────────────────────

    def _handle_history(self, n_steps: int) -> AdaptiveObservation:
        """
        Return last N step logs. This is the v2 evidence-gathering tool.
        Agent uses this to compare pre-drift vs post-drift responses.
        """
        recent = self._step_log[-n_steps:] if self._step_log else []
        history_str = str(recent)[:MAX_RESPONSE_BODY_LENGTH]

        self._step_log.append({
            "step":     self._state.step_count,
            "method":   "GET",
            "url":      "/query_history",
            "status":   200,
            "response": history_str[:200],
        })

        return self._make_obs(
            body=history_str,
            feedback=(
                f"Last {len(recent)} API interactions shown. "
                "Compare status codes and response bodies across steps "
                "to detect changes in the world state."
            ),
            episode_history=recent,
        )

    # ── internal: finalize episode ─────────────────────────────────────────────

    def _finalize(self, action: AdaptiveAction) -> AdaptiveObservation:
        final_belief = action.belief_state or self._state.agent_belief

        task_reward = GRADER.grade_task(
            task_completed=self._state.task_completed,
            steps_taken=self._state.step_count,
            max_steps=self._current_scenario["max_steps"],
            drift_detected=self._drift_was_detected(),
        )
        belief_accuracy = GRADER.grade_belief(
            agent_belief=final_belief,
            world_truth=self._state.world_truth,
            drift_type=self._current_scenario["drift_type"],
        )

        # Fallback: infer belief from action patterns if no explicit declaration
        if belief_accuracy == 0.0 and not final_belief:
            belief_accuracy = GRADER.infer_belief_from_actions(
                step_log=self._step_log,
                drift_type=self._current_scenario["drift_type"],
            )

        combined = round(task_reward * 0.7 + belief_accuracy * 0.3, 4)

        # v2: Update difficulty controller with this episode's result
        self._controller.record(
            drift_type=self._current_scenario["drift_type"],
            belief_accuracy=belief_accuracy,
        )

        # v2: Persist agent's beliefs if accuracy is high enough
        if belief_accuracy >= 0.70 and final_belief:
            domain = self._state.domain
            self._persistent_beliefs[domain] = copy.deepcopy(final_belief)

        return AdaptiveObservation(
            task_description=self._current_scenario["description"],
            domain=self._state.domain,
            current_step=self._state.step_count,
            max_steps=self._current_scenario["max_steps"],
            task_reward=round(task_reward, 4),
            belief_accuracy=round(belief_accuracy, 4),
            difficulty_level=self._controller.level,
            step_feedback=self._summary_message(task_reward, belief_accuracy),
            done=True,
            reward=combined,
        )

    # ── helpers ────────────────────────────────────────────────────────────────

    def _drift_was_detected(self) -> bool:
        """Did the agent call probe_schema or query_history AFTER drift was injected?"""
        if not self._state.drift_injected:
            return False
        for entry in self._step_log:
            if entry["step"] > self._state.drift_step:
                url = entry.get("url", "")
                if "/openapi" in url or "_history" in url or "openapi-schema" in url:
                    return True
        return False

    def _check_task_completion(self, resp_body: str = "") -> bool:
        goal = self._current_scenario.get("task_goal", "")
        if goal == "place_order":
            return "order_id" in resp_body and "status" in resp_body
        elif goal == "book_room":
            return "booking_id" in resp_body or "confirmation" in resp_body
        elif goal == "file_claim":
            return "claim_id" in resp_body
        elif goal == "book_flight":
            return "flight_id" in resp_body or "booking" in resp_body
        elif goal == "apply_discount":
            return "discount_applied" in resp_body
        elif goal == "search_products":
            body_lower = resp_body.lower()
            key1 = "products" in body_lower
            key2 = "items" in body_lower
            return key1 or key2
        return False

    def _get_feedback(self, status: int) -> str:
        hints = {
            401: "Authentication failed. Something about your credentials or auth scheme may have changed.",
            403: "Forbidden. A policy may have changed. Read the error body carefully.",
            404: "Endpoint not found. The path may have moved — check API schema.",
            405: "Method not allowed.",
            415: "Wrong Content-Type. Use application/json.",
            422: "Validation error. A field name, type, or required field may have changed.",
            429: "Rate limited. A rule about request limits may have changed.",
            503: "Service temporarily unavailable. Consider retrying before assuming a structural change.",
        }
        if status == 200:
            return "200 OK. Verify the response structure matches what you expected."
        return hints.get(status, f"Status {status}. Check the response body for clues.")

    def _make_obs(self, status: int = 200, body: str = "", feedback: str = "",
                  episode_history: list = None) -> AdaptiveObservation:
        return AdaptiveObservation(
            task_description=self._current_scenario["description"],
            domain=self._state.domain,
            current_step=self._state.step_count,
            max_steps=self._current_scenario["max_steps"],
            last_status_code=status,
            last_response_body=body,
            step_feedback=feedback,
            episode_history=episode_history or [],
            difficulty_level=self._controller.level,
            done=False,
            reward=0.001,
        )

    def _summary_message(self, task_reward: float, belief_accuracy: float) -> str:
        return (
            f"Episode complete. "
            f"Task reward: {task_reward:.3f}. "
            f"Belief accuracy: {belief_accuracy:.3f}. "
            f"Combined reward: {task_reward * 0.7 + belief_accuracy * 0.3:.3f}. "
            f"Difficulty level: {self._controller.level}/3."
        )

    def _find_scenario(self, scenario_id: str) -> dict:
        for pool in SCENARIO_REGISTRY.values():
            for s in pool:
                if s["id"] == scenario_id:
                    return s
        raise ValueError(f"Scenario not found: {scenario_id}. "
                         f"Call reset(scenario_id='auto') to use random selection.")

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