#!/usr/bin/env bash
set -euo pipefail
export DEBIAN_FRONTEND=noninteractive

#===============================================================================
# Ubuntu 22.04 MLX 및 mlx-lm 환경 자동 설정 스크립트 (최종)
#===============================================================================

# 1. 시스템 패키지 업데이트
sudo apt update

# 2. 필수 도구 설치 (curl 포함)
# apt install은 idempotent하므로 바로 실행
sudo apt install -y \\
  apt-transport-https \\
  ca-certificates \\
  gnupg \\
  software-properties-common \\
  wget \\
  curl

# 3. CMake 3.25+ 설치 (Kitware 레포지토리)
if ! command -v cmake &> /dev/null || [[ "$(cmake --version | head -n1)" != *"3.25"* && "$(cmake --version | head -n1)" != *"4."* ]]; then
  sudo rm -f /usr/share/keyrings/kitware-archive-keyring.gpg
  curl -fsSL <https://apt.kitware.com/keys/kitware-archive-latest.asc> \\
    | sudo gpg --batch --yes --dearmor -o /usr/share/keyrings/kitware-archive-keyring.gpg
  echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] <https://apt.kitware.com/ubuntu> jammy main' \\
    | sudo tee /etc/apt/sources.list.d/kitware.list > /dev/null
  sudo apt update
  sudo apt install -y cmake
else
  echo "[SKIP] cmake >=3.25 is already installed"
fi

# 4. 빌드 의존성 설치 (BLAS/LAPACK 포함)
sudo apt install -y \\
  build-essential \\
  git \\
  python3-dev \\
  python3-venv \\
  libomp-dev \\
  libopenblas-dev \\
  liblapack-dev \\
  liblapacke-dev

# 5. Python 가상환경 생성 및 활성화
VENV_DIR="$HOME/mlx-venv"
if [ -d "$VENV_DIR" ]; then
  echo "[SKIP] Virtual environment exists: $VENV_DIR"
else
  python3 -m venv "$VENV_DIR"
fi
# shellcheck disable=SC1090
source "$VENV_DIR/bin/activate"

# 6. pip 최신화
pip install --upgrade pip setuptools wheel

# 7. 환경 변수 설정 (Hugging Face 병렬 다운로드)
export HF_HUB_ENABLE_PARALLEL_DOWNLOAD=1
export HF_HUB_DOWNLOAD_NUM_THREADS=16
export HF_HUB_MAX_RETRIES=5
export HF_HOME="$HOME/.cache/huggingface"
export HF_HUB_CACHE="$HOME/.cache/huggingface"
mkdir -p "$HF_HOME"

# 8. MLX 소스코드 클론 또는 업데이트
WORKDIR="$HOME/mlx"
if [ -d "$WORKDIR" ]; then
  echo "[SKIP] MLX folder exists — pulling latest changes..."
  cd "$WORKDIR"
  git pull origin main
else
  git clone <https://github.com/ml-explore/mlx.git> "$WORKDIR"
  cd "$WORKDIR"
fi

# 9. C++ 코어 빌드 (이미 빌드된 경우 스킵)
BUILD_DIR="$WORKDIR/build"
if [ -f "$BUILD_DIR/libmlx.so" ] || [ -f "$BUILD_DIR/libmlx.dylib" ]; then
  echo "[SKIP] MLX C++ core already built in $BUILD_DIR"
else
  mkdir -p "$BUILD_DIR" && cd "$BUILD_DIR"
  cmake .. -DCMAKE_BUILD_TYPE=Release
  make -j"$(nproc)"
fi

# 10. (선택) 시스템 라이브러리로 설치 (선택) 시스템 라이브러리로 설치
sudo make install || echo "[WARN] make install failed or skipped"

# 11. Python 바인딩 설치
cd "$WORKDIR"
pip install -e .

# 12. mlx-lm 소스코드 클론 및 빌드 (CLI 및 라이브러리)
MLXLM_DIR="$HOME/mlx-lm"
if [ -d "$MLXLM_DIR" ]; then
  echo "[SKIP] mlx-lm source exists — pulling latest changes..."
  cd "$MLXLM_DIR"
  git pull origin main
else
  git clone <https://github.com/ml-explore/mlx-lm.git> "$MLXLM_DIR"
  cd "$MLXLM_DIR"
fi
# pip editable install for mlx-lm
pip install -e .

# 13. 추가 의존성 설치 (QLoRA/파인튜닝용) mlx-lm 설치 (CLI 및 라이브러리)
if pip show mlx-lm &> /dev/null; then
  echo "[SKIP] mlx-lm is already installed"
else
  pip install mlx-lm
fi

# 13. 추가 의존성 설치 (QLoRA/파인튜닝용)
REQUIREMENTS=(torch transformers bitsandbytes accelerate)
for pkg in "${REQUIREMENTS[@]}"; do
  if pip show "$pkg" &> /dev/null; then
    echo "[SKIP] $pkg is already installed"
  else
    pip install "$pkg"
  fi
done

# 14. 동작 확인
python3 - <<EOF
import mlx.core
import mlx_lm.convert
from huggingface_hub import snapshot_download
print("[OK] MLX core & Python bindings 정상 작동")
print("Hugging Face cache at: $HF_HOME")
EOF

# 완료 메시지
echo "[SUCCESS] Ubuntu 22.04에 MLX 및 mlx-lm 설치 완료. 활성화된 venv: $VENV_DIR"