core/logger.py

"""HF dataset logger, schema 1.2.

Schema 1.2 (socratic v4 — plan-complete turn type, end-session flush,
upload error capture, derived session summary):
  {
    schema_version: "1.2",
    session_id, condition, run_mode,
    started_at, ended_at,
    model, case_id,
    perfect_answer: {text, generated_at, char_count} or null,
    discussion_plan: {text, extracted_at} or null,
    manager_history: [verdict dicts, one per bot turn],
    turns: [
      {
        n, timestamp,
        user_message,
        manager_verdict: {live_subprobe_id, advanced_this_turn, turn_type,
                          directive, reason, latency_ms},
        plan_control: <verbatim PLAN CONTROL text injected into main bot>,
        assistant_response,
        scratch: {chosen_target, response_summary, directive_followed},
        latency_s: {ttft, total}
      }, ...
    ],
    feedback: [...],
    feedback_current: {message_index_str: latest_kind},
    argument_snapshots: [...],
    essays: {argument: {text, submitted_at, char_count}},
    upload_errors: [<error strings from failed upload attempts>],
    session_summary: {
      n_turns, n_transitions, n_plan_completes,
      subprobe_coverage: {sub_id: {engaging_default: K, other: M}, ...},
      last_subprobe_visited, total_subprobes_in_plan,
      outcome: "plan-complete-reached" | "ran-without-plan-complete",
    }
  }

What changed from 1.1:
  - schema_version: bumped to 1.2.
  - upload_errors: NEW top-level list; persists upload-failure messages so a
    truncated log shows WHY it truncated when re-uploaded.
  - session_summary: NEW top-level derived block, computed at every
    _build_payload, so coverage / transition / plan-complete stats are
    available without re-deriving from manager_history.
  - plan-complete: new valid turn_type (handled transparently in turns[]
    and manager_history; counted in session_summary).
  - End-of-session flush: SessionLogger.flush() does a synchronous upload
    with retries, called by app.on_end_session in the CONFIRMED branch so
    the final payload always lands in the dataset before the UI thanks the
    participant.
  - Upload retries: _upload now retries 3 times with exponential backoff
    (1s, 2s, 4s) before giving up and recording the error.

Path: {HF_DATASET}/{LOG_PATH}/{session_id}.json
"""

import io
import json
import queue
import secrets
import sys
import threading
import time
from datetime import datetime, timezone, timedelta
from core import config_loader
from core.config_loader import (
    HF_DATASET, MODEL, LOG_PATH,
    CONDITION, SCHEMA_VERSION, CASE_ID, RUN_MODE,
)

AEST = timezone(timedelta(hours=10))


def now_aest():
    """Wall-clock datetime in AEST. Single source of truth so display strings
    in app.py and logger timestamps stay aligned."""
    return datetime.now(AEST)


def _now_iso():
    return now_aest().isoformat()


def _unix_to_iso(unix_ts):
    """Convert unix epoch seconds (from time.time()) to AEST ISO 8601 string.
    Returns None if input is None."""
    if unix_ts is None:
        return None
    return datetime.fromtimestamp(unix_ts, tz=AEST).isoformat()


# Completion-code character set. Digits only (0-9) — short, unambiguous,
# easy for participants to type. 10^5 = 100,000 combinations — collision-
# safe for any study size in the hundreds.
_CODE_ALPHABET = "0123456789"


def _make_completion_code():
    """Generate a short, human-friendly completion code.
    Embedded verbatim in the log filename (as the digit suffix) and stored
    as a top-level field. The end-of-session modal shows the participant
    a presentation string built from this code (see _completion_modal_html
    in app.py) — the prefix decoration is added at display time and is NOT
    part of the stored code."""
    return "".join(secrets.choice(_CODE_ALPHABET) for _ in range(5))


