study_app/backend/server.py

"""
server.py
=========

FastAPI backend for the TFG decision-personalization study.

Run it (from the repository root) with:

    python -m uvicorn apps.user_study_app.backend.server:app --reload --port 8000

or use the helper scripts in ``apps/user_study_app/``
(``run_study_app.bat`` on Windows, ``run_study_app.sh`` elsewhere). Then open
http://127.0.0.1:8000/ in a browser.

Endpoints
---------
GET  /                 -> the study frontend (index.html)
GET  /api/health       -> liveness + whether the model loaded
GET  /api/config       -> public study config (model version, scenario count, ...)
POST /api/decisions    -> per-scenario A/B decisions (randomized slots, hidden mapping)
POST /api/save         -> persist a full session to CSV/JSON

The experimental condition behind each A/B slot is returned to the browser so it
can be stored, but it is NEVER displayed to participants (see frontend/app.js).
"""

from __future__ import annotations

import hashlib
import json
import os
import random
from pathlib import Path
from typing import Any

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from . import scenario_catalog
from .decision_text import render_decision_text, short_label
from .model_inference_adapter import (
    CONDITION_NON_PERSONALIZED,
    CONDITION_PERSONALIZED,
    build_adapter_from_config,
)
from .storage import StudyStorage
from .taxonomy import sanitize_snapshot


# --------------------------------------------------------------------------- #
# Paths & configuration
# --------------------------------------------------------------------------- #
APP_DIR = Path(__file__).resolve().parent.parent          # .../app
FRONTEND_DIR = APP_DIR / "frontend"

# Which config (model) this server instance serves. An env var lets you run a
# SECOND questionnaire (different model) without editing files: e.g. point it at
# config_new_model.json and launch on another port. Relative paths resolve from
# the app dir. Defaults to the original config.json so existing setups are
# unchanged. Inside a Space build the selected config is baked in as config.json.
_env_config = os.environ.get("STUDY_CONFIG_PATH", "").strip()
if _env_config:
    _cfg_path = Path(_env_config)
    if _cfg_path.is_absolute():
        CONFIG_PATH = _cfg_path
    elif (APP_DIR / _cfg_path).exists():
        # bare filename / app-dir-relative (e.g. "config_new_model.json")
        CONFIG_PATH = APP_DIR / _cfg_path
    else:
        # otherwise treat as relative to the current working directory (e.g. a
        # repo-root-relative "apps/user_study_app/config_new_model.json")
        CONFIG_PATH = _cfg_path
else:
    CONFIG_PATH = APP_DIR / "config.json"


def _load_config() -> dict[str, Any]:
    if CONFIG_PATH.exists():
        try:
            return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
        except json.JSONDecodeError as exc:  # pragma: no cover - config typo guard
            raise RuntimeError(f"Invalid config.json: {exc}") from exc
    return {}


CONFIG = _load_config()

from . import results_store  # noqa: E402
from .model_inference_adapter import REPO_ROOT  # noqa: E402  (REPO_ROOT computed there)

# Storage directory: an env var (used in cloud deployments) overrides config.json,
# which overrides the default local path. Absolute env paths are honored as-is.
_env_storage_dir = os.environ.get("STUDY_COLLECTED_DATA_DIR", "").strip()
_storage_dir = (
    CONFIG.get("storage", {}).get("collected_data_dir", "apps/user_study_app/collected_data")
)
if _env_storage_dir:
    STORAGE = StudyStorage(Path(_env_storage_dir))
else:
    STORAGE = StudyStorage(REPO_ROOT / _storage_dir)

# Optional remote sink (e.g. a private Hugging Face Dataset) for shared hosting.
# A no-op unless the relevant environment variables are set.
REMOTE_SINK = results_store.build_sink_from_env()

ADAPTER = build_adapter_from_config(CONFIG)

