#!/bin/bash
###############################################################################
# OpenClaw Agent Onboarding Script
#
# Fully automated: creates workspace, registers in agents.yaml +
# project_registry.yaml, installs systemd service, reloads monitor.
#
# Usage:
#   ./onboard.sh <agent_id> <agent_name> <project_id> [qdrant_host]
#
# Examples:
#   ./onboard.sh crypto "CryptoBot" crypto              # local agent
#   ./onboard.sh remote1 "RemoteBot" advert 100.115.94.1 # remote agent
###############################################################################

set -e

WORKSPACE="/root/.openclaw/workspace"
TEMPLATE_DIR="$WORKSPACE/templates"
AGENTS_YAML="$WORKSPACE/agents.yaml"
REGISTRY="$WORKSPACE/skills/mem0-integration/project_registry.yaml"

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 3 ]; then
    echo "Usage: $0 <agent_id> <agent_name> <project_id> [qdrant_host]"
    echo ""
    echo "  agent_id     Unique identifier (lowercase, no spaces)"
    echo "  agent_name   Display name for the agent"
    echo "  project_id   Project to register in (must exist in project_registry.yaml)"
    echo "  qdrant_host  Optional: Qdrant host for remote agents (default: localhost)"
    exit 1
fi

AGENT_ID="$1"
AGENT_NAME="$2"
PROJECT_ID="$3"
QDRANT_HOST="${4:-localhost}"
USER_ID="wang_yuanzhang"
DATE=$(date +%Y-%m-%d)

AGENT_WORKSPACE="$WORKSPACE/agents/${AGENT_ID}-workspace"
AGENT_CONFIG_DIR="/root/.openclaw-${AGENT_ID}"
SYSTEMD_UNIT="openclaw-gateway-${AGENT_ID}.service"

log_info "Onboarding agent: ${AGENT_NAME} (${AGENT_ID})"
log_info "Project: ${PROJECT_ID}, Qdrant: ${QDRANT_HOST}"

# Pre-check: ensure agent doesn't already exist
if python3 "$WORKSPACE/scripts/parse_agents.py" info "$AGENT_ID" >/dev/null 2>&1; then
    log_error "Agent '${AGENT_ID}' already exists in agents.yaml"
    exit 1
fi

# 1. Create workspace from templates
if [ -d "$AGENT_WORKSPACE" ]; then
    log_error "Workspace already exists: $AGENT_WORKSPACE"
    exit 1
fi

log_info "Creating workspace at $AGENT_WORKSPACE..."
mkdir -p "$AGENT_WORKSPACE/skills/mem0-integration"
mkdir -p "$AGENT_WORKSPACE/memory"

for tmpl in IDENTITY.md.template SOUL.md.template; do
    base="${tmpl%.template}"
    sed -e "s/{{AGENT_ID}}/${AGENT_ID}/g" \
        -e "s/{{AGENT_NAME}}/${AGENT_NAME}/g" \
        -e "s/{{AGENT_ROLE}}/TODO: define role/g" \
        -e "s/{{PROJECT_ID}}/${PROJECT_ID}/g" \
        -e "s/{{DATE}}/${DATE}/g" \
        "$TEMPLATE_DIR/agent-workspace/$tmpl" > "$AGENT_WORKSPACE/$base"
done

ln -sf "$WORKSPACE/USER.md" "$AGENT_WORKSPACE/USER.md"
ln -sf "$WORKSPACE/AGENTS.md" "$AGENT_WORKSPACE/AGENTS.md"

sed -e "s/{{AGENT_ID}}/${AGENT_ID}/g" \
    -e "s/{{AGENT_NAME}}/${AGENT_NAME}/g" \
    -e "s/{{QDRANT_HOST}}/${QDRANT_HOST}/g" \
    -e "s/{{USER_ID}}/${USER_ID}/g" \
    "$TEMPLATE_DIR/agent-workspace/skills/mem0-integration/config.yaml.template" \
    > "$AGENT_WORKSPACE/skills/mem0-integration/config.yaml"

log_success "Workspace created"

# 2. Register in agents.yaml (uses sys.argv to avoid shell injection)
if [ "$QDRANT_HOST" = "localhost" ]; then
    AGENT_TYPE="local-systemd"
    python3 - "$AGENTS_YAML" "$AGENT_ID" "$AGENT_NAME" "$AGENT_TYPE" \
              "$AGENT_CONFIG_DIR" "$AGENT_WORKSPACE" "$SYSTEMD_UNIT" "$PROJECT_ID" <<'PYEOF'
import sys, yaml
yaml_path, aid, aname, atype, profile, ws, unit, proj = sys.argv[1:9]
with open(yaml_path, 'r', encoding='utf-8') as f:
    data = yaml.safe_load(f)
