server/app.py

import os
import threading
import uuid
from typing import List, Optional

from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from openenv.core.env_server import create_app
from pydantic import BaseModel

from models import AdaptiveAction, AdaptiveObservation
from server.adaptive_world_environment import AdaptiveWorldEnvironment
from server.mock_api import router as mock_router

_HERE   = os.path.dirname(os.path.abspath(__file__))   # server/
_ROOT   = os.path.dirname(_HERE)                        # repo root
_STATIC = os.path.join(_ROOT, "static")                # repo root/static/

# In-memory store for stateful 2-phase episodes
_sessions: dict = {}
_sessions_lock = threading.Lock()
import asyncio as _asyncio
_episode_lock = _asyncio.Lock()

app = create_app(AdaptiveWorldEnvironment, AdaptiveAction, AdaptiveObservation,
                 env_name="adaptive_world_env")

app.include_router(mock_router)

# Serve static assets (SPA + chart images)
app.mount("/static", StaticFiles(directory=_STATIC), name="static")


@app.get("/", include_in_schema=False)
async def serve_spa():
    # HF/proxy/browser caches often keep old index.html after git push; force revalidation.
    path = os.path.join(_STATIC, "index.html")
    return FileResponse(
        path,
        headers={
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Pragma": "no-cache",
            "Expires": "0",
        },
    )


@app.get("/health")
async def health():
    return {"status": "ok", "environment": "adaptive-world-env", "version": "2.3"}


# ── /run_episode  ─────────────────────────────────────────────────────────────
# OpenEnv's /reset and /step are stateless (fresh env per request).
# This endpoint runs a full multi-step episode inside a SINGLE env instance
# so state (step_count, drift_injected, world_truth) accumulates correctly.

class EpisodeRequest(BaseModel):
    scenario_id: str = "auto"
    difficulty: str = "easy"
    actions: List[dict]


class StepResult(BaseModel):
    step: int
    action_type: str
    status_code: Optional[int] = None
    response_body: Optional[str] = None
    feedback: Optional[str] = None
    done: bool = False
    task_reward: Optional[float] = None
    belief_accuracy: Optional[float] = None
    reward: Optional[float] = None


class EpisodeResponse(BaseModel):
    task_reward: float
    belief_accuracy: float
    reward: float
    steps_taken: int
    task_completed: bool
    steps: List[StepResult]


@app.post("/run_episode", response_model=EpisodeResponse)
async def run_episode(req: EpisodeRequest) -> EpisodeResponse:
    """
    Run a full episode in a single env instance.

    Each env.step() is run in a thread-pool executor so the asyncio event
    loop stays free to process the internal localhost HTTP calls that
    _execute_api_call makes to /mock_api/*.  Without this, the async
    handler blocks the event loop and those self-calls deadlock (status=0).
    """
    import asyncio
    from concurrent.futures import ThreadPoolExecutor

    loop = asyncio.get_event_loop()
    # Dedicated single-thread executor keeps env calls sequential
    executor = ThreadPoolExecutor(max_workers=1)

    def _reset():
        env = AdaptiveWorldEnvironment()
        env.reset(scenario_id=req.scenario_id, difficulty=req.difficulty)
        return env

    env: AdaptiveWorldEnvironment = await loop.run_in_executor(executor, _reset)

    step_results: List[StepResult] = []
    final_task_reward = 0.0
    final_belief_accuracy = 0.0
    final_reward = 0.001
    final_done = False

    for raw_action in req.actions:
        try:
            action = AdaptiveAction(**raw_action)
        except Exception:
            continue

        obs: AdaptiveObservation = await loop.run_in_executor(
            executor, lambda a=action: env.step(a)
        )

        sr = StepResult(
            step=env.state.step_count,
            action_type=action.action_type,
            status_code=obs.last_status_code,
            response_body=(obs.last_response_body or "")[:500],
            feedback=(obs.step_feedback or "")[:300],
            done=obs.done,
        )

        if obs.done:
            sr.task_reward = obs.task_reward
            sr.belief_accuracy = obs.belief_accuracy
            sr.reward = obs.reward
            final_task_reward = float(obs.task_reward or 0.0)
            final_belief_accuracy = float(obs.belief_accuracy or 0.0)
            final_reward = float(obs.reward or 0.001)
            final_done = True

        step_results.append(sr)

        if final_done:
            break

    # If episode didn't end via submit_result, force-submit with no belief
    if not final_done:
        obs = await loop.run_in_executor(
            executor,
            lambda: env.step(AdaptiveAction(action_type="submit_result", belief_state={}))
        )
        final_task_reward = float(obs.task_reward or 0.0)
        final_belief_accuracy = float(obs.belief_accuracy or 0.0)
        final_reward = float(obs.reward or 0.001)
        step_results.append(StepResult(
            step=env.state.step_count,
            action_type="submit_result",
            done=True,
            task_reward=final_task_reward,
            belief_accuracy=final_belief_accuracy,
            reward=final_reward,
        ))

    executor.shutdown(wait=False)

    return EpisodeResponse(
        task_reward=final_task_reward,
        belief_accuracy=final_belief_accuracy,
        reward=final_reward,
        steps_taken=env.state.step_count,
        task_completed=env.state.task_completed,
        steps=step_results,
    )


