#!/bin/bash
# CLI replayer for Investibook workspaces
#
# Usage:
#   ./tools/replayer/replay.sh <workspace_id>                        # Rerun existing workspace
#   ./tools/replayer/replay.sh <workspace_id> --eval-only            # Just extract & evaluate (no rerun)
#   ./tools/replayer/replay.sh <workspace_id> --create <issue_id>    # Create new workspace + run
#   ./tools/replayer/replay.sh <workspace_id> --check                # Structured checklist (JSON)
#   ./tools/replayer/replay.sh <workspace_id> --followup "<message>" # Send follow-up message
#   ./tools/replayer/replay.sh --profile dev <ws_id> --check         # Use CLI profile for remote env
#
# Requires: backend running at localhost:8000 (or --profile for remote)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
BACKEND_DIR="$REPO_ROOT/backend"
BASE_URL="${INVESTIBOOK_BASE_URL:-http://localhost:8000}"
PROMPT="Use tidb-clinic-diagnosis to analyze this alert. Provide root cause, evidence chain, and recommendations."
MAX_TIME="${REPLAY_MAX_TIME:-600}"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

log()  { echo -e "${CYAN}[replay]${NC} $*"; }
ok()   { echo -e "${GREEN}[ok]${NC} $*"; }
warn() { echo -e "${YELLOW}[warn]${NC} $*"; }
err()  { echo -e "${RED}[error]${NC} $*" >&2; }

usage() {
    echo "Usage: $0 [--profile <name>] <workspace_id> [--eval-only | --create <issue_id> | --check | --followup <msg>]"
    echo ""
    echo "Modes:"
    echo "  <ws_id>                       Rerun: clear messages, send diagnosis prompt, wait, evaluate"
    echo "  <ws_id> --eval-only           Evaluate existing messages (no rerun)"
    echo "  <ws_id> --create <issue>      Create workspace, add Jira source, prepare, run diagnosis"
    echo "  <ws_id> --check               Structured checklist output (JSON) for agent evaluation"
    echo "  <ws_id> --followup <message>  Send follow-up message to workspace, wait for completion"
    echo ""
    echo "Options:"
    echo "  --profile <name>              Use CLI profile (reads token/url from ~/.investibook/config.json)"
    echo ""
    echo "Environment:"
    echo "  INVESTIBOOK_BASE_URL      Backend URL (default: http://localhost:8000)"
    echo "  REPLAY_MAX_TIME           Diagnosis timeout in seconds (default: 600)"
    exit 1
}

# --- Token ---
get_token() {
    cd "$BACKEND_DIR" && .venv/bin/python3 -c "
from investibook.db.queries import _db_init, _db_get_or_create_user, _db_create_user_session
_db_init()
user = _db_get_or_create_user('dev@localhost')
token = _db_create_user_session(user['id'])
print(token)
" 2>/dev/null
}