STUDY_CFG = CONFIG.get("study", {}) or {}
SCENARIOS_PER_SESSION = STUDY_CFG.get("scenarios_per_session", None)
SHUFFLE_SCENARIOS = bool(STUDY_CFG.get("shuffle_scenarios", True))
REQUIRE_DIFFERENT_AB = bool(STUDY_CFG.get("require_different_ab", True))
MIN_SCENARIOS = int(STUDY_CFG.get("min_scenarios", 0) or 0)
# Diversity: how many situations may share the SAME robot action (action_text).
# 1 = every shown situation has a distinct action (no "take out the trash" twice).
MAX_PER_ACTION = int(STUDY_CFG.get("max_situations_per_action", 1) or 1)
ENABLE_WRONG = bool(STUDY_CFG.get("enable_wrong_preferences_condition", False))

_MODEL_LOAD_ERROR: str | None = None


# --------------------------------------------------------------------------- #
# App
# --------------------------------------------------------------------------- #
app = FastAPI(title="TFG Decision Personalization Study", version="1.0.0")

# Permissive CORS so the frontend also works if opened from a different origin
# (e.g. a different local port). Safe for a local research tool.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
def _startup() -> None:
    global _MODEL_LOAD_ERROR
    STORAGE.ensure_dirs()
    # Always (re)write the static scenario catalog so analysts have it on disk.
    STORAGE.write_scenario_catalog(scenario_catalog.catalog_rows())
    # Eagerly load the model so the first /api/decisions request is fast. If it
    # fails (e.g. missing checkpoint), keep the server up and report via /health.
    try:
        ADAPTER.warm_up()
    except Exception as exc:  # pragma: no cover - depends on local checkpoint
        _MODEL_LOAD_ERROR = f"{type(exc).__name__}: {exc}"


# --------------------------------------------------------------------------- #
# Request models
# --------------------------------------------------------------------------- #
class PreferenceRow(BaseModel):
    signal_name: str
    polarity: str
    weight: float | None = None  # 1-10 importance; consumed by the model


class DecisionsRequest(BaseModel):
    session_id: str | None = None
    preference_snapshot: list[PreferenceRow] = Field(default_factory=list)
    scenario_ids: list[str] | None = None


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _seeded_rng(*parts: str) -> random.Random:
    digest = hashlib.sha256("::".join(parts).encode("utf-8")).hexdigest()
    return random.Random(int(digest[:16], 16))


def _slot_payload(slot: str, result_dict: dict[str, Any], decision_text: str) -> dict[str, Any]:
    return {
        "slot": slot,
        # condition is stored by the client but never rendered to the participant.
        "condition": result_dict["condition"],
        "model_action": result_dict["model_action"],
        "decision_text": decision_text,
        "short_label": short_label(result_dict["model_action"]),
        "model_confidence": result_dict["model_confidence"],
        "model_version": result_dict["model_version"],
    }


def _build_scenario_out(
    scenario: dict[str, Any],
    snapshot: list[dict[str, str]],
    session_id: str,
) -> tuple[dict[str, Any], bool]:
    """
    Run both conditions for one scenario, lay them out in randomized A/B slots,
    and report whether the two displayed decisions DIFFER.
    """
    personalized = ADAPTER.decide_personalized(
        scenario=scenario, preference_snapshot=snapshot, user_external_id=session_id
    ).to_dict()
    non_personalized = ADAPTER.decide_non_personalized(
        scenario=scenario, user_external_id=session_id
    ).to_dict()

    p_text = render_decision_text(scenario, personalized["model_action"])
    n_text = render_decision_text(scenario, non_personalized["model_action"])
    differs = p_text != n_text

    # Randomize which condition is shown in slot A vs B (seeded -> idempotent
    # across re-requests for the same session/scenario, so reloads are safe).
    rng = _seeded_rng(session_id, scenario["scenario_id"], "ab")
    if rng.random() < 0.5:
        slot_a_res, slot_a_text = personalized, p_text
        slot_b_res, slot_b_text = non_personalized, n_text
    else:
        slot_a_res, slot_a_text = non_personalized, n_text
        slot_b_res, slot_b_text = personalized, p_text

    scenario_out: dict[str, Any] = {
        **scenario_catalog.public_view(scenario),
        "slot_a": _slot_payload("A", slot_a_res, slot_a_text),
        "slot_b": _slot_payload("B", slot_b_res, slot_b_text),
        "model_version": ADAPTER.model_version,
    }

    if ENABLE_WRONG:
        try:
            wrong = ADAPTER.decide_wrong_preferences(
                scenario=scenario, preference_snapshot=snapshot, user_external_id=session_id
            ).to_dict()
            wrong["decision_text"] = render_decision_text(scenario, wrong["model_action"])
            scenario_out["wrong_preferences"] = wrong  # not shown in the UI
        except Exception:
            pass

    return scenario_out, differs


