#!/bin/bash
#
# OpenBrain MCP Deployment Script
# Deploys the OpenBrain MCP server to the VPS
#
# Usage: ./deploy.sh [options]
# Options:
#   --build-local    Build on local machine (requires cross-compilation)
#   --build-remote   Build on VPS (default)
#   --skip-model     Skip model download
#   --restart-only   Only restart the service
#

set -euo pipefail

# Configuration
VPS_HOST="${VPS_HOST:-}"
VPS_USER="${VPS_USER:-root}"
DEPLOY_DIR="/opt/openbrain-mcp"
SERVICE_NAME="openbrain-mcp"
SSH_KEY="${SSH_KEY:-/tmp/id_ed25519}"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

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

# Parse arguments
BUILD_REMOTE=true
SKIP_MODEL=false
RESTART_ONLY=false

for arg in "$@"; do
    case $arg in
        --build-local)  BUILD_REMOTE=false ;;
        --build-remote) BUILD_REMOTE=true ;;
        --skip-model)   SKIP_MODEL=true ;;
        --restart-only) RESTART_ONLY=true ;;
        *) log_error "Unknown argument: $arg"; exit 1 ;;
    esac
done

# Get script directory (where .gitea folder is)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

log_info "Project root: $PROJECT_ROOT"
if [ -z "$VPS_HOST" ]; then
    log_error "VPS_HOST is required. Export VPS_HOST before running deploy.sh"
    exit 1
fi
log_info "Deploying to: $VPS_USER@$VPS_HOST:$DEPLOY_DIR"

# SSH command helper
ssh_cmd() {
    ssh -i "$SSH_KEY" -o StrictHostKeyChecking=no "$VPS_USER@$VPS_HOST" "$@"
}

scp_cmd() {
    scp -i "$SSH_KEY" -o StrictHostKeyChecking=no "$@"
}

# Restart only mode
if [ "$RESTART_ONLY" = true ]; then
    log_info "Restarting service only..."
    ssh_cmd "systemctl restart $SERVICE_NAME"
    ssh_cmd "systemctl status $SERVICE_NAME --no-pager"
    exit 0
fi

# Step 1: Create deployment directory on VPS
log_info "Creating deployment directory on VPS..."
ssh_cmd "mkdir -p $DEPLOY_DIR/{src,models,logs,lib,.gitea}"

# Step 2: Sync source code to VPS
log_info "Syncing source code to VPS..."
rsync -avz --delete \
    -e "ssh -i $SSH_KEY -o StrictHostKeyChecking=no" \
    --exclude 'target/' \
    --exclude '.git/' \
    --exclude '.a0proj/' \
    --exclude 'models/' \
    --exclude '*.md' \
    "$PROJECT_ROOT/" \
    "$VPS_USER@$VPS_HOST:$DEPLOY_DIR/"

# Step 3: Copy .env if it doesn't exist on VPS
if ! ssh_cmd "test -f $DEPLOY_DIR/.env"; then
    log_warn ".env not found on VPS. Copying .env.example..."
    ssh_cmd "cp $DEPLOY_DIR/.env.example $DEPLOY_DIR/.env"
    log_warn "Please edit $DEPLOY_DIR/.env on VPS with actual credentials!"
fi

# Step 4: Download model if needed
if [ "$SKIP_MODEL" = false ]; then
    log_info "Checking/downloading embedding model..."
    ssh_cmd "bash $DEPLOY_DIR/.gitea/download-model.sh"
fi

# Step 5: Build on VPS
if [ "$BUILD_REMOTE" = true ]; then
    log_info "Building on VPS (this may take a while on first run)..."
    ssh_cmd "cd $DEPLOY_DIR && \
        source ~/.cargo/env 2>/dev/null || true && \
        cargo build --release 2>&1"
else
    log_error "Local cross-compilation not yet implemented"
    exit 1
fi

# Step 5b: Install the built binary where systemd expects it
log_info "Installing built binary..."
ssh_cmd "cp $DEPLOY_DIR/target/release/openbrain-mcp $DEPLOY_DIR/openbrain-mcp && chmod +x $DEPLOY_DIR/openbrain-mcp"

# Step 5c: Bootstrap runtime prerequisites
log_info "Bootstrapping runtime prerequisites..."
ssh -i "$SSH_KEY" -o StrictHostKeyChecking=no "$VPS_USER@$VPS_HOST" \
    "DEPLOY_DIR=$DEPLOY_DIR SERVICE_USER=openbrain SERVICE_GROUP=openbrain ORT_VERSION=1.24.3 bash -s" <<'EOS'
set -euo pipefail