# --- Load token and base_url from CLI profile ---
load_profile() {
    local profile_name="$1"
    local config_file="$HOME/.investibook/config.json"
    if [ ! -f "$config_file" ]; then
        err "Config file not found: $config_file"
        err "Run: investibook init --profile $profile_name --api-host <url> --api-key <key>"
        exit 1
    fi
    local result
    result=$(python3 -c "
import json, sys
with open('$config_file') as f:
    cfg = json.load(f)
profiles = cfg.get('profiles', {})
p = profiles.get('$profile_name')
if not p:
    print('ERROR|Profile \"$profile_name\" not found. Available: ' + ', '.join(profiles.keys()), file=sys.stderr)
    sys.exit(1)
api = p.get('api', '')
token = p.get('token', '')
if not api or not token:
    print('ERROR|Profile \"$profile_name\" missing api or token', file=sys.stderr)
    sys.exit(1)
print(f'{api}|{token}')
" 2>&1) || { err "Failed to load profile: $result"; exit 1; }
    BASE_URL=$(echo "$result" | cut -d'|' -f1)
    TOKEN=$(echo "$result" | cut -d'|' -f2)
}

# --- Health check ---
check_backend() {
    local health
    health=$(curl -s --max-time 5 "$BASE_URL/health" 2>/dev/null || true)
    if ! echo "$health" | python3 -c "import sys,json; json.load(sys.stdin)" 2>/dev/null; then
        err "Backend not running at $BASE_URL"
        exit 1
    fi
}

# --- Clear messages ---
clear_messages() {
    local ws_id="$1" token="$2"
    curl -s -X DELETE "$BASE_URL/api/workspace/$ws_id/messages" \
        -H "Authorization: Bearer $token" | python3 -c "
import sys, json
data = json.load(sys.stdin)
print(data.get('status', 'unknown'))" 2>/dev/null
}

# --- Start diagnosis via auto-chat API (non-blocking, persists messages) ---
run_diagnosis() {
    local ws_id="$1" token="$2"

    log "Starting auto-chat diagnosis..."
    local resp
    resp=$(curl -s -X POST "$BASE_URL/api/workspaces/$ws_id/auto-chat/start" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d "{\"questions\": [\"$PROMPT\"]}")

    local status
    status=$(echo "$resp" | python3 -c "import sys,json; print(json.load(sys.stdin).get('status','unknown'))" 2>/dev/null)

    if [ "$status" != "running" ]; then
        err "Failed to start auto-chat: $resp"
        return 1
    fi
    ok "Auto-chat started"

    # Poll until complete
    local elapsed=0
    while [ "$elapsed" -lt "$MAX_TIME" ]; do
        local diag_status
        diag_status=$(curl -s "$BASE_URL/api/workspaces/$ws_id/diagnostics/status" \
            -H "Authorization: Bearer $token" 2>/dev/null \
            | python3 -c "
import sys,json
d=json.load(sys.stdin)
s=d.get('status','unknown')
idx=d.get('current_index',0)
total=d.get('total_questions',0)
reason=d.get('intervention_reason','')
print(f'{s}|{idx}|{total}|{reason}')
" 2>/dev/null || echo "error|0|0|")

        local ds di dt dr
        ds=$(echo "$diag_status" | cut -d'|' -f1)
        di=$(echo "$diag_status" | cut -d'|' -f2)
        dt=$(echo "$diag_status" | cut -d'|' -f3)
        dr=$(echo "$diag_status" | cut -d'|' -f4)

        case "$ds" in
            completed|resolved)
                ok "Diagnosis completed (question $di/$dt)"
                return 0
                ;;
            intervention_needed|needs_intervention)
                warn "Intervention needed: $dr"
                return 0
                ;;
            error)
                err "Diagnosis error"
                return 1
                ;;
        esac

        printf "\r  ... status=%s question=%s/%s elapsed=%ss" "$ds" "$di" "$dt" "$elapsed"
        sleep 5
        elapsed=$((elapsed + 5))
    done
    echo ""
    warn "Timed out after ${MAX_TIME}s (status=$ds)"
}

# --- Create workspace + source ---
create_workspace() {
    local issue_id="$1" token="$2"

    log "Creating workspace for $issue_id..." >&2
    local ws_id
    ws_id=$(curl -s -X POST "$BASE_URL/api/workspaces" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d "{\"name\": \"$issue_id postmortem\"}" \
        | python3 -c "import sys,json; print(json.load(sys.stdin)['id'])")
    ok "Workspace: $ws_id" >&2

    log "Adding Jira source..." >&2
    curl -s -X POST "$BASE_URL/api/workspace/$ws_id/sources" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d "{\"type\":\"jira_issue\",\"name\":\"$issue_id\",\"issueId\":\"$issue_id\",\"status\":\"loading\"}" \
        > /dev/null

    log "Preparing source data..." >&2
    local task_id
    task_id=$(curl -s -X POST "$BASE_URL/api/sources/prepare" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d "{\"workspace_id\":\"$ws_id\",\"issue_id\":\"$issue_id\"}" \
        | python3 -c "import sys,json; print(json.load(sys.stdin)['task_id'])")

    # Poll until ready
    while true; do
        local status
        status=$(curl -s "$BASE_URL/api/sources/prepare/$task_id" \
            -H "Authorization: Bearer $token" \
            | python3 -c "import sys,json; print(json.load(sys.stdin)['status'])")
        case "$status" in
            completed) ok "Source prepared" >&2; break ;;
            error)     err "Source preparation failed"; return 1 ;;
            *)         sleep 3 ;;
        esac
    done

    echo "$ws_id"
}

