#!/usr/bin/env bash
# =============================================================================
# Pyvorin Edge SDK — Raspberry Pi 5 Installer
# =============================================================================
# Download and run on your Pi:
#   curl -fsSL https://pypi.pyvorin.com/pifiles/install-pyvorin-edge.sh | bash
#
# Or save first, then run:
#   curl -fsSL https://pypi.pyvorin.com/pifiles/install-pyvorin-edge.sh -o install.sh
#   bash install.sh
#
# With sudo for system-wide install (/opt/pyvorin-edge):
#   curl -fsSL ... | sudo bash
#
# With options:
#   bash install.sh --system --branch main --skip-tests
# =============================================================================

set -euo pipefail

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
SCRIPT_NAME="${0##*/}"
INSTALL_USER="${SUDO_USER:-$USER}"
INSTALL_HOME="$(eval echo ~"$INSTALL_USER")"

# Detect if we have sudo / are root
if [[ "$(id -u)" -eq 0 ]]; then
    HAS_ROOT=1
else
    HAS_ROOT=0
fi

# Defaults
USE_SYSTEM=0
SKIP_TESTS=0
SKIP_DEMO=0
SKIP_BENCHMARK=0
BUNDLE_URL="https://pypi.pyvorin.com/pifiles/pyvorin-edge-latest.tar.gz"
NO_DOWNLOAD=0
INSTALL_DIR=""
PY_VERSION=""
PY_BIN="python3"

# ---------------------------------------------------------------------------
# Auto-detect Python
# ---------------------------------------------------------------------------
detect_python() {
    # Prefer python3 if it exists and is >= 3.10
    if command -v python3 &>/dev/null; then
        local ver
        ver=$(python3 --version 2>&1 | awk '{print $2}')
        local major_minor
        major_minor=$(echo "$ver" | cut -d. -f1,2)
        # Check version is at least 3.10
        if printf '%s\n' "3.10" "$major_minor" | sort -V -C; then
            PY_BIN="python3"
            PY_VERSION="$ver"
            return 0
        fi
    fi
    # Fallback: try specific versions
    for try in python3.13 python3.12 python3.11 python3.10; do
        if command -v "$try" &>/dev/null; then
            PY_BIN="$try"
            PY_VERSION=$("$try" --version 2>&1 | awk '{print $2}')
            return 0
        fi
    done
    log_err "No suitable Python (>= 3.10) found. Install python3 and re-run."
    exit 1
}

detect_python

# ---------------------------------------------------------------------------
# Colours
# ---------------------------------------------------------------------------
if [[ -t 1 ]]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    BLUE='\033[0;34m'
    BOLD='\033[1m'
    NC='\033[0m'
else
    RED='' GREEN='' YELLOW='' 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
# ---------------------------------------------------------------------------
usage() {
cat <<EOF
Pyvorin Edge SDK — Raspberry Pi 5 Installer

Usage:
  curl -fsSL https://pypi.pyvorin.com/pifiles/install-pyvorin-edge.sh | bash
  bash ${SCRIPT_NAME} [OPTIONS]

Options:
  --system           Install to /opt/pyvorin-edge (requires sudo)
  --user             Install to ~/pyvorin-edge (default if no sudo)
  --dir PATH         Override install directory
  --bundle-url URL   Tarball URL to download (default: https://pypi.pyvorin.com/pifiles/pyvorin-edge-latest.tar.gz)
  --no-download      Skip download; use existing repo in current directory
  --py-version VER   Python version (default: auto-detect)
  --skip-tests       Skip pytest suite
  --skip-demo        Skip housing safety demo
  --skip-benchmark   Skip benchmark CLI
  --help             Show this help

Folders:
  User install  (default) : ~/pyvorin-edge/
  System install (--system) : /opt/pyvorin-edge/
  Config                : <install-dir>/config.toml
  Database              : <install-dir>/data/edge.db
  Logs                  : <install-dir>/logs/

Prerequisites on Pi:
  sudo apt install -y python3.12 python3.12-venv python3.12-dev build-essential libffi-dev

What this script does:
  1. Verify ARM64 architecture
  2. Install system packages (apt)
  3. Clone or use existing repo
  4. Create Python venv
  5. pip install Edge SDK + dependencies
  6. Build ARM64 NEON C extension
  7. Run 423 test suite
  8. Run housing safety demo
  9. Run benchmark CLI
  10. Print PASS/FAIL summary
EOF
}

# ---------------------------------------------------------------------------
# Parse args
# ---------------------------------------------------------------------------
while [[ $# -gt 0 ]]; do
    case "$1" in
        --system)      USE_SYSTEM=1; shift ;;
        --user)        USE_SYSTEM=0; shift ;;
        --dir)         INSTALL_DIR="$2"; shift 2 ;;
        --bundle-url)  BUNDLE_URL="$2"; shift 2 ;;
        --no-download) NO_DOWNLOAD=1; shift ;;
        --py-version)  PY_VERSION="$2"; PY_BIN="python${2}"; shift 2 ;;
        --skip-tests)  SKIP_TESTS=1; shift ;;
        --skip-demo)   SKIP_DEMO=1; shift ;;
        --skip-benchmark) SKIP_BENCHMARK=1; shift ;;
        --help|-h)     usage; exit 0 ;;
        *) log_err "Unknown option: $1"; usage; exit 1 ;;
    esac
