#!/usr/bin/env bash
# =============================================================================
# Pyvorin Edge SDK — In-place Updater
# =============================================================================
# Run from your existing install directory:
#   cd ~/pyvorin-edge && bash update-pyvorin-edge.sh
#
# Or specify the directory:
#   bash update-pyvorin-edge.sh --dir /opt/pyvorin-edge
# =============================================================================

set -euo pipefail

BUNDLE_URL="https://pypi.pyvorin.com/pifiles/pyvorin-edge-latest.tar.gz"
INSTALL_DIR=""
SKIP_NEON=0
SKIP_TESTS=0

# Colours
if [[ -t 1 ]]; then
    GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; BLUE='\033[0;34m'; BOLD='\033[1m'; NC='\033[0m'
else
    GREEN=''; YELLOW=''; RED=''; BLUE=''; BOLD=''; NC=''
fi

log_info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
log_ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
log_err()   { echo -e "${RED}[ERROR]${NC} $*" >&2; }
log_title() { echo -e "\n${BOLD}$*${NC}"; }

usage() {
cat <<'USAGE'
Pyvorin Edge SDK — In-place Updater

Usage:
  cd ~/pyvorin-edge && bash update-pyvorin-edge.sh
  bash update-pyvorin-edge.sh --dir /opt/pyvorin-edge

Options:
  --dir PATH       Install directory (default: auto-detect from CWD)
  --skip-neon      Skip ARM64 NEON rebuild
  --skip-tests     Skip smoke tests
  --help           Show this help
USAGE
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --dir)       INSTALL_DIR="$2"; shift 2 ;;
        --skip-neon) SKIP_NEON=1; shift ;;
        --skip-tests) SKIP_TESTS=1; shift ;;
        --help|-h)   usage; exit 0 ;;
        *) log_err "Unknown option: $1"; usage; exit 1 ;;
    esac
done

# Auto-detect install dir
if [[ -z "${INSTALL_DIR}" ]]; then
    if [[ -f "pyproject.toml" && -d "edge_sdk" && -d "venv" ]]; then
        INSTALL_DIR="$(pwd)"
    elif [[ -d "${HOME}/pyvorin-edge" ]]; then
        INSTALL_DIR="${HOME}/pyvorin-edge"
    else
        log_err "Cannot find a pyvorin-edge install."
        log_err "Run from inside the install directory, or use --dir."
        exit 1
    fi
fi

if [[ ! -d "${INSTALL_DIR}/venv" ]]; then
    log_err "No venv found in ${INSTALL_DIR} — this doesn't look like a pyvorin-edge install."
    exit 1
fi

cd "${INSTALL_DIR}"
log_title "Pyvorin Edge SDK — In-place Update"
log_info "Directory:  ${INSTALL_DIR}"
log_info "Bundle:     ${BUNDLE_URL}"

# ---------------------------------------------------------------------------
# Step 1 — Backup config
# ---------------------------------------------------------------------------
log_title "Step 1/5 — Backup config.toml"
if [[ -f "config.toml" ]]; then
    cp config.toml "config.toml.backup.$(date +%Y%m%d_%H%M%S)"
    log_ok "Backed up config.toml"
else
    log_warn "No config.toml found — will use defaults after update."
fi

# ---------------------------------------------------------------------------
# Step 2 — Download new bundle
# ---------------------------------------------------------------------------
log_title "Step 2/5 — Download latest bundle"
TMP_TAR="$(mktemp /tmp/pyvorin-edge-update.XXXXXX.tar.gz)"
trap 'rm -f "${TMP_TAR}"' EXIT

if ! curl -fsSL "${BUNDLE_URL}" -o "${TMP_TAR}"; then
    log_err "Failed to download bundle."
    exit 1
fi
log_ok "Bundle downloaded ($(du -h "${TMP_TAR}" | cut -f1))"

# ---------------------------------------------------------------------------
# Step 3 — Extract over existing code
# ---------------------------------------------------------------------------
log_title "Step 3/5 — Apply update (preserves config, data, logs)"
tar xzf "${TMP_TAR}" --overwrite \
    edge_sdk/ edge_runtime/ examples/ tools/ pifiles/ tests/ kernels/
rm -f "${TMP_TAR}"
log_ok "Source code updated."

# ---------------------------------------------------------------------------
# Step 4 — Rebuild / reinstall
# ---------------------------------------------------------------------------
log_title "Step 4/5 — Reinstall Python package"
source "${INSTALL_DIR}/venv/bin/activate"
pip install --quiet --upgrade pip setuptools wheel 2>/dev/null || true
pip install --quiet -e ".[dev]" 2>/dev/null || {
    log_warn "Editable install of pyproject.toml failed (expected on Pi without compiler source)."
    log_info "Installing dev dependencies directly ..."
    pip install --quiet pytest pytest-cov pytest-timeout 2>/dev/null || true
}
log_ok "Python dependencies ready."

# NEON rebuild
ARCH="$(uname -m)"
if [[ "${SKIP_NEON}" -eq 0 && ( "${ARCH}" == "aarch64" || "${ARCH}" == "arm64" ) ]]; then
    log_info "Rebuilding ARM64 NEON extension ..."
    cd "${INSTALL_DIR}/kernels/arm64_neon"
    python setup.py build_ext --inplace 2>&1 | tail -3
    cd "${INSTALL_DIR}"
    log_ok "NEON extension rebuilt."
elif [[ "${SKIP_NEON}" -eq 1 ]]; then
    log_warn "Skipping NEON rebuild (--skip-neon)"
else
    log_warn "Not on ARM64 — skipping NEON rebuild."
fi

# ---------------------------------------------------------------------------
# Step 5 — Smoke tests
# ---------------------------------------------------------------------------
if [[ "${SKIP_TESTS}" -eq 0 ]]; then
    log_title "Step 5/5 — Smoke tests"
    python -c "from edge_sdk.pyvorin_edge.cost_model import CostModel; print('CostModel OK')" || true
    python -c "from edge_sdk.pyvorin_edge.cost_model import CostModel; print('Currency:', CostModel().to_dict()['pricing']['currency'])" 2>/dev/null || true
    pytest tests/edge_sdk/ tests/edge_runtime/ tests/kernels/ -q --tb=line 2>&1 | tail -15 || true
    log_ok "Smoke tests complete."
else
    log_warn "Skipping smoke tests (--skip-tests)"
fi

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
log_title "============================================================"
log_title "                    UPDATE COMPLETE"
log_title "============================================================"
echo ""
echo "  Directory:   ${INSTALL_DIR}"
echo "  Python:      $(python --version 2>&1)"
echo "  Architecture: ${ARCH}"
echo ""
log_ok "Pyvorin Edge SDK updated to latest."
echo ""
echo "  Quick checks:"
echo "    python -c \"from edge_sdk.pyvorin_edge import CostModel; print(CostModel)\""
echo "    python examples/industry_benchmarks/run_all.py"
echo "    bash pifiles/run-pi-benchmarks.sh"
echo ""

exit 0