# --- Extract & evaluate ---
extract_and_evaluate() {
    local ws_id="$1" token="$2"
    local msg_file="/tmp/messages-${ws_id}.json"

    log "Fetching messages..."
    curl -s "$BASE_URL/api/workspace/$ws_id/messages" \
        -H "Authorization: Bearer $token" > "$msg_file"

    python3 -c "
import json, sys

def normalize_message_parts(raw_parts, content, role):
    \"\"\"Build a list of part dicts; tolerate legacy/malformed parts (str items, JSON string blobs).\"\"\"
    if raw_parts is not None:
        if isinstance(raw_parts, str):
            try:
                parsed = json.loads(raw_parts)
            except json.JSONDecodeError:
                parsed = [{'type': 'text', 'text': raw_parts}]
        else:
            parsed = raw_parts
        if isinstance(parsed, dict):
            parts = [parsed]
        elif isinstance(parsed, list):
            parts = []
            for item in parsed:
                if isinstance(item, dict):
                    parts.append(item)
                elif isinstance(item, str):
                    parts.append({'type': 'text', 'text': item})
        else:
            parts = []
    elif content and role == 'assistant':
        parts = [{'type': 'text', 'text': content}]
    elif content and role == 'user':
        parts = [{'type': 'text', 'text': content}]
    else:
        parts = []
    return parts

with open('$msg_file') as f:
    data = json.load(f)
if isinstance(data, list):
    msgs = data
elif isinstance(data, dict):
    if data.get('detail') is not None and 'messages' not in data:
        print('API error: ' + str(data.get('detail')), file=sys.stderr)
        sys.exit(1)
    rawm = data.get('messages')
    msgs = rawm if isinstance(rawm, list) else []
else:
    msgs = []

all_tools = []
all_texts = []
human_interventions = 0

for m in msgs:
    if not isinstance(m, dict):
        continue
    role = m.get('role', '')
    raw_parts = m.get('parts')
    content = m.get('content', '')
    parts = normalize_message_parts(raw_parts, content, role)

    if role == 'user' and parts:
        human_interventions += 1
    if role != 'assistant':
        continue
    for p in parts:
        if not isinstance(p, dict):
            continue
        if p.get('type') == 'tool':
            state = p.get('state', {})
            inp = state.get('input', {})
            all_tools.append({
                'title': state.get('title', ''),
                'description': inp.get('description', '') if isinstance(inp, dict) else '',
                'status': state.get('status', ''),
            })
        elif p.get('type') == 'text':
            all_texts.append(p.get('text', ''))

total = len(all_tools)
completed = sum(1 for t in all_tools if t['status'] == 'completed')
stuck = sum(1 for t in all_tools if t['status'] == 'running')

# Cumulative text: scan ALL assistant messages for data coverage
final = all_texts[-1] if all_texts else ''
all_text_lower = ' '.join(all_texts).lower()

# Data source coverage: check tool calls first, fallback to cumulative text content
has_topsql_tidb = any('tidb' in (t['description']+t['title']).lower() and 'topsql' in (t['description']+t['title']).lower() for t in all_tools) or ('topsql' in all_text_lower and 'tidb' in all_text_lower)
has_topsql_tikv = any('tikv' in (t['description']+t['title']).lower() and 'topsql' in (t['description']+t['title']).lower() for t in all_tools) or ('topsql' in all_text_lower and 'tikv' in all_text_lower)
has_slowquery = any('slow' in (t['description']+t['title']).lower() for t in all_tools) or ('slow' in all_text_lower and 'query' in all_text_lower)
has_plan = any('plan' in (t['description']+t['title']).lower() or 'decoded_plan' in (t['description']+t['title']).lower() for t in all_tools) or 'decoded_plan' in all_text_lower or 'execution plan' in all_text_lower
has_logs = any('log' in (t['description']+t['title']).lower() or 'loki' in (t['description']+t['title']).lower() for t in all_tools) or 'loki' in all_text_lower or 'error log' in all_text_lower
has_events = any('event' in (t['description']+t['title']).lower() or 'activity' in (t['description']+t['title']).lower() for t in all_tools) or 'cluster event' in all_text_lower or 'activity' in all_text_lower
has_metrics = any('cpu' in (t['description']+t['title']).lower() or 'latency' in (t['description']+t['title']).lower() or 'ebs' in (t['description']+t['title']).lower() or 'iops' in (t['description']+t['title']).lower() for t in all_tools) or ('cpu' in all_text_lower and ('latency' in all_text_lower or 'metric' in all_text_lower))

print(f'''
{'='*60}
  Workspace: $ws_id
  Messages: {len(msgs)}
  Human interventions: {human_interventions}
  Tool calls: {total} (completed={completed}, stuck={stuck})
{'='*60}

Data Source Coverage:
  Metrics (EBS/CPU/latency):  {'YES' if has_metrics else 'NO'}
  TopSQL TiDB:                {'YES' if has_topsql_tidb else 'NO'}
  TopSQL TiKV:                {'YES' if has_topsql_tikv else 'NO'}
  SlowQuery:                  {'YES' if has_slowquery else 'NO'}
  Execution Plan:             {'YES' if has_plan else 'NO'}
  Loki Logs:                  {'YES' if has_logs else 'NO'}
  Events/Activity:            {'YES' if has_events else 'NO'}

Tool call sequence:''')

for i, t in enumerate(all_tools):
    label = t['description'] or t['title']
    print(f'  {i+1:2d}. [{t[\"status\"]}] {label[:100]}')

print(f'''
{'='*60}
  Final diagnosis (last 2000 chars):
{'='*60}
{final[-2000:]}
''')
"
}

