#!/usr/bin/env bash
# memory-import.sh — Import a Neuron .neuronmem bundle onto this device
#
# Usage:
#   ./tools/memory-import.sh input.neuronmem [--passphrase "your passphrase"]
#   ./tools/memory-import.sh input.neuronmem [--dry-run]   # verify only, no changes
#
# The script will:
#   1. Decrypt and unpack the .neuronmem file
#   2. Validate the checksum and version
#   3. Back up the current snapshot.json
#   4. Stop the soul service
#   5. Replace snapshot.json
#   6. Restart the soul service
#   7. Verify the soul came back up

set -euo pipefail

# ── Config ─────────────────────────────────────────────────────────────────────
ENGRAM_SNAPSHOT="${HOME}/.neuron/engram/snapshot.json"
SOUL_SERVICE="ai.neurontechnologies.soul"
SOUL_PORT="7770"
SOUL_STARTUP_TIMEOUT=30   # seconds to wait for soul to come back

# ── Parse args ─────────────────────────────────────────────────────────────────
INPUT_PATH=""
PASSPHRASE=""
PASSPHRASE_SET=0
DRY_RUN=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --passphrase)
      PASSPHRASE="$2"
      PASSPHRASE_SET=1
      shift 2
      ;;
    --passphrase=*)
      PASSPHRASE="${1#*=}"
      PASSPHRASE_SET=1
      shift
      ;;
    --dry-run)
      DRY_RUN=1
      shift
      ;;
    -*)
      echo "Unknown option: $1" >&2
      echo "Usage: $0 input.neuronmem [--passphrase \"...\"] [--dry-run]" >&2
      exit 1
      ;;
    *)
      if [[ -z "$INPUT_PATH" ]]; then
        INPUT_PATH="$1"
      else
        echo "Unexpected argument: $1" >&2
        exit 1
      fi
      shift
      ;;
  esac
done

if [[ -z "$INPUT_PATH" ]]; then
  echo "ERROR: No input file specified." >&2
  echo "Usage: $0 input.neuronmem [--passphrase \"...\"] [--dry-run]" >&2
  exit 1
fi

if [[ ! -f "$INPUT_PATH" ]]; then
  echo "ERROR: Input file not found: $INPUT_PATH" >&2
  exit 1
fi

echo "Neuron Memory Import"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "Source:   $INPUT_PATH"
echo "Target:   $ENGRAM_SNAPSHOT"
if [[ $DRY_RUN -eq 1 ]]; then
  echo "Mode:     DRY RUN (no changes will be made)"
fi
echo ""

# ── Prompt for passphrase if needed ───────────────────────────────────────────
if [[ $PASSPHRASE_SET -eq 0 ]]; then
  read -r -s -p "Enter passphrase: " PASSPHRASE
  echo ""
  if [[ -z "$PASSPHRASE" ]]; then
    echo "ERROR: Passphrase cannot be empty." >&2
    exit 1
  fi
fi

# ── Decrypt to temp dir ────────────────────────────────────────────────────────
WORK_DIR="$(mktemp -d)"
CLEANUP() {
  rm -rf "$WORK_DIR"
}
trap CLEANUP EXIT

TAR_PATH="${WORK_DIR}/bundle.tar.gz"

echo "Decrypting..."
if ! openssl enc -d -aes-256-cbc \
    -pbkdf2 \
    -iter 600000 \
    -in "$INPUT_PATH" \
    -out "$TAR_PATH" \
    -pass "pass:${PASSPHRASE}" 2>/dev/null; then
  echo "ERROR: Decryption failed. Wrong passphrase or corrupted file." >&2
  exit 1
fi
echo "  Decrypted successfully."

# ── Unpack ─────────────────────────────────────────────────────────────────────
echo "Unpacking..."
(cd "$WORK_DIR" && tar xzf "$TAR_PATH") || {
  echo "ERROR: Failed to unpack bundle. File may be corrupted." >&2
  exit 1
}