done

# ---------------------------------------------------------------------------
# Determine install directory
# ---------------------------------------------------------------------------
if [[ -n "${INSTALL_DIR}" ]]; then
    : # user override
elif [[ "${USE_SYSTEM}" -eq 1 ]]; then
    if [[ "${HAS_ROOT}" -eq 0 ]]; then
        log_err "--system requires root/sudo. Run with sudo or omit --system for user install."
        exit 1
    fi
    INSTALL_DIR="/opt/pyvorin-edge"
else
    INSTALL_DIR="${INSTALL_HOME}/pyvorin-edge"
fi

log_title "Pyvorin Edge SDK — Raspberry Pi 5 Installer"
log_info "Architecture:    $(uname -m)"
log_info "Install dir:     ${INSTALL_DIR}"
log_info "Python binary:   ${PY_BIN}"
log_info "Python version:  ${PY_VERSION}"
log_info "Run as root:     ${HAS_ROOT}"
log_info "Skip tests:      ${SKIP_TESTS}"
log_info "Skip demo:       ${SKIP_DEMO}"
log_info "Skip benchmark:  ${SKIP_BENCHMARK}"

# ---------------------------------------------------------------------------
# Step 1 — Verify ARM64
# ---------------------------------------------------------------------------
log_title "Step 1/10 — Verify ARM64 architecture"
ARCH="$(uname -m)"
if [[ "${ARCH}" != "aarch64" && "${ARCH}" != "arm64" ]]; then
    log_warn "Architecture is '${ARCH}', not aarch64/arm64."
    log_warn "NEON kernels will use scalar fallback. Continuing anyway..."
else
    log_ok "ARM64 confirmed (${ARCH})."
fi

# ---------------------------------------------------------------------------
# Step 2 — System packages
# ---------------------------------------------------------------------------
log_title "Step 2/10 — Install system packages"

ensure_apt_pkg() {
    local pkg="$1"
    if dpkg -s "${pkg}" &>/dev/null; then
        log_ok "${pkg} already installed"
        return 0
    fi
    if [[ "${HAS_ROOT}" -eq 0 ]]; then
        log_warn "Missing apt package: ${pkg}. Run with sudo or install manually:"
        log_warn "  sudo apt install -y ${pkg}"
        return 1
    fi
    log_info "Installing ${pkg} ..."
    apt-get update -qq
    apt-get install -y -qq "${pkg}"
}

NEEDED_PKGS=("python3" "python3-venv" "python3-dev" "build-essential" "libffi-dev" "git" "curl")
MISSING_APT=0
for pkg in "${NEEDED_PKGS[@]}"; do
    if ! ensure_apt_pkg "${pkg}"; then
        MISSING_APT=1
    fi
done

if [[ "${MISSING_APT}" -eq 1 ]]; then
    log_err "Some system packages are missing. Install them and re-run."
    exit 1
fi
log_ok "System packages ready."

# ---------------------------------------------------------------------------
# Step 3 — Get code (download bundle or use existing)
# ---------------------------------------------------------------------------
log_title "Step 3/10 — Get Pyvorin Edge SDK code"