# --- Send follow-up message via auto-chat (add question + continue) ---
send_followup() {
    local ws_id="$1" token="$2" message="$3"

    log "Sending follow-up via auto-chat..."

    # Start a new auto-chat round with the follow-up as the question
    local resp
    resp=$(curl -s -X POST "$BASE_URL/api/workspaces/$ws_id/auto-chat/start" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d "$(python3 -c "import json,sys; print(json.dumps({'questions': [sys.argv[1]]}))" "$message")")

    local status
    status=$(echo "$resp" | python3 -c "import sys,json; print(json.load(sys.stdin).get('status','unknown'))" 2>/dev/null)

    if [ "$status" != "running" ]; then
        err "Failed to start follow-up: $resp"
        return 1
    fi
    ok "Follow-up started"

    # Poll until complete
    local elapsed=0
    while [ "$elapsed" -lt "$MAX_TIME" ]; do
        local diag_status
        diag_status=$(curl -s "$BASE_URL/api/workspaces/$ws_id/diagnostics/status" \
            -H "Authorization: Bearer $token" 2>/dev/null \
            | python3 -c "
import sys,json
d=json.load(sys.stdin)
s=d.get('status','unknown')
print(s)
" 2>/dev/null || echo "error")

        case "$diag_status" in
            completed|resolved)
                ok "Follow-up completed"
                return 0
                ;;
            intervention_needed|needs_intervention)
                warn "Agent needs intervention"
                return 0
                ;;
            error)
                err "Follow-up error"
                return 1
                ;;
        esac

        printf "\r  ... follow-up status=%s elapsed=%ss" "$diag_status" "$elapsed"
        sleep 5
        elapsed=$((elapsed + 5))
    done
    echo ""
    warn "Follow-up timed out after ${MAX_TIME}s"
}

