#!/bin/bash
# lib-state.sh — Bash helpers for reading/writing agent state files.
# Source this in pipeline scripts: source ops/agent-state/lib-state.sh
#
# All writes use atomic rename (write to .tmp, then mv) to prevent corruption.
# All reads return valid JSON or empty string on missing/corrupt files.

STATE_ROOT="${TELEO_STATE_ROOT:-/opt/teleo-eval/agent-state}"

# --- Internal helpers ---

_state_dir() {
    local agent="$1"
    echo "$STATE_ROOT/$agent"
}

# --- Report (current status) ---

state_read_report() {
    local agent="$1"
    local file="$(_state_dir "$agent")/report.json"
    [ -f "$file" ] && cat "$file" || echo "{}"
}

state_update_report() {
    local agent="$1"
    local status="$2"
    local summary="$3"
    local file="$(_state_dir "$agent")/report.json"

    _STATE_FILE="$file" _STATE_AGENT="$agent" _STATE_STATUS="$status" \
    _STATE_SUMMARY="$summary" _STATE_TS="$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    python3 -c "
import json, os
try:
    with open(os.environ['_STATE_FILE']) as f:
        data = json.load(f)
except:
    data = {'agent': os.environ['_STATE_AGENT']}
data['status'] = os.environ['_STATE_STATUS']
data['summary'] = os.environ['_STATE_SUMMARY']
data['updated_at'] = os.environ['_STATE_TS']
print(json.dumps(data, indent=2))
" | _atomic_write_stdin "$file"
}

# Variant that takes full JSON from stdin
_atomic_write_stdin() {
    local filepath="$1"
    local tmpfile="${filepath}.tmp.$$"
    cat > "$tmpfile"
    mv -f "$tmpfile" "$filepath"
}

# Full report update with session info (called at session end)
state_finalize_report() {
    local agent="$1"
    local status="$2"
    local summary="$3"
    local session_id="$4"
    local started_at="$5"
    local ended_at="$6"
    local outcome="$7"
    local sources="$8"
    local branch="$9"
    local pr_number="${10}"
    local next_priority="${11:-null}"
    local file="$(_state_dir "$agent")/report.json"

    _STATE_FILE="$file" _STATE_AGENT="$agent" _STATE_STATUS="$status" \
    _STATE_SUMMARY="$summary" _STATE_SESSION_ID="$session_id" \
    _STATE_STARTED="$started_at" _STATE_ENDED="$ended_at" \
    _STATE_OUTCOME="$outcome" _STATE_SOURCES="$sources" \
    _STATE_BRANCH="$branch" _STATE_PR="$pr_number" \
    _STATE_NEXT="$next_priority" \
    python3 -c "
import json, os
e = os.environ
sources = int(e['_STATE_SOURCES']) if e['_STATE_SOURCES'].isdigit() else 0
pr = int(e['_STATE_PR']) if e['_STATE_PR'].isdigit() else None
next_p = None if e['_STATE_NEXT'] == 'null' else e['_STATE_NEXT']
data = {
    'agent': e['_STATE_AGENT'],
    'updated_at': e['_STATE_ENDED'],
    'status': e['_STATE_STATUS'],
    'summary': e['_STATE_SUMMARY'],
    'current_task': None,
    'last_session': {
        'id': e['_STATE_SESSION_ID'],
        'started_at': e['_STATE_STARTED'],
        'ended_at': e['_STATE_ENDED'],
        'outcome': e['_STATE_OUTCOME'],
        'sources_archived': sources,
        'branch': e['_STATE_BRANCH'],
        'pr_number': pr
    },
    'blocked_by': None,
    'next_priority': next_p
}
print(json.dumps(data, indent=2))
" | _atomic_write_stdin "$file"
}

# --- Session ---

state_start_session() {
    local agent="$1"
    local session_id="$2"
    local type="$3"
    local domain="$4"
    local branch="$5"
    local model="${6:-sonnet}"
    local timeout="${7:-5400}"
    local started_at
    started_at="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    local file="$(_state_dir "$agent")/session.json"

    _STATE_FILE="$file" _STATE_AGENT="$agent" _STATE_SID="$session_id" \
    _STATE_STARTED="$started_at" _STATE_TYPE="$type" _STATE_DOMAIN="$domain" \
    _STATE_BRANCH="$branch" _STATE_MODEL="$model" _STATE_TIMEOUT="$timeout" \
    python3 -c "
import json, os
e = os.environ
data = {
    'agent': e['_STATE_AGENT'],
    'session_id': e['_STATE_SID'],
    'started_at': e['_STATE_STARTED'],
    'ended_at': None,
    'type': e['_STATE_TYPE'],
    'domain': e['_STATE_DOMAIN'],
    'branch': e['_STATE_BRANCH'],
    'status': 'running',
    'model': e['_STATE_MODEL'],
    'timeout_seconds': int(e['_STATE_TIMEOUT']),
    'research_question': None,
    'belief_targeted': None,
    'disconfirmation_target': None,
    'sources_archived': 0,
    'sources_expected': 0,
    'tokens_used': None,
    'cost_usd': None,
    'errors': [],
    'handoff_notes': None
}
print(json.dumps(data, indent=2))
" | _atomic_write_stdin "$file"

    echo "$started_at"
}