data['agents'][aid] = {
    'name': aname, 'type': atype,
    'profile_dir': profile, 'workspace': ws,
    'service': {'unit': unit},
    'env_file': f'{aid}-gateway.env',
    'projects': [proj],
}
with open(yaml_path, 'w', encoding='utf-8') as f:
    yaml.dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
PYEOF
else
    AGENT_TYPE="remote-http"
    python3 - "$AGENTS_YAML" "$AGENT_ID" "$AGENT_NAME" "$AGENT_TYPE" \
              "$AGENT_WORKSPACE" "$QDRANT_HOST" "$PROJECT_ID" <<'PYEOF'
import sys, yaml
yaml_path, aid, aname, atype, ws, qhost, proj = sys.argv[1:8]
with open(yaml_path, 'r', encoding='utf-8') as f:
    data = yaml.safe_load(f)
data['agents'][aid] = {
    'name': aname, 'type': atype,
    'workspace': ws,
    'service': {'health_url': f'http://{qhost}:18789/health', 'timeout': 5000},
    'projects': [proj],
    'qdrant_host': qhost,
}
with open(yaml_path, 'w', encoding='utf-8') as f:
    yaml.dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
PYEOF
fi
log_success "Registered in agents.yaml (type: ${AGENT_TYPE})"

# 3. Register in project_registry.yaml
if grep -q "\"${AGENT_ID}\"" "$REGISTRY" 2>/dev/null; then
    log_warning "Agent ${AGENT_ID} already in project registry"
else
    if grep -q "^  ${PROJECT_ID}:" "$REGISTRY"; then
        sed -i "/^  ${PROJECT_ID}:/,/owner:/ {
            /members:/a\\      - \"${AGENT_ID}\"
        }" "$REGISTRY"
        log_success "Registered ${AGENT_ID} in project ${PROJECT_ID}"
    else
        log_warning "Project ${PROJECT_ID} not found in registry. Add manually."
    fi
fi

# 4. Generate systemd service + env files (local agents only)
if [ "$AGENT_TYPE" = "local-systemd" ]; then
    SERVICE_FILE="$WORKSPACE/systemd/${SYSTEMD_UNIT}"
    sed -e "s/{{AGENT_ID}}/${AGENT_ID}/g" \
        -e "s/{{AGENT_NAME}}/${AGENT_NAME}/g" \
        "$TEMPLATE_DIR/systemd/agent-gateway.service.template" > "$SERVICE_FILE"
    log_success "Service file: $SERVICE_FILE"
    
    ENV_FILE="$WORKSPACE/systemd/${AGENT_ID}-gateway.env"
    sed -e "s/{{AGENT_ID}}/${AGENT_ID}/g" \
        -e "s/{{AGENT_NAME}}/${AGENT_NAME}/g" \
        -e "s/{{QDRANT_HOST}}/${QDRANT_HOST}/g" \
        "$TEMPLATE_DIR/systemd/agent-gateway.env.template" > "$ENV_FILE"
    chmod 600 "$ENV_FILE"
    log_success "Env file: $ENV_FILE"
    
    # 5. Install and start the service
    setup_user_env
    mkdir -p ~/.config/systemd/user/
    cp "$SERVICE_FILE" "$HOME/.config/systemd/user/${SYSTEMD_UNIT}"
    systemctl --user daemon-reload
    systemctl --user enable "${SYSTEMD_UNIT}"
    log_success "Service installed and enabled"
    
    # 6. Create OpenClaw profile directory
    mkdir -p "$AGENT_CONFIG_DIR"
    log_info "Profile directory created: $AGENT_CONFIG_DIR"
    log_warning "You must create $AGENT_CONFIG_DIR/openclaw.json before starting"
fi

# 7. Reload the monitor to pick up the new agent
systemctl restart openclaw-agent-monitor 2>/dev/null && log_success "Monitor reloaded" || log_warning "Monitor reload failed (may not be running)"

echo ""
log_success "Onboarding complete for ${AGENT_NAME} (${AGENT_ID})"
echo ""

if [ "$AGENT_TYPE" = "local-systemd" ]; then
    log_info "Remaining steps:"
    echo "  1. Edit agent identity: vim $AGENT_WORKSPACE/IDENTITY.md"
    echo "  2. Create openclaw.json: vim $AGENT_CONFIG_DIR/openclaw.json"
    echo "     (copy from /root/.openclaw/openclaw.json and modify)"
    echo "  3. Start: systemctl --user start ${SYSTEMD_UNIT}"
    echo ""
elif [ "$AGENT_TYPE" = "remote-http" ]; then
    log_info "Remaining steps:"
    echo "  1. Deploy the agent on the remote server at ${QDRANT_HOST}"
    echo "  2. Ensure Tailscale connectivity to ${QDRANT_HOST}:6333 (Qdrant)"
    echo "  3. Configure the remote agent to use Qdrant collection: mem0_v4_shared"
    echo ""
fi