# Check if we're already inside an extracted repo
if [[ -f "pyproject.toml" && -d "edge_sdk" && -d "kernels" && "${NO_DOWNLOAD}" -eq 1 ]]; then
    REPO_DIR="$(pwd)"
    log_info "Using existing repo at ${REPO_DIR} (--no-download)"
else
    REPO_DIR="${INSTALL_DIR}"
    # Clean old extract if present
    if [[ -d "${REPO_DIR}" && "${NO_DOWNLOAD}" -eq 0 ]]; then
        log_info "Removing old install at ${REPO_DIR} ..."
        rm -rf "${REPO_DIR}"
    fi
    mkdir -p "${REPO_DIR}"

    if [[ "${NO_DOWNLOAD}" -eq 1 ]]; then
        log_err "--no-download set but no repo found in current directory."
        log_err "Run this script from inside the pyvorin-edge directory, or omit --no-download."
        exit 1
    fi

    log_info "Downloading bundle from ${BUNDLE_URL} ..."
    TMP_TAR="$(mktemp /tmp/pyvorin-edge.XXXXXX.tar.gz)"
    trap 'rm -f "${TMP_TAR}"' EXIT

    if ! curl -fsSL "${BUNDLE_URL}" -o "${TMP_TAR}"; then
        log_err "Failed to download bundle from ${BUNDLE_URL}"
        log_err "Check the URL and network connectivity."
        exit 1
    fi

    log_info "Extracting to ${REPO_DIR} ..."
    tar xzf "${TMP_TAR}" -C "${REPO_DIR}"
    rm -f "${TMP_TAR}"
fi

cd "${REPO_DIR}"
log_ok "Code ready at $(pwd)"

# ---------------------------------------------------------------------------
# Step 4 — Python venv
# ---------------------------------------------------------------------------
log_title "Step 4/10 — Create Python virtual environment"

VENV_DIR="${REPO_DIR}/venv"
if [[ ! -d "${VENV_DIR}" ]]; then
    "${PY_BIN}" -m venv "${VENV_DIR}"
    log_ok "Created venv at ${VENV_DIR}"
else
    log_ok "Using existing venv at ${VENV_DIR}"
fi

source "${VENV_DIR}/bin/activate"
python --version

# ---------------------------------------------------------------------------
# Step 5 — pip install
# ---------------------------------------------------------------------------
log_title "Step 5/10 — Install Python dependencies"

pip install --quiet --upgrade pip setuptools wheel
pip install --quiet -e ".[dev]"
log_ok "Python dependencies installed."

# ---------------------------------------------------------------------------
# Step 6 — Build NEON extension
# ---------------------------------------------------------------------------
log_title "Step 6/10 — Build ARM64 NEON extension"

if [[ "${ARCH}" == "aarch64" || "${ARCH}" == "arm64" ]]; then
    cd "${REPO_DIR}/kernels/arm64_neon"
    python setup.py build_ext --inplace 2>&1 | tail -5
    cd "${REPO_DIR}"
    log_ok "NEON extension built."
else
    log_warn "Not on ARM64 — skipping NEON build (scalar fallback only)."
fi

# ---------------------------------------------------------------------------
# Step 7 — Run tests
# ---------------------------------------------------------------------------
TEST_RC=0
if [[ "${SKIP_TESTS}" -eq 1 ]]; then
    log_warn "Skipping tests (--skip-tests)"
else
    log_title "Step 7/10 — Run test suite"
    # On the Pi we skip integration tests — they need the pyvorin compiler
    # which is a dev/build tool, not deployed to the gateway.
    pytest tests/edge_sdk/ tests/edge_runtime/ tests/kernels/ \
        -v --tb=short --confcutdir=tests/edge_sdk 2>&1 | tail -30 || TEST_RC=$?
    TEST_RC=${TEST_RC:-0}
    if [[ "${TEST_RC}" -eq 0 ]]; then
        log_ok "All tests passed."
    else
        log_warn "Some tests failed (exit ${TEST_RC})."
    fi
fi

# ---------------------------------------------------------------------------
# Step 8 — Run demo
# ---------------------------------------------------------------------------
DEMO_RC=0
if [[ "${SKIP_DEMO}" -eq 1 ]]; then
    log_warn "Skipping demo (--skip-demo)"