# --- Structured checklist (JSON output for agent consumption) ---
check_diagnosis() {
    local ws_id="$1" token="$2"
    local msg_file="/tmp/messages-${ws_id}.json"

    curl -s "$BASE_URL/api/workspace/$ws_id/messages" \
        -H "Authorization: Bearer $token" > "$msg_file"

    python3 -c "
import json, sys

def normalize_message_parts(raw_parts, content, role):
    if raw_parts is not None:
        if isinstance(raw_parts, str):
            try:
                parsed = json.loads(raw_parts)
            except json.JSONDecodeError:
                parsed = [{'type': 'text', 'text': raw_parts}]
        else:
            parsed = raw_parts
        if isinstance(parsed, dict):
            parts = [parsed]
        elif isinstance(parsed, list):
            parts = []
            for item in parsed:
                if isinstance(item, dict):
                    parts.append(item)
                elif isinstance(item, str):
                    parts.append({'type': 'text', 'text': item})
        else:
            parts = []
    elif content and role == 'assistant':
        parts = [{'type': 'text', 'text': content}]
    elif content and role == 'user':
        parts = [{'type': 'text', 'text': content}]
    else:
        parts = []
    return parts

with open('$msg_file') as f:
    data = json.load(f)
if isinstance(data, list):
    msgs = data
elif isinstance(data, dict):
    if data.get('detail') is not None and 'messages' not in data:
        print('API error: ' + str(data.get('detail')), file=sys.stderr)
        sys.exit(1)
    rawm = data.get('messages')
    msgs = rawm if isinstance(rawm, list) else []
else:
    msgs = []

all_tools = []
all_texts = []
human_msgs = []

for m in msgs:
    if not isinstance(m, dict):
        continue
    role = m.get('role', '')
    raw_parts = m.get('parts')
    content = m.get('content', '')
    parts = normalize_message_parts(raw_parts, content, role)

    if role == 'user':
        for p in parts:
            if not isinstance(p, dict):
                continue
            if p.get('type') == 'text':
                human_msgs.append(p.get('text', ''))
    if role != 'assistant':
        continue
    for p in parts:
        if not isinstance(p, dict):
            continue
        if p.get('type') == 'tool':
            state = p.get('state', {})
            inp = state.get('input', {})
            all_tools.append({
                'title': state.get('title', ''),
                'description': inp.get('description', '') if isinstance(inp, dict) else '',
                'status': state.get('status', ''),
            })
        elif p.get('type') == 'text':
            all_texts.append(p.get('text', ''))

def match(keywords, tools):
    \"\"\"Check if any tool description/title contains ALL keywords.\"\"\"
    for t in tools:
        combined = (t['description'] + ' ' + t['title']).lower()
        if all(k in combined for k in keywords):
            return True
    return False

final_text = all_texts[-1] if all_texts else ''
final_lower = final_text.lower()

# Cumulative text: scan ALL assistant messages for data coverage (not just last)
all_text = ' '.join(all_texts).lower()

# Data source checks (tool calls first, fallback to cumulative text content)
checks = {
    'metrics':        match(['cpu'], all_tools) or match(['latency'], all_tools) or match(['ebs'], all_tools) or match(['iops'], all_tools) or ('cpu' in all_text and ('latency' in all_text or 'metric' in all_text)),
    'topsql_tidb':    match(['topsql', 'tidb'], all_tools) or ('topsql' in all_text and 'tidb' in all_text),
    'topsql_tikv':    match(['topsql', 'tikv'], all_tools) or ('topsql' in all_text and 'tikv' in all_text),
    'slowquery':      match(['slow'], all_tools) or ('slow' in all_text and 'query' in all_text),
    'decoded_plan':   match(['decoded_plan'], all_tools) or match(['plan'], all_tools) or 'decoded_plan' in all_text or 'execution plan' in all_text,
    'logs':           match(['log'], all_tools) or match(['loki'], all_tools) or 'loki' in all_text,
    'events':         match(['event'], all_tools) or match(['activity'], all_tools) or 'cluster event' in all_text or 'activity' in all_text,
}

# Quality checks on final diagnosis text (support both EN and CN phrasing)
quality = {
    'has_root_cause':    any(kw in final_lower for kw in ['root cause', 'primary cause', 'root-cause', 'causality', 'causal chain']) or any(kw in final_text for kw in ['根因', '根本原因', '主要原因', '瓶颈']),
    'has_evidence':      any(kw in final_lower for kw in ['evidence', 'confirmed by', 'observed', 'metrics', 'topsql', 'slowquery']) or any(kw in final_text for kw in ['证据', '确认', '观察到', '数据显示']),
    'has_recommendation': any(kw in final_lower for kw in ['recommend', 'suggestion', 'remediation', 'binding', 'index']) or any(kw in final_text for kw in ['建议', '优化', '修复', '索引']),
    'has_specific_table': any(kw in final_lower for kw in ['table', 'idx_', 'index(']) or any(kw in final_text for kw in ['表', '索引']),
    'has_complete_sql':  ('create' in final_lower and ('binding' in final_lower or 'index' in final_lower) and '...' not in final_text) or ('CREATE' in final_text and 'INDEX' in final_text and '...' not in final_text),
    'no_stuck_tools':    all(t['status'] != 'running' for t in all_tools),
}

total_tools = len(all_tools)
completed_tools = sum(1 for t in all_tools if t['status'] == 'completed')

# Identify gaps for follow-up suggestions
gaps = []
if not checks['topsql_tidb']:
    gaps.append({'check': 'topsql_tidb', 'followup': 'Please query TopSQL with component=tidb to identify top CPU-consuming SQL digests.'})
if not checks['topsql_tikv']:
    gaps.append({'check': 'topsql_tikv', 'followup': 'Please query TopSQL with component=tikv to check TiKV-side CPU attribution and identify hotspot instances.'})
if not checks['decoded_plan']:
    gaps.append({'check': 'decoded_plan', 'followup': 'Please fetch decoded_plan from slow_query_logs for the top slow query digest. Use the decoded_plan field, not the plan field.'})
if not checks['slowquery']:
    gaps.append({'check': 'slowquery', 'followup': 'Please query slow_query_logs to find slow query digests and their execution statistics.'})
if not quality['has_complete_sql'] and quality['has_recommendation']:
    gaps.append({'check': 'complete_sql', 'followup': 'Your remediation includes SQL statements with ellipsis (...). Please provide complete, executable SQL for CREATE INDEX and CREATE BINDING.'})
if not quality['has_specific_table']:
    gaps.append({'check': 'specific_table', 'followup': 'Please specify the exact table and column causing the performance issue in your root cause analysis.'})
if not quality['no_stuck_tools']:
    gaps.append({'check': 'stuck_tools', 'followup': 'Some tool calls appear stuck. Please continue your analysis or provide a conclusion based on data collected so far.'})

result = {
    'workspace_id': '$ws_id',
    'message_count': len(msgs),
    'human_interventions': len(human_msgs),
    'tool_calls': {'total': total_tools, 'completed': completed_tools, 'stuck': total_tools - completed_tools},
    'data_coverage': checks,
    'quality': quality,
    'gaps': gaps,
    'all_passed': len(gaps) == 0,
    'final_text_length': len(final_text),
}

print(json.dumps(result, indent=2))
"
}