class SessionLogger:
    """Schema 1.1 session log. Coalesces uploads via a single background
    worker thread + 1-slot queue: a burst of N writes produces at most ~2
    actual HF uploads (first + final) instead of N redundant overwrites."""

    def __init__(self):
        self.session_id = None
        self.completion_code = None    # 5-char code shown to participant at end
        self.assigned_key_idx = None   # 1-10 (pool key) / 0 (fallback API_KEY) / None
        self.start_time = None
        self.walkthrough_started_at = None  # epoch seconds (time.time()) when
                                             # participant clicked "Start
                                             # walkthrough" on the welcome
                                             # overlay; set ONCE, preserved
                                             # across start_session() resets.
        self.turns = []
        self.essays = {}                # {"argument": {...}}
        self.feedback = []
        self.perfect_answer = None
        self.discussion_plan = None
        self.dominant_framework = None  # one of: consequence, duty_rights, character
        self.argument_snapshots = []
        self.manager_history = []       # cumulative verdict dicts
        self.upload_errors = []         # accumulated upload-failure messages
        self._upload_queue = queue.Queue(maxsize=1)
        self._upload_worker = None
        self._upload_lock = threading.Lock()  # serializes flush vs background worker
        # State tracking for verify-after-upload + recovery backup.
        # _last_confirmed_state is the state-hash of the most recent payload
        # that was uploaded AND verified to exist on HF. If at flush time
        # _last_confirmed_state lags current state, we write a recovery file.
        self._last_confirmed_state = None

    def start_session(self):
        self.start_time = now_aest()
        self.session_id = self.start_time.strftime("%Y%m%d_%H%M%S")
        self.completion_code = _make_completion_code()
        self.assigned_key_idx = None
        self.turns = []
        self.essays = {}
        self.feedback = []
        self.perfect_answer = None
        self.discussion_plan = None
        self.dominant_framework = None
        self.argument_snapshots = []
        self.manager_history = []
        self.upload_errors = []
        self._last_confirmed_state = None
        self._ensure_worker()

    @classmethod
    def from_dict(cls, data):
        """Rebuild a SessionLogger from a parsed HF JSON payload, used by the
        on_recover_session handler when a participant returns to a session
        whose tab was closed before End Session.

        Preserves the original session_id + completion_code so subsequent
        uploads land in the SAME HF file. All accumulated state (essays,
        perfect_answer, discussion_plan, dominant_framework, feedback,
        argument_snapshots, upload_errors) is restored verbatim. turns and
        manager_history are typically empty in the no_ai condition (reading
        only, no chat) but restored for symmetry with chat-based spaces.

        Returns the logger, OR None if the payload is missing the minimum
        fields (session_id + completion_code) — caller falls back to fresh.
        """
        if not isinstance(data, dict):
            return None
        session_id = data.get("session_id")
        completion_code = data.get("completion_code")
        if not session_id or not completion_code:
            return None

        logger = cls()
        logger.session_id = session_id
        logger.completion_code = completion_code
        logger.assigned_key_idx = data.get("assigned_key_idx")

        started_at_str = data.get("started_at")
        try:
            logger.start_time = (
                datetime.fromisoformat(started_at_str)
                if started_at_str else now_aest()
            )
        except (ValueError, TypeError):
            logger.start_time = now_aest()

        # walkthrough_started_at is the only unix-epoch field on the no_ai
        # logger (no argument_submitted_at gate in this condition).
        v = data.get("walkthrough_started_at")
        if v:
            try:
                logger.walkthrough_started_at = (
                    datetime.fromisoformat(v).timestamp()
                )
            except (ValueError, TypeError):
                pass

        logger.turns = list(data.get("turns") or [])
        logger.manager_history = list(data.get("manager_history") or [])
        logger.argument_snapshots = list(data.get("argument_snapshots") or [])
        logger.feedback = list(data.get("feedback") or [])
        logger.essays = dict(data.get("essays") or {})
        logger.perfect_answer = data.get("perfect_answer")
        logger.discussion_plan = data.get("discussion_plan")
        logger.dominant_framework = data.get("dominant_framework")
        logger.upload_errors = list(data.get("upload_errors") or [])

        try:
            latest_arg = (logger.essays.get("argument") or {}).get("text", "")
            logger.argument_snapshots.append({
                "n": len(logger.argument_snapshots) + 1,
                "timestamp": _now_iso(),
                "trigger": "session_resumed",
                "text": latest_arg,
                "char_count": len(latest_arg),
                "changed_since_prev": False,
            })
        except Exception:
            pass

        logger._last_confirmed_state = logger._state_hash()
        logger._ensure_worker()
        logger._enqueue_upload()
        return logger

    def _filename(self):
        """Filename written to HF dataset for this session. Embeds both the
        session timestamp (for chronological sorting) and the completion code
        (for joining sessions to Qualtrics responses by code)."""
        if not self.session_id:
            return None
        if self.completion_code:
            return f"{LOG_PATH}/{self.session_id}_{self.completion_code}.json"
        return f"{LOG_PATH}/{self.session_id}.json"

    def _state_hash(self):
        """Tuple snapshot of how much data has been accumulated this session.
        Used to compare 'what we last successfully uploaded' against 'what we
        have now'. Monotonically grows; bigger tuple = newer state."""
        return (
            len(self.turns),
            len(self.manager_history),
            len(self.argument_snapshots),
            len(self.feedback),
        )

    def _ensure_worker(self):
        if self._upload_worker is None or not self._upload_worker.is_alive():
            self._upload_worker = threading.Thread(
                target=self._upload_loop, daemon=True
            )
            self._upload_worker.start()

    def set_essay(self, key, text):
        if not self.session_id:
            self.start_session()
        self.essays[key] = {
            "text": text,
            "submitted_at": _now_iso(),
            "char_count": len(text),
        }
        self._enqueue_upload()

    def set_discussion_plan(self, text):
        if not self.session_id:
            return
        self.discussion_plan = {
            "text": text,
            "extracted_at": _now_iso(),
        }
        self._enqueue_upload()

    def set_perfect_answer(self, text):
        if not self.session_id:
            return
        self.perfect_answer = {
            "text": text,
            "generated_at": _now_iso(),
            "char_count": len(text),
        }
        self._enqueue_upload()

    def set_dominant_framework(self, framework):
        """Record which of the 3 frameworks (consequence / duty_rights /
        character) the perfect-answer generator chose for this session.
        Stored for log analysis."""
        if not self.session_id:
            return
        self.dominant_framework = framework
        self._enqueue_upload()

    def add_feedback(self, message_index, kind, message_excerpt=""):
        """Record a participant aha! / ugh. reaction on a bot message.

        kind: "aha" (made something click),
              "stress" (made participant uncomfortable),
              "none" (cancel, second click on the already-active emoji).
        Every click (including cancels) is appended with its own timestamp.
        """
        if not self.session_id:
            return
        self.feedback.append({
            "n": len(self.feedback) + 1,
            "message_index": message_index,
            "kind": kind,
            "timestamp": _now_iso(),
            "message_excerpt": (message_excerpt or "")[:200],
        })
        self._enqueue_upload()

    def get_current_kind(self, message_index):
        """Return the current ('aha' / 'stress' / None) reaction for a given
        message index, derived from the latest event for that index."""
        for event in reversed(self.feedback):
            if event.get("message_index") == message_index:
                k = event.get("kind")
                return None if k == "none" else k
        return None

    def add_argument_snapshot(self, trigger, text, context=None):
        """Snapshot the participant's notepad textbox at an interaction event.
        The bot NEVER sees this; logger-only side-channel."""
        if not self.session_id:
            return
        text = text or ""
        prev = self.argument_snapshots[-1] if self.argument_snapshots else None
        snap = {
            "n": len(self.argument_snapshots) + 1,
            "timestamp": _now_iso(),
            "trigger": trigger,
            "text": text,
            "char_count": len(text),
            "changed_since_prev": (prev["text"] != text) if prev else None,
        }
        if context:
            snap["context"] = context
        self.argument_snapshots.append(snap)
        self._enqueue_upload()

    def add_manager_verdict(self, verdict_dict):
        """Append a manager verdict to the session-level history list.
        Called once per chat turn, right after the manager LLM returns.

        `verdict_dict` should be the dict form of `ManagerVerdict.to_dict()`.
        """
        if not self.session_id:
            return
        self.manager_history.append(dict(verdict_dict))
        self._enqueue_upload()

    def add_turn(self, user_message, manager_verdict, plan_control,
                 assistant_response, scratch=None, latency_s=None):
        """Append a chat turn. All fields together describe one full round:
        what the participant said, what the manager decided, what was
        injected into the main bot, what the main bot said back."""
        if not self.session_id:
            self.start_session()
        turn = {
            "n": len(self.turns) + 1,
            "timestamp": _now_iso(),
            "user_message": user_message,
            "manager_verdict": dict(manager_verdict) if manager_verdict else None,
            "plan_control": plan_control,
            "assistant_response": assistant_response,
        }
        if scratch:
            turn["scratch"] = dict(scratch)
        if latency_s:
            turn["latency_s"] = latency_s
        self.turns.append(turn)
        self._enqueue_upload()

    def _derive_feedback_current(self):
        """Collapse feedback events to {message_index: latest_kind},
        omitting messages whose latest event is a cancel ('none')."""
        latest = {}
        for event in self.feedback:
            idx = event.get("message_index")
            if idx is None:
                continue
            latest[idx] = event.get("kind")
        return {idx: k for idx, k in latest.items() if k and k != "none"}

    def _derive_session_summary(self):
        """Compute coverage / transition / plan-complete stats from
        manager_history. Returns None if no manager turns yet."""
        if not self.manager_history:
            return None
        by_sub = {}
        n_transitions = 0
        n_plan_completes = 0
        for v in self.manager_history:
            sub = v.get("live_subprobe_id")
            tt = v.get("turn_type")
            if sub is not None:
                if sub not in by_sub:
                    by_sub[sub] = {"engaging_default": 0, "other": 0}
                if tt == "engaging-default":
                    by_sub[sub]["engaging_default"] += 1
                else:
                    by_sub[sub]["other"] += 1
            if v.get("advanced_this_turn"):
                n_transitions += 1
            if tt == "plan-complete":
                n_plan_completes += 1
        # Total sub-probes in the plan, derived from discussion_plan text.
        total_in_plan = None
        if self.discussion_plan and self.discussion_plan.get("text"):
            import re
            ids = re.findall(
                r"^\s*-\s*id\s*:\s*(\d+)",
                self.discussion_plan["text"],
                re.IGNORECASE | re.MULTILINE,
            )
            if ids:
                total_in_plan = max(int(x) for x in ids)
        outcome = (
            "plan-complete-reached" if n_plan_completes > 0
            else "ran-without-plan-complete"
        )
        return {
            "n_turns": len(self.turns),
            "n_manager_verdicts": len(self.manager_history),
            "n_transitions": n_transitions,
            "n_plan_completes": n_plan_completes,
            "subprobe_coverage": by_sub,
            "last_subprobe_visited": max(by_sub) if by_sub else None,
            "total_subprobes_in_plan": total_in_plan,
            "outcome": outcome,
        }

    def set_assigned_key(self, key_idx):
        """Record which OpenAI key index was assigned to this session
        (1-10 for pool keys, 0 for fallback API_KEY)."""
        self.assigned_key_idx = key_idx

    def _build_payload(self):
        return {
            "schema_version": SCHEMA_VERSION,
            "session_id": self.session_id,
            "completion_code": self.completion_code,
            "assigned_key_idx": self.assigned_key_idx,
            "condition": CONDITION,
            "run_mode": RUN_MODE,
            "walkthrough_started_at": _unix_to_iso(self.walkthrough_started_at),
            "started_at": self.start_time.isoformat() if self.start_time else None,
            "ended_at": _now_iso(),
            "model": MODEL,
            "case_id": CASE_ID,
            "perfect_answer": self.perfect_answer,
            "discussion_plan": self.discussion_plan,
            "dominant_framework": self.dominant_framework,
            "manager_history": list(self.manager_history),
            "turns": list(self.turns),
            "feedback": list(self.feedback),
            "feedback_current": self._derive_feedback_current(),
            "argument_snapshots": list(self.argument_snapshots),
            "essays": self.essays,
            "upload_errors": list(self.upload_errors),
            "session_summary": self._derive_session_summary(),
        }

    def _enqueue_upload(self):
        """Push the latest payload onto the upload queue. If a payload is
        already queued, it gets superseded (only the newest survives).

        Payload is built CUMULATIVELY (contains all turns so far). This is
        the natural 'defer on failure' mechanism: if T_n upload failed
        silently and _last_confirmed_state didn't advance, T_{n+1}'s payload
        still contains T_n's data and a retry will catch it up."""
        if not config_loader.hf_api or not self.session_id:
            return
        payload = self._build_payload()
        # Snapshot state hash at payload-build time so verify-after-upload
        # can mark the precise level of progress this upload represents.
        state_snapshot = self._state_hash()
        try:
            self._upload_queue.get_nowait()
        except queue.Empty:
            pass
        self._upload_queue.put((payload, state_snapshot))
        self._ensure_worker()

    def _upload_loop(self):
        while True:
            payload, state_snapshot = self._upload_queue.get()
            self._upload(payload, expected_state=state_snapshot)

    def _verify_uploaded(self, filename):
        """Lightweight verify-after-upload check. We don't trust upload_file's
        return alone because HF can return 200 OK but silently drop the commit
        under rate-limit or backend-conflict conditions. Confirm the file is
        actually present by listing the dataset's files.

        Returns True iff `filename` shows up in the dataset listing right after
        upload. False on any failure (network error, file missing).
        """
        try:
            files = config_loader.hf_api.list_repo_files(
                repo_id=HF_DATASET, repo_type="dataset"
            )
            return filename in files
        except Exception:
            return False

    def _upload(self, payload, max_retries=3, expected_state=None):
        """Upload payload to HF dataset with exponential backoff retry +
        verify-after-upload. Returns True on success, False if all retries
        failed or upload silently failed verification.

        `expected_state` is the _state_hash() snapshot taken at the moment
        this payload was built; if the upload succeeds AND verifies, we mark
        `_last_confirmed_state = expected_state` so the caller (or flush)
        can tell whether the most recent state has been committed to HF.

        On failure: records the error in self.upload_errors and prints to
        stderr; does NOT update _last_confirmed_state, which means the next
        cumulative upload (next turn) will naturally retry this turn's data.
        """
        last_err = None
        filename = self._filename()
        for attempt in range(max_retries):
            try:
                content = json.dumps(payload, ensure_ascii=False, indent=2)
                with self._upload_lock:
                    config_loader.hf_api.upload_file(
                        path_or_fileobj=io.BytesIO(content.encode("utf-8")),
                        path_in_repo=filename,
                        repo_id=HF_DATASET,
                        repo_type="dataset",
                    )
                # Verify the commit actually applied — HF can silently drop
                # uploads under rate pressure. If the file isn't visible
                # right after, treat as failure and retry.
                if self._verify_uploaded(filename):
                    if expected_state is not None:
                        self._last_confirmed_state = expected_state
                    return True
                last_err = "verify-after-upload failed (file not visible after commit)"
            except Exception as e:
                last_err = e
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # 1s, 2s before final attempt
        err_str = (
            f"[{_now_iso()}] upload failed after {max_retries} attempts: "
            f"{type(last_err).__name__ if isinstance(last_err, Exception) else 'verify'}: {last_err}"
        )[:500]
        self.upload_errors.append(err_str)
        print(f"[logger] {err_str}", file=sys.stderr, flush=True)
        return False

    def _upload_recovery_backup(self, payload):
        """End-of-session safety net: if the primary upload couldn't be
        confirmed even after the extended retry schedule in flush(), write a
        clearly-named recovery file to a distinct sub-folder. Naming:

            <LOG_PATH>/_recovery/<session_id>.json

        The `_recovery/` prefix sorts separately from normal logs so an
        analyst can immediately see 'these sessions had upload trouble'.
        Inside the file, we tag the payload with a `__recovery_backup__`
        marker noting reason + intended primary path + timestamp.

        Uses the SAME extended backoff schedule as the primary flush
        (2,4,8,16,30,45,60s, ~165s worst case), since this is the absolute
        last line of defence — give HF every chance to accept the write.
        """
        suffix = (f"{self.session_id}_{self.completion_code}.json"
                  if self.completion_code else f"{self.session_id}.json")
        backup_path = f"{LOG_PATH}/_recovery/{suffix}"
        primary_path = f"{LOG_PATH}/{suffix}"
        backup_payload = dict(payload)
        backup_payload["__recovery_backup__"] = {
            "reason": "primary upload could not be confirmed before end_session",
            "written_at": _now_iso(),
            "primary_path": primary_path,
            "last_confirmed_state": self._last_confirmed_state,
            "current_state": self._state_hash(),
        }
        last_err = None
        # Extended backoff schedule matching flush() primary path.
        backoffs = [0, 2, 4, 8, 16, 30, 45]
        for i, wait in enumerate(backoffs):
            if wait > 0:
                time.sleep(wait)
            try:
                content = json.dumps(backup_payload, ensure_ascii=False, indent=2)
                with self._upload_lock:
                    config_loader.hf_api.upload_file(
                        path_or_fileobj=io.BytesIO(content.encode("utf-8")),
                        path_in_repo=backup_path,
                        repo_id=HF_DATASET,
                        repo_type="dataset",
                    )
                if self._verify_uploaded(backup_path):
                    print(f"[logger] recovery backup saved to {backup_path}",
                          file=sys.stderr, flush=True)
                    return True
                last_err = "backup verify-after-upload failed"
            except Exception as e:
                last_err = e
        print(f"[logger] CRITICAL: recovery backup ALSO failed for session "
              f"{self.session_id}: {last_err}", file=sys.stderr, flush=True)
        return False

    def flush(self):
        """End-of-session synchronous upload with extended retry budget.

        Designed for the case where per-turn uploads may have failed silently
        during the conversation (HF rate-limit / commit conflict under heavy
        concurrency). This is the LAST CHANCE to get the session data into
        the dataset, so we retry much harder than per-turn uploads:

          - Primary upload: up to 8 attempts with backoffs of
            2, 4, 8, 16, 30, 45, 60s (total ~165s worst case).
          - If primary still fails, write recovery backup to
            <LOG_PATH>/_recovery/<session_id>.json with 7 attempts at the
            same backoff schedule.

        Returns a status dict:
          {"status": "uploaded_primary",  "primary": True,  "recovery": False}
          {"status": "uploaded_recovery", "primary": False, "recovery": True}
          {"status": "all_failed",        "primary": False, "recovery": False}
          {"status": "no_logger",         "primary": False, "recovery": False}

        UI handlers should display the status to the participant so they know
        whether their data made it. `session_id` is exposed so a researcher
        can recover later if all_failed."""
        if not config_loader.hf_api or not self.session_id:
            return {"status": "no_logger", "primary": False, "recovery": False,
                    "session_id": self.session_id,
                    "completion_code": self.completion_code}
        # Drain any pending async payload — the sync flush supersedes it.
        try:
            self._upload_queue.get_nowait()
        except queue.Empty:
            pass
        payload = self._build_payload()
        current_state = self._state_hash()
        # Extended backoff schedule for end-of-session: 2,4,8,16,30,45,60s.
        backoffs = [0, 2, 4, 8, 16, 30, 45, 60]
        primary_ok = False
        for i, wait in enumerate(backoffs):
            if wait > 0:
                time.sleep(wait)
            # Single-attempt path that respects verify but does NOT do its own
            # retries (we control retry timing here).
            if self._upload_single_attempt(payload, expected_state=current_state):
                primary_ok = True
                break
        if primary_ok:
            return {"status": "uploaded_primary", "primary": True, "recovery": False,
                    "session_id": self.session_id,
                    "completion_code": self.completion_code}
        # Primary still failed — write recovery backup with extended retries.
        recovery_ok = self._upload_recovery_backup(payload)
        if recovery_ok:
            return {"status": "uploaded_recovery", "primary": False, "recovery": True,
                    "session_id": self.session_id,
                    "completion_code": self.completion_code}
        return {"status": "all_failed", "primary": False, "recovery": False,
                "session_id": self.session_id,
                "completion_code": self.completion_code}

    def _upload_single_attempt(self, payload, expected_state=None):
        """One upload + verify, no internal retry. Caller controls retry timing.
        Used by flush() which has its own extended backoff schedule."""
        filename = self._filename()
        try:
            content = json.dumps(payload, ensure_ascii=False, indent=2)
            with self._upload_lock:
                config_loader.hf_api.upload_file(
                    path_or_fileobj=io.BytesIO(content.encode("utf-8")),
                    path_in_repo=filename,
                    repo_id=HF_DATASET,
                    repo_type="dataset",
                )
            if self._verify_uploaded(filename):
                if expected_state is not None:
                    self._last_confirmed_state = expected_state
                return True
            self.upload_errors.append(
                f"[{_now_iso()}] flush attempt: verify-after-upload failed"[:500]
            )
            return False
        except Exception as e:
            self.upload_errors.append(
                f"[{_now_iso()}] flush attempt: {type(e).__name__}: {e}"[:500]
            )
            print(f"[logger] flush attempt failed: {type(e).__name__}: {e}",
                  file=sys.stderr, flush=True)
            return False