#!/bin/bash
###############################################################################
# OpenClaw Agent Offboarding / Removal Script
#
# Cleanly removes an agent: stops service, removes from agents.yaml,
# project_registry.yaml, optionally deletes workspace, profile, and Qdrant data.
#
# Usage:
#   ./offboard.sh <agent_id> [--keep-data]
#
# Options:
#   --keep-data   Keep workspace and profile directories (only unregister)
#
# Examples:
#   ./offboard.sh crypto              # full removal
#   ./offboard.sh crypto --keep-data  # keep files, just unregister
###############################################################################

set -e

WORKSPACE="/root/.openclaw/workspace"
AGENTS_YAML="$WORKSPACE/agents.yaml"
REGISTRY="$WORKSPACE/skills/mem0-integration/project_registry.yaml"
PARSE_AGENTS="python3 $WORKSPACE/scripts/parse_agents.py"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info()    { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warning() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error()   { echo -e "${RED}[ERROR]${NC} $1"; }

setup_user_env() {
    export XDG_RUNTIME_DIR=/run/user/$(id -u)
    export DBUS_SESSION_BUS_ADDRESS="unix:path=/run/user/$(id -u)/bus"
}

if [ $# -lt 1 ]; then
    echo "Usage: $0 <agent_id> [--keep-data]"
    exit 1
fi

AGENT_ID="$1"
KEEP_DATA=false
[ "$2" = "--keep-data" ] && KEEP_DATA=true

if [ "$AGENT_ID" = "main" ]; then
    log_error "Cannot remove the main (hub) agent"
    exit 1
fi

# Validate agent exists
if ! $PARSE_AGENTS info "$AGENT_ID" >/dev/null 2>&1; then
    log_error "Agent '${AGENT_ID}' not found in agents.yaml"
    exit 1
fi

# Load agent info
eval $($PARSE_AGENTS info "$AGENT_ID")
log_info "Offboarding: ${AGENT_NAME} (${AGENT_ID}), type: ${AGENT_TYPE}"

echo ""
log_warning "This will remove agent '${AGENT_ID}' from the system."
if [ "$KEEP_DATA" = "true" ]; then
    log_info "Mode: keep data (unregister only)"
else
    log_warning "Mode: FULL removal (workspace, profile, Qdrant data will be DELETED)"
fi
read -p "Continue? (y/N): " confirm
if [[ ! $confirm =~ ^[Yy]$ ]]; then
    log_info "Cancelled."
    exit 0
fi

# 1. Stop and disable the service
setup_user_env
if [ "$AGENT_TYPE" = "local-systemd" ] && [ -n "$SYSTEMD_UNIT" ]; then
    log_info "Stopping service: $SYSTEMD_UNIT"
    systemctl --user stop "$SYSTEMD_UNIT" 2>/dev/null || true
    systemctl --user disable "$SYSTEMD_UNIT" 2>/dev/null || true
    rm -f "$HOME/.config/systemd/user/$SYSTEMD_UNIT"
    systemctl --user daemon-reload 2>/dev/null
    log_success "Service stopped and removed"
fi

# 2. Remove from agents.yaml
python3 - "$AGENTS_YAML" "$AGENT_ID" <<'PYEOF'
import sys, yaml
yaml_path, aid = sys.argv[1:3]
with open(yaml_path, 'r', encoding='utf-8') as f:
    data = yaml.safe_load(f)
if aid in data.get('agents', {}):
    del data['agents'][aid]
    with open(yaml_path, 'w', encoding='utf-8') as f:
        yaml.dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
    print('removed from agents.yaml')
else:
    print('not found in agents.yaml')
PYEOF
log_success "Removed from agents.yaml"

# 3. Remove from project_registry.yaml
if grep -q "\"${AGENT_ID}\"" "$REGISTRY" 2>/dev/null; then
    sed -i "/\"${AGENT_ID}\"/d" "$REGISTRY"
    log_success "Removed from project_registry.yaml"
else
    log_info "Not found in project_registry.yaml (skip)"
fi

# 4. Delete files (unless --keep-data)
if [ "$KEEP_DATA" = "false" ]; then
    AGENT_WORKSPACE="$WORKSPACE/agents/${AGENT_ID}-workspace"
    AGENT_CONFIG_DIR="/root/.openclaw-${AGENT_ID}"
    SVC_FILE="$WORKSPACE/systemd/openclaw-gateway-${AGENT_ID}.service"
    ENV_FILE="$WORKSPACE/systemd/${AGENT_ID}-gateway.env"
    LOGS_DIR="$WORKSPACE/logs/agents/${AGENT_ID}"
    RUNTIME_DIR="/root/.openclaw/agents/${AGENT_ID}"

    [ -d "$AGENT_WORKSPACE" ] && rm -rf "$AGENT_WORKSPACE" && log_success "Deleted workspace: $AGENT_WORKSPACE"
    [ -d "$AGENT_CONFIG_DIR" ] && rm -rf "$AGENT_CONFIG_DIR" && log_success "Deleted profile: $AGENT_CONFIG_DIR"
    [ -f "$SVC_FILE" ] && rm -f "$SVC_FILE" && log_success "Deleted service file"
    [ -f "$ENV_FILE" ] && rm -f "$ENV_FILE" && log_success "Deleted env file"
    [ -d "$LOGS_DIR" ] && rm -rf "$LOGS_DIR" && log_success "Deleted logs"
    [ -d "$RUNTIME_DIR" ] && rm -rf "$RUNTIME_DIR" && log_success "Deleted runtime data"

    # 5. Clean Qdrant data
    log_info "Cleaning Qdrant memories for agent_id=${AGENT_ID}..."
    python3 -c "
try:
    from qdrant_client import QdrantClient
    from qdrant_client.models import Filter, FieldCondition, MatchValue, FilterSelector
    client = QdrantClient(host='localhost', port=6333)
    result = client.delete(
        collection_name='mem0_v4_shared',
        points_selector=FilterSelector(filter=Filter(must=[
            FieldCondition(key='agent_id', match=MatchValue(value='${AGENT_ID}'))
        ]))
    )
    print(f'Deleted Qdrant memories: {result.status}')
except Exception as e:
    print(f'Qdrant cleanup skipped: {e}')
" 2>/dev/null
    log_success "Qdrant memories cleaned"
fi

# 6. Reload monitor
systemctl restart openclaw-agent-monitor 2>/dev/null && log_success "Monitor reloaded" || log_warning "Monitor reload failed"

echo ""
log_success "Agent '${AGENT_ID}' has been fully removed."
echo ""
log_info "Summary:"
echo "  - Service:              removed"
echo "  - agents.yaml:          removed"
echo "  - project_registry:     removed"
if [ "$KEEP_DATA" = "false" ]; then
    echo "  - Workspace + profile:  deleted"
    echo "  - Qdrant memories:      deleted"
else
    echo "  - Workspace + profile:  kept (--keep-data)"
    echo "  - Qdrant memories:      kept (--keep-data)"
fi