else
    log_title "Step 8/10 — Run housing safety demo"
    python examples/housing_safety/pipeline.py 2>&1 | tail -20 || DEMO_RC=$?
    DEMO_RC=${DEMO_RC:-0}
    if [[ "${DEMO_RC}" -eq 0 ]]; then
        log_ok "Demo completed."
    else
        log_warn "Demo failed (exit ${DEMO_RC})."
    fi
fi

# ---------------------------------------------------------------------------
# Step 9 — Run benchmark
# ---------------------------------------------------------------------------
BENCH_RC=0
if [[ "${SKIP_BENCHMARK}" -eq 1 ]]; then
    log_warn "Skipping benchmark (--skip-benchmark)"
else
    log_title "Step 9/10 — Run benchmark CLI"
    python tools/pyv-edge-benchmark compare --suite sensor_pipeline --iterations 100 2>&1 | tail -30 || BENCH_RC=$?
    BENCH_RC=${BENCH_RC:-0}
    if [[ "${BENCH_RC}" -eq 0 ]]; then
        log_ok "Benchmark completed."
    else
        log_warn "Benchmark failed (exit ${BENCH_RC})."
    fi
fi

# ---------------------------------------------------------------------------
# Step 10 — Systemd service (optional, only with --system)
# ---------------------------------------------------------------------------
log_title "Step 10/10 — Optional systemd service"

if [[ "${USE_SYSTEM}" -eq 1 && "${HAS_ROOT}" -eq 1 ]]; then
    if [[ -f "tools/pyvorin-edge.service" ]]; then
        cp tools/pyvorin-edge.service /etc/systemd/system/pyvorin-edge.service
        systemctl daemon-reload
        log_ok "Installed systemd service: pyvorin-edge"
        log_info "Start it with: sudo systemctl start pyvorin-edge"
        log_info "Enable auto-start: sudo systemctl enable pyvorin-edge"
    fi
else
    log_info "Skipping systemd install (use --system with sudo to install service)."
fi

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
log_title "============================================================"
log_title "                    INSTALL SUMMARY"
log_title "============================================================"
echo ""
echo "  Install directory : ${INSTALL_DIR}"
echo "  Virtual env       : ${VENV_DIR}"
echo "  Python            : $(python --version 2>&1)"
echo "  Architecture      : ${ARCH}"
echo ""
echo "  Tests             : $([[ ${SKIP_TESTS} -eq 1 ]] && echo SKIPPED || ([[ ${TEST_RC} -eq 0 ]] && echo PASS || echo "FAIL (exit ${TEST_RC})"))"
echo "  Demo              : $([[ ${SKIP_DEMO} -eq 1 ]] && echo SKIPPED || ([[ ${DEMO_RC} -eq 0 ]] && echo PASS || echo "FAIL (exit ${DEMO_RC})"))"
echo "  Benchmark         : $([[ ${SKIP_BENCHMARK} -eq 1 ]] && echo SKIPPED || ([[ ${BENCH_RC} -eq 0 ]] && echo PASS || echo "FAIL (exit ${BENCH_RC})"))"
echo ""

OVERALL_RC=0
if [[ "${SKIP_TESTS}" -eq 0 && "${TEST_RC}" -ne 0 ]]; then OVERALL_RC=1; fi
if [[ "${SKIP_DEMO}" -eq 0 && "${DEMO_RC}" -ne 0 ]]; then OVERALL_RC=1; fi
if [[ "${SKIP_BENCHMARK}" -eq 0 && "${BENCH_RC}" -ne 0 ]]; then OVERALL_RC=1; fi

if [[ "${OVERALL_RC}" -eq 0 ]]; then
    log_ok "INSTALL COMPLETE — Pyvorin Edge SDK is ready on your Pi 5."
else
    log_warn "INSTALL COMPLETE WITH WARNINGS — check output above."
fi

echo ""
echo "  Next steps:"
echo "    cd ${INSTALL_DIR}"
echo "    source venv/bin/activate"
echo "    python tools/pyv-edge-benchmark --help"
echo "    python examples/housing_safety/pipeline.py"
echo ""
echo "  For OEM integration docs:"
echo "    cat docs/edge_sdk/OEM_INTEGRATION.md"
echo ""

exit "${OVERALL_RC}"