# --- Wait for stream to finish ---
wait_for_completion() {
    local ws_id="$1" token="$2"
    log "Waiting for diagnosis to complete..."
    local prev_count=0
    local stable=0
    while true; do
        local count
        count=$(curl -s "$BASE_URL/api/workspace/$ws_id/messages" \
            -H "Authorization: Bearer $token" | python3 -c "
import sys, json
data = json.load(sys.stdin)
if isinstance(data, list):
    msgs = data
elif isinstance(data, dict):
    if data.get('detail') is not None and 'messages' not in data:
        msgs = []
    else:
        rawm = data.get('messages')
        msgs = rawm if isinstance(rawm, list) else []
else:
    msgs = []
total = 0
for m in msgs:
    if not isinstance(m, dict) or m.get('role') != 'assistant':
        continue
    raw = m.get('parts')
    if raw is None:
        parts = []
    elif isinstance(raw, str):
        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError:
            parsed = []
        parts = [parsed] if isinstance(parsed, dict) else (parsed if isinstance(parsed, list) else [])
    elif isinstance(raw, dict):
        parts = [raw]
    elif isinstance(raw, list):
        parts = [x for x in raw if isinstance(x, dict)]
    else:
        parts = []
    for p in parts:
        if p.get('type') == 'tool':
            total += 1
print(total)" 2>/dev/null || echo "0")

        local sessions
        sessions=$(curl -s "$BASE_URL/health" 2>/dev/null | python3 -c "import sys,json; print(json.load(sys.stdin).get('active_sessions',0))" 2>/dev/null || echo "1")

        if [ "$sessions" = "0" ] || ([ "$count" = "$prev_count" ] && [ "$count" != "0" ] && [ "$sessions" = "0" ]); then
            ok "Diagnosis complete (${count} tool calls)"
            break
        fi

        if [ "$count" = "$prev_count" ] && [ "$count" != "0" ]; then
            stable=$((stable + 1))
            if [ "$stable" -ge 6 ]; then
                ok "Diagnosis appears complete (${count} tool calls, stable for 30s)"
                break
            fi
        else
            stable=0
        fi

        prev_count=$count
        printf "\r  ... %s tool calls, %s active sessions" "$count" "$sessions"
        sleep 5
    done
    echo ""
}

# ======================
# Main
# ======================

[ $# -lt 1 ] && usage

PROFILE=""
MODE="rerun"
WS_ID=""

# Parse all arguments
while [ $# -gt 0 ]; do
    case "$1" in
        --profile)
            PROFILE="${2:-}"
            [ -z "$PROFILE" ] && err "--profile requires <name>" && exit 1
            shift 2
            ;;
        --eval-only) MODE="eval"; shift ;;
        --check) MODE="check"; shift ;;
        --create)
            MODE="create"
            ISSUE_ID="${2:-}"
            [ -z "$ISSUE_ID" ] && err "--create requires <issue_id>" && exit 1
            shift 2
            ;;
        --followup)
            MODE="followup"
            FOLLOWUP_MSG="${2:-}"
            [ -z "$FOLLOWUP_MSG" ] && err "--followup requires <message>" && exit 1
            shift 2
            ;;
        -*)
            err "Unknown option: $1"; usage ;;
        *)
            [ -z "$WS_ID" ] && WS_ID="$1" || { err "Unexpected argument: $1"; usage; }
            shift
            ;;
    esac