# Locate the bundle directory (neuronmem-v1/)
BUNDLE_DIR=""
for d in "${WORK_DIR}"/neuronmem-v*/; do
  if [[ -d "$d" ]]; then
    BUNDLE_DIR="$d"
    break
  fi
done

if [[ -z "$BUNDLE_DIR" ]]; then
  echo "ERROR: Bundle directory not found. Invalid .neuronmem file." >&2
  exit 1
fi

METADATA_FILE="${BUNDLE_DIR}metadata.json"
NODES_FILE="${BUNDLE_DIR}nodes.json"

if [[ ! -f "$METADATA_FILE" ]]; then
  echo "ERROR: metadata.json missing from bundle." >&2
  exit 1
fi

if [[ ! -f "$NODES_FILE" ]]; then
  echo "ERROR: nodes.json missing from bundle." >&2
  exit 1
fi

# ── Validate metadata ──────────────────────────────────────────────────────────
echo "Validating metadata..."
FORMAT_VERSION="$(python3 -c "import json; d=json.load(open('${METADATA_FILE}')); print(d.get('version','?'))")"
EXPORTED_AT="$(python3 -c "import json; d=json.load(open('${METADATA_FILE}')); print(d.get('exported_at','?'))")"
EXPECTED_COUNT="$(python3 -c "import json; d=json.load(open('${METADATA_FILE}')); print(d.get('node_count','?'))")"
STORED_CHECKSUM="$(python3 -c "import json; d=json.load(open('${METADATA_FILE}')); print(d.get('sha256','?'))")"
SOURCE_HOST="$(python3 -c "import json; d=json.load(open('${METADATA_FILE}')); print(d.get('source_host','?'))")"

echo "  Format version: ${FORMAT_VERSION}"
echo "  Exported at:    ${EXPORTED_AT}"
echo "  Source host:    ${SOURCE_HOST}"
echo "  Expected nodes: ${EXPECTED_COUNT}"

if [[ "$FORMAT_VERSION" != "1" ]]; then
  echo "ERROR: Unsupported bundle format version: ${FORMAT_VERSION}" >&2
  echo "       This tool supports version 1 only." >&2
  exit 1
fi

# ── Validate checksum ──────────────────────────────────────────────────────────
echo "Verifying checksum..."
ACTUAL_CHECKSUM="$(openssl dgst -sha256 "$NODES_FILE" | awk '{print $NF}')"

if [[ "$ACTUAL_CHECKSUM" != "$STORED_CHECKSUM" ]]; then
  echo "ERROR: Checksum mismatch!" >&2
  echo "  Expected: ${STORED_CHECKSUM}" >&2
  echo "  Got:      ${ACTUAL_CHECKSUM}" >&2
  echo "  The bundle may be corrupted." >&2
  exit 1
fi
echo "  Checksum OK: ${ACTUAL_CHECKSUM:0:16}..."

# ── Verify node count ──────────────────────────────────────────────────────────
echo "Verifying node count..."
ACTUAL_COUNT="$(python3 -c "
import json
with open('${NODES_FILE}') as f:
    d = json.load(f)
nodes = d.get('nodes', d if isinstance(d, list) else [])
print(len(nodes) if isinstance(nodes, list) else len(nodes))
" 2>/dev/null || echo "unknown")"

echo "  Found ${ACTUAL_COUNT} nodes (expected ${EXPECTED_COUNT})"

if [[ "$ACTUAL_COUNT" != "$EXPECTED_COUNT" && "$EXPECTED_COUNT" != "unknown" ]]; then
  echo "WARNING: Node count mismatch (expected ${EXPECTED_COUNT}, found ${ACTUAL_COUNT})." >&2
  echo "         Proceeding anyway — count may differ if nodes were deduplicated." >&2
fi

