#!/usr/bin/env bash
# wm-collect: Harvest results from completed wm-port agent sessions.
#
# Scans all wm-port-* sessions. For each idle session, reads the result,
# updates done.txt, saves the log, and kills the session.
#
# Usage:
#   wm-collect              # collect all finished sessions
#   wm-collect --dry-run    # show what would be collected, don't act

set -euo pipefail

ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
TOOLS="$ROOT/tools/wm-port"
MOUNT="${OLLIE:-$HOME/mnt/ollie}"
DONE_FILE="$TOOLS/done.txt"
LOG_DIR="$TOOLS/logs"

source "$TOOLS/task-map.sh"
mkdir -p "$LOG_DIR"
touch "$DONE_FILE"

dry_run=0
[ "${1:-}" = "--dry-run" ] && dry_run=1
[ "${1:-}" = "--help" ] || [ "${1:-}" = "-h" ] && {
    sed -n '2,/^$/{ s/^# //; s/^#//; p; }' "$0"; exit 0
}

collected=0
skipped=0

while IFS=$'\t' read -r sid state cwd backend model; do
    [[ "$sid" == wm-port-* ]] || continue
    tid="${sid#wm-port-}"
    session_path="$MOUNT/s/$sid"

    # Skip sessions still running
    if [ "$state" != "idle" ]; then
        (( ++skipped ))
        continue
    fi

    # Skip already-collected tasks
    task_is_done "$tid" "$DONE_FILE" 2>/dev/null && {
        if [ "$dry_run" -eq 0 ]; then
            echo kill > "$session_path/ctl" 2>/dev/null || true
        fi
        continue
    }

    # Read the result
    local_offset=$(cat "$session_path/offset" 2>/dev/null || echo 0)
    result=$(tail -c +$((local_offset + 1)) "$session_path/chat" 2>/dev/null || true)

    # Determine status
    status="UNKNOWN"
    if [ -z "$result" ] || echo "$result" | head -1 | grep -q '^user:'; then
        status="FAIL"
    elif echo "$result" | grep -qiE '\*\*DONE\*\*|^DONE:|^DONE\b|Status:.*DONE|DONE: task completed|DONE: verified'; then
        status="DONE"
    elif echo "$result" | grep -qiE '\*\*STOP\*\*|^STOP:|^STOP\b|Status:.*STOP|STOP: blocked'; then
        status="STOP"
    elif echo "$result" | grep -qiE '\*\*FAIL\*\*|^FAIL:|^FAIL\b|Status:.*FAIL|FAIL: could not'; then
        status="FAIL"
    elif echo "$result" | grep -qiE 'compiles clean|no action needed'; then
        status="DONE"
    fi

    # Report
    case "$status" in
        DONE)
            printf '\033[1;32m\u2713 %s: %s\033[0m\n' "$tid" "$status"
            ;;
        STOP)
            printf '\033[1;33m\u26a0 %s: %s\033[0m\n' "$tid" "$status"
            ;;
        FAIL)
            printf '\033[1;31m\u2717 %s: %s\033[0m\n' "$tid" "$status"
            ;;
        *)
            printf '\033[1;35m? %s: %s\033[0m\n' "$tid" "$status"
            ;;
    esac

    if [ "$dry_run" -eq 1 ]; then
        (( ++collected ))
        continue
    fi

    # Save log
    printf '%s\n' "$result" > "$LOG_DIR/${tid}.log"

    # Update done.txt for successful tasks
    case "$status" in
        DONE)
            echo "$tid" >> "$DONE_FILE"
            echo kill > "$session_path/ctl" 2>/dev/null || true
            ;;
        *)
            # Keep session alive for inspection
            printf '  Session: %s\n' "$session_path"
            ;;
    esac

    (( ++collected ))
done < <(cat "$MOUNT/s/idx" 2>/dev/null)

if [ "$collected" -eq 0 ] && [ "$skipped" -eq 0 ]; then
    echo "No wm-port sessions found."
elif [ "$collected" -eq 0 ]; then
    printf '%d session(s) still running.\n' "$skipped"
else
    printf '\nCollected %d result(s)' "$collected"
    [ "$skipped" -gt 0 ] && printf ', %d still running' "$skipped"
    printf '.\n'
fi