# ── /start_episode + /finish_episode  ─────────────────────────────────────────
# 3-phase belief training:
#   Phase 1 (notebook): model generates task action
#   Phase 2 (server):   /start_episode runs pre-drift + query_history + probe_schema
#                       returns session_id + probe evidence to notebook
#   Phase 3 (notebook): model generates belief FROM probe evidence
#                       calls /finish_episode with belief → gets final scores

class StartEpisodeRequest(BaseModel):
    scenario_id: str = "auto"
    difficulty:  str = "easy"
    task_action: dict


class StartEpisodeResponse(BaseModel):
    session_id:       str
    probe_response:   str
    history_response: str
    pre_drift_ok:     bool


class FinishEpisodeRequest(BaseModel):
    session_id:    str
    task_action:   dict
    belief_state:  dict


@app.post("/start_episode", response_model=StartEpisodeResponse)
async def start_episode(req: StartEpisodeRequest) -> StartEpisodeResponse:
    """
    Phase 1 of 2-phase episode.
    Runs: N × task_action → query_history → probe_schema
    Stores the live env in memory; returns probe evidence to the notebook.
    """
    import asyncio
    from concurrent.futures import ThreadPoolExecutor

    loop     = asyncio.get_event_loop()
    executor = ThreadPoolExecutor(max_workers=1)
    pre_n    = {"easy": 3, "medium": 4, "hard": 4}.get(req.difficulty, 3)

    def _run():
        env = AdaptiveWorldEnvironment()
        env.reset(scenario_id=req.scenario_id, difficulty=req.difficulty)

        task_action   = AdaptiveAction(**req.task_action)
        pre_drift_ok  = False

        for _ in range(pre_n):
            obs = env.step(task_action)
            if obs.last_status_code and obs.last_status_code < 300:
                pre_drift_ok = True

        history_obs      = env.step(AdaptiveAction(action_type="query_history", history_steps=3))
        probe_obs        = env.step(AdaptiveAction(action_type="probe_schema"))
        history_response = (history_obs.last_response_body or "")[:500]
        probe_response   = (probe_obs.last_response_body   or "")[:500]
        return env, probe_response, history_response, pre_drift_ok

    async with _episode_lock:
        env, probe_response, history_response, pre_drift_ok = await loop.run_in_executor(executor, _run)

    sid = str(uuid.uuid4())[:8]
    with _sessions_lock:
        _sessions[sid] = {"env": env, "executor": executor}

    return StartEpisodeResponse(
        session_id       = sid,
        probe_response   = probe_response,
        history_response = history_response,
        pre_drift_ok     = pre_drift_ok,
    )


@app.post("/finish_episode", response_model=EpisodeResponse)
async def finish_episode(req: FinishEpisodeRequest) -> EpisodeResponse:
    """
    Phase 2 of 2-phase episode.
    Runs: corrected task_action → submit_result with model's belief_state.
    Returns final task_reward + belief_accuracy.
    """
    import asyncio

    with _sessions_lock:
        session = _sessions.pop(req.session_id, None)

    if session is None:
        return EpisodeResponse(
            task_reward=0.0, belief_accuracy=0.0, reward=0.0,
            steps_taken=0, task_completed=False, steps=[],
        )

    env      = session["env"]
    executor = session["executor"]
    loop     = asyncio.get_event_loop()

    def _finish():
        # One corrected call after probing
        env.step(AdaptiveAction(**req.task_action))
        # Submit belief
        obs = env.step(AdaptiveAction(
            action_type  = "submit_result",
            belief_state = req.belief_state,
        ))
        return obs

    obs      = await loop.run_in_executor(executor, _finish)
    executor.shutdown(wait=False)

    return EpisodeResponse(
        task_reward     = float(obs.task_reward     or 0.0),
        belief_accuracy = float(obs.belief_accuracy or 0.0),
        reward          = float(obs.reward          or 0.0),
        steps_taken     = env.state.step_count,
        task_completed  = env.state.task_completed,
        steps           = [],
    )


def main():
    import uvicorn
    uvicorn.run("server.app:app", host="0.0.0.0", port=7860)


if __name__ == "__main__":
    main()