# ── Dry run exit ───────────────────────────────────────────────────────────────
if [[ $DRY_RUN -eq 1 ]]; then
  echo ""
  echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
  echo "DRY RUN complete. Bundle is valid."
  echo "  Nodes: ${ACTUAL_COUNT}"
  echo "  Checksum: verified"
  echo "  Run without --dry-run to import."
  echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
  exit 0
fi

# ── Safety confirmation ────────────────────────────────────────────────────────
echo ""
echo "WARNING: This will replace your current Neuron memory store."
echo "  Current snapshot: $ENGRAM_SNAPSHOT"
echo "  A backup will be created before replacing."
echo ""
read -r -p "Type 'yes' to continue: " CONFIRM
if [[ "$CONFIRM" != "yes" ]]; then
  echo "Aborted."
  exit 0
fi

# ── Backup existing snapshot ───────────────────────────────────────────────────
BACKUP_TIMESTAMP="$(date -u +"%Y%m%dT%H%M%SZ")"
ENGRAM_DIR="$(dirname "$ENGRAM_SNAPSHOT")"
BACKUP_PATH="${HOME}/.neuron/engram-backup-${BACKUP_TIMESTAMP}.tar.gz"

echo ""
echo "Backing up current snapshot..."
if [[ -f "$ENGRAM_SNAPSHOT" ]]; then
  (cd "$HOME/.neuron" && tar czf "$BACKUP_PATH" "$(basename "$ENGRAM_DIR")/snapshot.json" 2>/dev/null) || \
  cp "$ENGRAM_SNAPSHOT" "${ENGRAM_SNAPSHOT}.backup-${BACKUP_TIMESTAMP}"
  echo "  Backup: $BACKUP_PATH"
else
  echo "  No existing snapshot to back up."
fi

# ── Stop soul service ──────────────────────────────────────────────────────────
echo "Stopping soul service (${SOUL_SERVICE})..."
launchctl stop "$SOUL_SERVICE" 2>/dev/null || true
# Also stop engram service if running
launchctl stop "ai.neuron.engram" 2>/dev/null || true
sleep 2
echo "  Soul stopped."

# ── Replace snapshot.json ──────────────────────────────────────────────────────
echo "Installing new snapshot..."
cp "$NODES_FILE" "$ENGRAM_SNAPSHOT"
echo "  snapshot.json replaced ($(du -sh "$ENGRAM_SNAPSHOT" | cut -f1))"

# ── Restart soul service ───────────────────────────────────────────────────────
echo "Restarting soul service..."
launchctl start "$SOUL_SERVICE" 2>/dev/null || true
launchctl start "ai.neuron.engram" 2>/dev/null || true

# ── Wait for soul to come up ───────────────────────────────────────────────────
echo "Waiting for soul to come up on port ${SOUL_PORT}..."
ELAPSED=0
SOUL_UP=0
while [[ $ELAPSED -lt $SOUL_STARTUP_TIMEOUT ]]; do
  if curl -sf "http://localhost:${SOUL_PORT}/" > /dev/null 2>&1; then
    SOUL_UP=1
    break
  fi
  # Try a known endpoint that returns any response (even 404 means it's up)
  HTTP_CODE="$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:${SOUL_PORT}/api/neuron/memory" 2>/dev/null || echo "000")"
  if [[ "$HTTP_CODE" != "000" ]]; then
    SOUL_UP=1
    break
  fi
  sleep 1
  ELAPSED=$((ELAPSED + 1))
done

if [[ $SOUL_UP -eq 1 ]]; then
  echo "  Soul is up (responded in ${ELAPSED}s)."
else
  echo "  WARNING: Soul did not respond within ${SOUL_STARTUP_TIMEOUT}s."
  echo "  The service may still be starting. Check: launchctl list | grep soul"
fi

# ── Final report ───────────────────────────────────────────────────────────────
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "Import complete."
echo "  Nodes imported: ${ACTUAL_COUNT}"
echo "  Exported at:    ${EXPORTED_AT}"
echo "  Source host:    ${SOURCE_HOST}"
echo "  Backup:         ${BACKUP_PATH}"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"