#!/usr/bin/env bash
# Lambda Mindlink Memotron — Linux installer
# Supports: Debian/Ubuntu/Mint (apt) and Fedora/RHEL (dnf/yum)
#
# Usage:
#   ./install.sh               # interactive menu
#   ./install.sh cuda          # non-interactive: cuda | rocm | cpu
set -euo pipefail

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; NC='\033[0m'
info() { printf "${GREEN}[INFO]${NC} %s\n" "$*"; }
warn() { printf "${YELLOW}[WARN]${NC} %s\n" "$*"; }
die()  { printf "${RED}[ERROR]${NC} %s\n" "$*" >&2; exit 1; }

REPO_URL="https://huggingface.co/AIMindLink/lambda-mindlink-memotron"

# ── 1. Detect package manager ──────────────────────────────────────────
if command -v apt-get >/dev/null 2>&1; then
    PKGMGR="apt"
elif command -v dnf >/dev/null 2>&1; then
    PKGMGR="dnf"
elif command -v yum >/dev/null 2>&1; then
    PKGMGR="yum"
else
    die "Neither apt nor dnf/yum found — this installer supports Debian/Ubuntu/Mint and Fedora/RHEL."
fi
info "Detected package manager: $PKGMGR"

# ── 2. Choose backend ──────────────────────────────────────────────────
if [[ $# -ge 1 ]]; then
    BACKEND="${1,,}"
else
    # Suggest a default from detected hardware
    SUGGEST=3
    if command -v nvidia-smi >/dev/null 2>&1; then
        SUGGEST=1
    elif command -v rocminfo >/dev/null 2>&1; then
        SUGGEST=2
    elif command -v lspci >/dev/null 2>&1 && lspci | grep -iE 'vga|3d|display' | grep -qi amd; then
        SUGGEST=2
    fi
    echo
    echo "Select llama-cpp-python backend:"
    echo "  1) CUDA  — NVIDIA GPU"
    echo "  2) ROCm  — AMD Ryzen iGPU / AMD dGPU"
    echo "  3) CPU   — no GPU (slow)"
    read -rp "Choice [1-3, default $SUGGEST]: " BACKEND
    BACKEND="${BACKEND:-$SUGGEST}"
fi
BACKEND="${BACKEND,,}"
case "$BACKEND" in
    1|cuda) BACKEND="cuda" ;;
    2|rocm) BACKEND="rocm" ;;
    3|cpu)  BACKEND="cpu"  ;;
    *) die "Unknown backend '$BACKEND' (use: cuda, rocm or cpu)" ;;
esac
info "Backend: $BACKEND"

# ── 3. System build tools ──────────────────────────────────────────────
case "$PKGMGR" in
    apt)
        info "Installing build tools (build-essential, cmake, python3-dev, python3-venv, git)..."
        sudo apt-get update
        sudo apt-get install -y build-essential cmake python3-dev python3-venv git
        ;;
    dnf|yum)
        info "Installing build tools (gcc-c++, cmake, python3-devel, git)..."
        sudo "$PKGMGR" install -y gcc-c++ cmake python3 python3-devel git
        ;;
esac

# ── 4. Python version check (project targets 3.11 / 3.12) ─────────────
if ! python3 -c 'import sys; sys.exit(0 if sys.version_info[:2] in ((3, 11), (3, 12)) else 1)'; then
    warn "Found $(python3 --version 2>&1) — the project targets 3.11/3.12. Continuing anyway."
fi

# ── 5. Get the repo (skip clone if script is run from inside the repo) ─
if [[ ! -f main.py ]]; then
    info "main.py not found here — cloning repository..."
    git clone "$REPO_URL"
    cd lambda-mindlink-memotron
fi

# ── 6. Virtual environment ────────────────────────────────────────────
if [[ ! -d .venv ]]; then
    info "Creating virtual environment..."
    python3 -m venv .venv
fi
source .venv/bin/activate
info "Upgrading pip..."
pip install --upgrade pip

# ── 7. llama-cpp-python for the chosen backend ─────────────────────────
case "$BACKEND" in
    cuda)
        command -v nvcc >/dev/null 2>&1 || \
            warn "nvcc not found — install the NVIDIA CUDA Toolkit 12.x first, or this build will fail."
        info "Installing llama-cpp-python with CUDA support..."
        CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python \
            --upgrade --force-reinstall --no-cache-dir
        ;;
    rocm)
        if [[ "$PKGMGR" == "apt" ]]; then
            read -rp "Also install ROCm dev packages via apt (hipblas-dev rocblas-dev)? [y/N] " yn
            if [[ "${yn,,}" == "y" || "${yn,,}" == "yes" ]]; then
                sudo apt-get install -y hipblas-dev rocblas-dev
            fi
        fi
        info "Installing llama-cpp-python with ROCm support..."
        CMAKE_ARGS="-DGGML_HIPBLAS=on" pip install llama-cpp-python \
            --upgrade --force-reinstall --no-cache-dir
        ;;
    cpu)
        info "Installing llama-cpp-python (CPU only)..."
        pip install llama-cpp-python --upgrade --force-reinstall --no-cache-dir
        ;;
esac

# ── 8. Remaining dependencies + ai folder ──────────────────────────────
info "Installing requirements.txt..."
pip install -r requirements.txt
mkdir -p ai

# ── 9. Done ────────────────────────────────────────────────────────────
echo
info "Installation complete."
echo "Next steps:"
echo "  1. Copy your .gguf files into: $(pwd)/ai/"
echo "  2. Edit config.py (_ALPHA_INTELLIGENCE_TO_LOAD, _MODEL_TYPE)"
echo "  3. Start the app:"
echo "        source .venv/bin/activate"
echo "        python main.py"