# --------------------------------------------------------------------------- #
# API
# --------------------------------------------------------------------------- #
@app.get("/api/health")
def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "model_loaded": ADAPTER.is_loaded,
        "model_load_error": _MODEL_LOAD_ERROR,
        "model_version": ADAPTER.model_version,
        "checkpoint_path": str(ADAPTER.checkpoint_path),
        "remote_sink": REMOTE_SINK.status(),
        "collected_data_dir": str(STORAGE.root),
    }


@app.get("/api/config")
def public_config() -> dict[str, Any]:
    return {
        "app_name": STUDY_CFG.get("app_name") or "Assistive Robot Study",
        "theme": STUDY_CFG.get("theme") or "default",
        "model_version": ADAPTER.model_version,
        "num_scenarios": len(scenario_catalog.all_scenarios()),
        "scenarios_per_session": SCENARIOS_PER_SESSION,
        "shuffle_scenarios": SHUFFLE_SCENARIOS,
        "enable_wrong_preferences_condition": ENABLE_WRONG,
        "decision_metrics": [
            "perceived_appropriateness",
            "preference_respect",
            "usefulness",
            "trust",
            "satisfaction",
            "intrusiveness",
            "perceived_safety",
        ],
    }


@app.post("/api/decisions")
def decisions(req: DecisionsRequest) -> JSONResponse:
    """
    Compute the personalized vs non-personalized decision for each scenario and
    return them in randomized A/B slots. The mapping (which slot is which
    condition) is included so the client can store it, but the frontend must not
    display it.
    """
    if _MODEL_LOAD_ERROR is not None and not ADAPTER.is_loaded:
        return JSONResponse(
            status_code=503,
            content={
                "error": "model_not_loaded",
                "detail": _MODEL_LOAD_ERROR,
                "hint": "Fix model.checkpoint_path in config.json and restart the backend.",
            },
        )

    session_id = req.session_id or "anonymous_session"
    snapshot = sanitize_snapshot([row.model_dump() for row in req.preference_snapshot])

    # Build a candidate pool TAILORED to the participant: situations anchored to
    # the signals they declared a preference on come first (most likely to make
    # personalized vs non-personalized differ), then the rest as fallback.
    explicit = bool(req.scenario_ids)
    if explicit:
        pool = [s for sid in req.scenario_ids if (s := scenario_catalog.get_scenario(sid))]
    else:
        pool = scenario_catalog.select_for_snapshot(
            snapshot,
            seed=int(_seeded_rng(session_id, "order").random() * 1_000_000),
            shuffle=SHUFFLE_SCENARIOS,
        )

    target = SCENARIOS_PER_SESSION if SCENARIOS_PER_SESSION else len(pool)

    # We "reroll" through the pool, running the model on each candidate, keeping
    # only those whose two decisions differ, until we have `target`. Evaluation is
    # capped so the request stays fast even in the rare low-differ case.
    max_eval = max(target * 20, 160)

    def _action_key(scn: dict[str, Any]) -> str:
        return str((scn.get("action_input", {}) or {}).get("action_text", "")).strip().lower()

    out_scenarios: list[dict[str, Any]] = []
    fallback: list[tuple[str, dict[str, Any]]] = []  # (action_key, scenario_out)
    action_counts: dict[str, int] = {}
    evaluated = 0
    for scenario in pool:
        if len(out_scenarios) >= target or evaluated >= max_eval:
            break
        act = _action_key(scenario)
        # Diversity guard: never show the same robot action more than MAX_PER_ACTION
        # times (so the set isn't "take out the trash" at 10 different hours).
        if not explicit and action_counts.get(act, 0) >= MAX_PER_ACTION:
            continue
        evaluated += 1
        try:
            scenario_out, differs = _build_scenario_out(scenario, snapshot, session_id)
        except Exception as exc:
            return JSONResponse(
                status_code=500,
                content={
                    "error": "inference_failed",
                    "scenario_id": scenario["scenario_id"],
                    "detail": f"{type(exc).__name__}: {exc}",
                },
            )
        if REQUIRE_DIFFERENT_AB and not explicit and not differs:
            fallback.append((act, scenario_out))
            continue
        out_scenarios.append(scenario_out)
        action_counts[act] = action_counts.get(act, 0) + 1

    # Floor: if too few differed, top up to min_scenarios with coinciding ones,
    # still respecting the per-action diversity cap.
    need = max(min(MIN_SCENARIOS, target), 1 if not out_scenarios else 0)
    for act, sc_out in fallback:
        if len(out_scenarios) >= need:
            break
        if action_counts.get(act, 0) >= MAX_PER_ACTION:
            continue
        out_scenarios.append(sc_out)
        action_counts[act] = action_counts.get(act, 0) + 1

    return JSONResponse(
        content={
            "model_version": ADAPTER.model_version,
            "session_id": session_id,
            "snapshot_size": len(snapshot),
            "scenarios": out_scenarios,
            "num_differing": sum(
                1 for s in out_scenarios
                if s["slot_a"]["decision_text"] != s["slot_b"]["decision_text"]
            ),
            "num_evaluated": evaluated,
        }
    )