done

[ -z "$WS_ID" ] && err "Missing workspace_id" && usage

# Authenticate: profile or local
if [ -n "$PROFILE" ]; then
    load_profile "$PROFILE"
    ok "Profile '$PROFILE' loaded ($BASE_URL)"
else
    check_backend
    TOKEN=$(get_token)
    ok "Token acquired"
fi

case "$MODE" in
    eval)
        extract_and_evaluate "$WS_ID" "$TOKEN"
        ;;
    check)
        check_diagnosis "$WS_ID" "$TOKEN"
        ;;
    followup)
        send_followup "$WS_ID" "$TOKEN" "$FOLLOWUP_MSG"
        ok "Follow-up stream finished"

        check_diagnosis "$WS_ID" "$TOKEN"
        ;;
    rerun)
        log "Clearing messages for $WS_ID..."
        clear_messages "$WS_ID" "$TOKEN"
        ok "Messages cleared"

        run_diagnosis "$WS_ID" "$TOKEN"
        ok "Stream finished"

        extract_and_evaluate "$WS_ID" "$TOKEN"
        ;;
    create)
        WS_ID=$(create_workspace "$ISSUE_ID" "$TOKEN")
        run_diagnosis "$WS_ID" "$TOKEN"
        ok "Stream finished"

        extract_and_evaluate "$WS_ID" "$TOKEN"
        ;;
esac