state_end_session() {
    local agent="$1"
    local outcome="$2"
    local sources="${3:-0}"
    local pr_number="${4:-null}"
    local file="$(_state_dir "$agent")/session.json"

    _STATE_FILE="$file" _STATE_OUTCOME="$outcome" _STATE_SOURCES="$sources" \
    _STATE_PR="$pr_number" _STATE_TS="$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    python3 -c "
import json, os
e = os.environ
with open(e['_STATE_FILE']) as f:
    data = json.load(f)
data['ended_at'] = e['_STATE_TS']
data['status'] = e['_STATE_OUTCOME']
data['sources_archived'] = int(e['_STATE_SOURCES']) if e['_STATE_SOURCES'].isdigit() else 0
pr = e.get('_STATE_PR', 'null')
data['pr_number'] = int(pr) if pr.isdigit() else None
print(json.dumps(data, indent=2))
" | _atomic_write_stdin "$file"
}

# --- Journal (append-only JSONL) ---

state_journal_append() {
    local agent="$1"
    local event="$2"
    shift 2
    # Remaining args are key=value pairs for extra fields
    local file="$(_state_dir "$agent")/journal.jsonl"

    _STATE_TS="$(date -u +%Y-%m-%dT%H:%M:%SZ)" _STATE_EVT="$event" \
    python3 -c "
import json, os, sys
entry = {'ts': os.environ['_STATE_TS'], 'event': os.environ['_STATE_EVT']}
for pair in sys.argv[1:]:
    k, _, v = pair.partition('=')
    if k:
        entry[k] = v
print(json.dumps(entry))
" "$@" >> "$file"
}

# --- Metrics ---

state_update_metrics() {
    local agent="$1"
    local outcome="$2"
    local sources="${3:-0}"
    local file="$(_state_dir "$agent")/metrics.json"

    _STATE_FILE="$file" _STATE_AGENT="$agent" _STATE_OUTCOME="$outcome" \
    _STATE_SOURCES="$sources" _STATE_TS="$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    python3 -c "
import json, os
e = os.environ
try:
    with open(e['_STATE_FILE']) as f:
        data = json.load(f)
except:
    data = {'agent': e['_STATE_AGENT'], 'lifetime': {}, 'rolling_30d': {}}

lt = data.setdefault('lifetime', {})
lt['sessions_total'] = lt.get('sessions_total', 0) + 1
outcome = e['_STATE_OUTCOME']
if outcome == 'completed':
    lt['sessions_completed'] = lt.get('sessions_completed', 0) + 1
elif outcome == 'timeout':
    lt['sessions_timeout'] = lt.get('sessions_timeout', 0) + 1
elif outcome == 'error':
    lt['sessions_error'] = lt.get('sessions_error', 0) + 1
lt['sources_archived'] = lt.get('sources_archived', 0) + (int(e['_STATE_SOURCES']) if e['_STATE_SOURCES'].isdigit() else 0)

data['updated_at'] = e['_STATE_TS']
print(json.dumps(data, indent=2))
" | _atomic_write_stdin "$file"
}

# --- Inbox ---

state_check_inbox() {
    local agent="$1"
    local inbox="$(_state_dir "$agent")/inbox"
    [ -d "$inbox" ] && ls "$inbox"/*.json 2>/dev/null || true
}

state_send_message() {
    local from="$1"
    local to="$2"
    local type="$3"
    local subject="$4"
    local body="$5"
    local inbox="$(_state_dir "$to")/inbox"
    local msg_id="msg-$(date +%s)-$$"
    local file="$inbox/${msg_id}.json"

    mkdir -p "$inbox"
    _STATE_FILE="$file" _STATE_MSGID="$msg_id" _STATE_FROM="$from" \
    _STATE_TO="$to" _STATE_TS="$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    _STATE_TYPE="$type" _STATE_SUBJECT="$subject" _STATE_BODY="$body" \
    python3 -c "
import json, os
e = os.environ
data = {
    'id': e['_STATE_MSGID'],
    'from': e['_STATE_FROM'],
    'to': e['_STATE_TO'],
    'created_at': e['_STATE_TS'],
    'type': e['_STATE_TYPE'],
    'priority': 'normal',
    'subject': e['_STATE_SUBJECT'],
    'body': e['_STATE_BODY'],
    'source_ref': None,
    'expires_at': None
}
print(json.dumps(data, indent=2))
" | _atomic_write_stdin "$file"
    echo "$msg_id"
}

# --- State directory check ---

state_ensure_dir() {
    local agent="$1"
    local dir="$(_state_dir "$agent")"
    if [ ! -d "$dir" ]; then
        echo "ERROR: Agent state not initialized for $agent. Run bootstrap.sh first." >&2
        return 1
    fi
}