@app.post("/api/save")
def save(payload: dict[str, Any]) -> JSONResponse:
    """
    Persist a session. Two independent sinks are attempted:
      - local CSV/JSON (always; on a cloud host this disk is ephemeral),
      - the remote sink (e.g. a private HF Dataset) if configured.
    The save is considered successful if EITHER sink succeeded, so a participant
    on a cloud host is fine even if the ephemeral local disk write fails.
    """
    participant = payload.get("participant") or {}
    session_id = str(participant.get("session_id") or payload.get("session_id") or "")
    counts = {
        "preferences": len(payload.get("preferences") or []),
        "decision_trials": len(payload.get("decision_trials") or []),
        "ab_comparisons": len(payload.get("ab_comparisons") or []),
    }

    local_result: dict[str, Any] = {"ok": False}
    try:
        local_summary = STORAGE.save_session(payload)
        local_result = {"ok": True, "export_file": local_summary.get("export_file")}
        session_id = local_summary.get("session_id", session_id)
    except Exception as exc:  # pragma: no cover - e.g. read-only fs
        local_result = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}

    remote_result: dict[str, Any] = {"ok": False, "skipped": True}
    if REMOTE_SINK.configured:
        try:
            remote_result = REMOTE_SINK.save_session(payload)
        except Exception as exc:
            remote_result = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}

    ok = bool(local_result.get("ok") or remote_result.get("ok"))
    status_code = 200 if ok else 500
    return JSONResponse(
        status_code=status_code,
        content={
            "ok": ok,
            "session_id": session_id,
            "counts": counts,
            "local": local_result,
            "remote": remote_result,
        },
    )


# --------------------------------------------------------------------------- #
# Frontend (served last so /api/* takes precedence)
# --------------------------------------------------------------------------- #
@app.get("/")
def index() -> FileResponse:
    return FileResponse(FRONTEND_DIR / "index.html")


if FRONTEND_DIR.is_dir():
    app.mount("/", StaticFiles(directory=str(FRONTEND_DIR), html=True), name="frontend")