DEPLOY_DIR="${DEPLOY_DIR:-/opt/openbrain-mcp}"
SERVICE_USER="${SERVICE_USER:-openbrain}"
SERVICE_GROUP="${SERVICE_GROUP:-openbrain}"
ORT_VERSION="${ORT_VERSION:-1.24.3}"

if command -v apt-get >/dev/null 2>&1; then
    export DEBIAN_FRONTEND=noninteractive
    apt-get update
    apt-get install -y --no-install-recommends ca-certificates curl tar libssl3
fi

if ! getent group "$SERVICE_GROUP" >/dev/null 2>&1; then
    groupadd --system "$SERVICE_GROUP"
fi

if ! id -u "$SERVICE_USER" >/dev/null 2>&1; then
    useradd --system --gid "$SERVICE_GROUP" --home "$DEPLOY_DIR" --shell /usr/sbin/nologin "$SERVICE_USER"
fi

install -d -m 0755 "$DEPLOY_DIR" "$DEPLOY_DIR/models" "$DEPLOY_DIR/logs" "$DEPLOY_DIR/lib"

ARCH="$(uname -m)"
case "$ARCH" in
    x86_64) ORT_ARCH="x64" ;;
    aarch64|arm64) ORT_ARCH="aarch64" ;;
    *) echo "Unsupported arch: $ARCH"; exit 1 ;;
esac

if [[ ! -f "$DEPLOY_DIR/lib/libonnxruntime.so" ]]; then
    TMP_DIR="$(mktemp -d)"
    ORT_TGZ="onnxruntime-linux-${ORT_ARCH}-${ORT_VERSION}.tgz"
    ORT_URL="https://github.com/microsoft/onnxruntime/releases/download/v${ORT_VERSION}/${ORT_TGZ}"
    curl -fL "$ORT_URL" -o "$TMP_DIR/$ORT_TGZ"
    tar -xzf "$TMP_DIR/$ORT_TGZ" -C "$TMP_DIR"
    ORT_ROOT="$TMP_DIR/onnxruntime-linux-${ORT_ARCH}-${ORT_VERSION}"
    cp "$ORT_ROOT/lib/libonnxruntime.so" "$DEPLOY_DIR/lib/libonnxruntime.so"
    cp "$ORT_ROOT/lib/libonnxruntime.so.${ORT_VERSION}" "$DEPLOY_DIR/lib/libonnxruntime.so.${ORT_VERSION}" || true
    rm -rf "$TMP_DIR"
fi

ENV_FILE="$DEPLOY_DIR/.env"
if [[ ! -f "$ENV_FILE" ]]; then
    if [[ -f "$DEPLOY_DIR/.env.example" ]]; then
        cp "$DEPLOY_DIR/.env.example" "$ENV_FILE"
    else
        touch "$ENV_FILE"
    fi
fi

upsert_env() {
    local key="$1"
    local value="$2"
    if grep -qE "^${key}=" "$ENV_FILE"; then
        sed -i "s|^${key}=.*|${key}=${value}|" "$ENV_FILE"
    else
        printf '%s=%s\n' "$key" "$value" >> "$ENV_FILE"
    fi
}

upsert_env "OPENBRAIN__EMBEDDING__MODEL_PATH" "$DEPLOY_DIR/models/all-MiniLM-L6-v2"
upsert_env "ORT_DYLIB_PATH" "$DEPLOY_DIR/lib/libonnxruntime.so"
upsert_env "OPENBRAIN__SERVER__HOST" "0.0.0.0"

chmod +x "$DEPLOY_DIR/openbrain-mcp" "$DEPLOY_DIR/.gitea/download-model.sh"
chown -R "$SERVICE_USER:$SERVICE_GROUP" "$DEPLOY_DIR"
EOS

# Step 5d: Run database migrations with the newly deployed binary
log_info "Running database migrations..."
ssh_cmd "cd $DEPLOY_DIR && ./openbrain-mcp migrate"

# Step 6: Install systemd service
log_info "Installing systemd service..."
scp_cmd "$SCRIPT_DIR/openbrain.service" "$VPS_USER@$VPS_HOST:/etc/systemd/system/$SERVICE_NAME.service"
ssh_cmd "systemctl daemon-reload"
ssh_cmd "systemctl enable $SERVICE_NAME"

# Step 7: Restart service
log_info "Restarting service..."
ssh_cmd "systemctl restart $SERVICE_NAME"
sleep 2

# Step 8: Check status
log_info "Checking service status..."
ssh_cmd "systemctl status $SERVICE_NAME --no-pager" || true

log_info "Deployment complete!"
log_info "Service URL: http://$VPS_HOST:3100/mcp/health"