#!/bin/bash

# GPU 학습용 Ubuntu 22.04 Docker 서버 세팅 스크립트
# 실행 방법: sudo bash setup_gpu_docker.sh

set -e

# 색상 설정
PURPLE='\033[0;35m'
NC='\033[0m' # No Color


echo "=== GPU 학습용 Ubuntu 22.04 Docker 서버 세팅 시작 ==="

# 시스템 업데이트
echo "시스템 업데이트 중..."
apt update && apt upgrade -y

# 기본 패키지 설치
echo "기본 패키지 설치 중..."
apt install -y \
    curl \
    wget \
    vim \
    git \
    htop \
    tree \
    unzip \
    software-properties-common \
    apt-transport-https \
    ca-certificates \
    gnupg \
    lsb-release

# NVIDIA 드라이버 설치
echo "NVIDIA 드라이버 설치 중..."
# 기존 NVIDIA 드라이버 제거
apt remove --purge -y nvidia*

# NVIDIA 드라이버 저장소 추가
add-apt-repository ppa:graphics-drivers/ppa -y
apt update

# 권장 드라이버 설치 (자동으로 최신 stable 버전 설치)
ubuntu-drivers autoinstall

# Docker 설치
echo "Docker 설치 중..."
# 기존 Docker 제거
apt remove -y docker docker-engine docker.io containerd runc

# Docker 공식 설치 스크립트 사용
wget -O get-docker.sh https://get.docker.com 
sh ./get-docker.sh

# Docker 서비스 시작 및 부팅시 자동 시작 설정
systemctl start docker
systemctl enable docker

# 현재 사용자를 docker 그룹에 추가
if [ "$SUDO_USER" ]; then
    usermod -aG docker $SUDO_USER
    echo "사용자 $SUDO_USER를 docker 그룹에 추가했습니다."
else
    echo "경고: SUDO_USER가 설정되지 않음. 수동으로 사용자를 docker 그룹에 추가하세요: sudo usermod -aG docker \$USER"
fi

# NVIDIA Container Toolkit 설치
echo "NVIDIA Container Toolkit 설치 중..."
# NVIDIA Container Toolkit 저장소 설정
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# NVIDIA Container Toolkit 설치
apt update
apt install -y nvidia-container-toolkit

# Docker에 NVIDIA 런타임 설정
echo "Docker에 NVIDIA 런타임 설정 중..."
nvidia-ctk runtime configure --runtime=docker

# Docker 데몬 재시작
systemctl restart docker

# 시스템 재부팅 후 실행할 검증 스크립트 생성
cat > /home/$SUDO_USER/verify_gpu_setup.sh << 'EOF'
#!/bin/bash

echo "=== GPU 및 Docker 설정 검증 ==="

echo "1. NVIDIA 드라이버 확인:"
nvidia-smi

echo -e "\n2. Docker 버전 확인:"
docker --version

echo -e "\n3. Docker Compose 확인:"
docker compose version

echo -e "\n4. NVIDIA Container Toolkit 테스트:"
docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

echo -e "\n5. 간단한 GPU 테스트 (PyTorch):"
docker run --rm --gpus all pytorch/pytorch:latest python -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA version: {torch.version.cuda}')
    print(f'GPU device count: {torch.cuda.device_count()}')
    print(f'GPU device name: {torch.cuda.get_device_name(0)}')
    # 간단한 GPU 연산 테스트
    x = torch.rand(1000, 1000).cuda()
    y = torch.rand(1000, 1000).cuda()
    z = torch.mm(x, y)
    print('GPU 연산 테스트 성공!')
else:
    print('GPU를 사용할 수 없습니다.')
"

echo -e "\n=== 검증 완료 ==="
EOF

chmod +x /home/$SUDO_USER/verify_gpu_setup.sh
chown $SUDO_USER:$SUDO_USER /home/$SUDO_USER/verify_gpu_setup.sh

# 유용한 Docker Compose 예제 생성
cat > /home/$SUDO_USER/docker-compose-gpu-example.yml << 'EOF'
version: '3.8'

services:
  jupyter-gpu:
    image: tensorflow/tensorflow:latest-gpu-jupyter
    ports:
      - "8888:8888"
    volumes:
      - ./notebooks:/tf/notebooks
      - ./data:/tf/data
    environment:
      - JUPYTER_ENABLE_LAB=yes
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    command: >
      bash -c "
        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 &&
        jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root --notebook-dir=/tf
      "

  pytorch-gpu:
    image: pytorch/pytorch:latest
    volumes:
      - ./workspace:/workspace
    working_dir: /workspace
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    command: python -c "import torch; print('PyTorch GPU available:', torch.cuda.is_available())"
EOF

chown $SUDO_USER:$SUDO_USER /home/$SUDO_USER/docker-compose-gpu-example.yml

# 디렉토리 생성
mkdir -p /home/$SUDO_USER/notebooks
mkdir -p /home/$SUDO_USER/data
mkdir -p /home/$SUDO_USER/workspace
chown -R $SUDO_USER:$SUDO_USER /home/$SUDO_USER/notebooks
chown -R $SUDO_USER:$SUDO_USER /home/$SUDO_USER/data
chown -R $SUDO_USER:$SUDO_USER /home/$SUDO_USER/workspace

echo "=== 설치 완료 ==="
echo ""
echo "다음 단계를 진행하세요:"
echo "1. 시스템을 재부팅하세요: sudo reboot"
echo "2. 재부팅 후 검증 스크립트를 실행하세요: ./verify_gpu_setup.sh"
echo "3. Docker Compose 예제를 사용하여 GPU 학습 환경을 시작하세요:"
echo "   cd ~ && docker compose -f docker-compose-gpu-example.yml up -d jupyter-gpu"
echo ""
echo "백업 및 복구 기능:"
echo "4. 현재 환경 전체 백업: ./backup_gpu_environment.sh"
echo "5. 빠른 백업 (설정만): ./quick_backup.sh"
echo "6. 다른 서버에 복구: sudo bash restore_gpu_environment.sh"
echo ""
echo "주의사항:"
echo "- 재부팅 후 새로운 터미널에서 docker 명령어를 사용하세요"
echo "- GPU 메모리 사용량을 모니터링하려면 'watch -n 1 nvidia-smi' 명령어를 사용하세요"
echo "- Jupyter Lab 접속: http://localhost:8888 (토큰은 docker logs로 확인)"
echo "- 백업은 정기적으로 실행하여 환경을 보존하세요"

# 백업 및 복구 스크립트 생성
cat > /home/$SUDO_USER/backup_gpu_environment.sh << 'EOF'
#!/bin/bash

# GPU 서버 환경 백업 스크립트
# 실행 방법: ./backup_gpu_environment.sh

set -e

BACKUP_DIR="gpu_server_backup_$(date +%Y%m%d_%H%M%S)"
BACKUP_PATH="$HOME/$BACKUP_DIR"

echo "=== GPU 서버 환경 백업 시작 ==="
echo "백업 경로: $BACKUP_PATH"

mkdir -p "$BACKUP_PATH"

# 시스템 정보 백업
echo "1. 시스템 정보 백업 중..."
lsb_release -a > "$BACKUP_PATH/system_info.txt"
uname -a >> "$BACKUP_PATH/system_info.txt"
nvidia-smi > "$BACKUP_PATH/nvidia_info.txt"
docker --version > "$BACKUP_PATH/docker_info.txt"

# 설치된 패키지 목록 백업
echo "2. 설치된 패키지 목록 백업 중..."
dpkg -l > "$BACKUP_PATH/installed_packages.txt"
apt list --installed > "$BACKUP_PATH/apt_packages.txt"

# Docker 이미지 및 컨테이너 백업
echo "3. Docker 이미지 목록 백업 중..."
docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.ID}}\t{{.Size}}" > "$BACKUP_PATH/docker_images.txt"
docker ps -a --format "table {{.Names}}\t{{.Image}}\t{{.Status}}" > "$BACKUP_PATH/docker_containers.txt"

# Docker 이미지 실제 백업 (선택사항)
echo "4. 주요 Docker 이미지 백업 중..."
mkdir -p "$BACKUP_PATH/docker_images"
for image in $(docker images --format "{{.Repository}}:{{.Tag}}" | grep -E "(pytorch|tensorflow|jupyter|cuda)" | head -10); do
    echo "  백업 중: $image"
    safe_name=$(echo "$image" | sed 's/[^a-zA-Z0-9._-]/_/g')
    docker save "$image" | gzip > "$BACKUP_PATH/docker_images/${safe_name}.tar.gz"
done

# Python 환경 백업
echo "5. Python 환경 백업 중..."
if command -v python3 &> /dev/null; then
    python3 -m pip list > "$BACKUP_PATH/pip_packages.txt"
fi

if command -v conda &> /dev/null; then
    conda env list > "$BACKUP_PATH/conda_envs.txt"
    conda list > "$BACKUP_PATH/conda_packages.txt"
fi

# 설정 파일 백업
echo "6. 설정 파일 백업 중..."
mkdir -p "$BACKUP_PATH/configs"
cp -r ~/.bashrc "$BACKUP_PATH/configs/" 2>/dev/null || true
cp -r ~/.profile "$BACKUP_PATH/configs/" 2>/dev/null || true
cp -r ~/.ssh "$BACKUP_PATH/configs/" 2>/dev/null || true
cp -r ~/.gitconfig "$BACKUP_PATH/configs/" 2>/dev/null || true
cp -r /etc/docker/daemon.json "$BACKUP_PATH/configs/" 2>/dev/null || true

# 중요한 디렉토리 백업
echo "7. 중요한 디렉토리 백업 중..."
mkdir -p "$BACKUP_PATH/user_data"
cp -r ~/notebooks "$BACKUP_PATH/user_data/" 2>/dev/null || true
cp -r ~/workspace "$BACKUP_PATH/user_data/" 2>/dev/null || true
cp -r ~/data "$BACKUP_PATH/user_data/" 2>/dev/null || true
cp -r ~/scripts "$BACKUP_PATH/user_data/" 2>/dev/null || true

# 복구 스크립트 생성
cat > "$BACKUP_PATH/restore_gpu_environment.sh" << 'RESTORE_EOF'
#!/bin/bash

# GPU 서버 환경 복구 스크립트
# 실행 방법: sudo bash restore_gpu_environment.sh

set -e

BACKUP_DIR="$(dirname "$0")"
echo "=== GPU 서버 환경 복구 시작 ==="
echo "복구 소스: $BACKUP_DIR"

# 패키지 복구
echo "1. 패키지 복구 중..."
if [ -f "$BACKUP_DIR/apt_packages.txt" ]; then
    # 기본 패키지 설치
    apt update
    grep -E "^ii" "$BACKUP_DIR/installed_packages.txt" | awk '{print $2}' | xargs apt install -y --ignore-missing
fi

# Docker 이미지 복구
echo "2. Docker 이미지 복구 중..."
if [ -d "$BACKUP_DIR/docker_images" ]; then
    for image_file in "$BACKUP_DIR/docker_images"/*.tar.gz; do
        if [ -f "$image_file" ]; then
            echo "  복구 중: $image_file"
            gunzip -c "$image_file" | docker load
        fi
    done
fi

# Python 패키지 복구
echo "3. Python 패키지 복구 중..."
if [ -f "$BACKUP_DIR/pip_packages.txt" ]; then
    pip3 install -r <(grep -v "^#" "$BACKUP_DIR/pip_packages.txt" | awk '{print $1}') --ignore-installed
fi

# 설정 파일 복구
echo "4. 설정 파일 복구 중..."
if [ -d "$BACKUP_DIR/configs" ]; then
    cp -r "$BACKUP_DIR/configs/.bashrc" ~/ 2>/dev/null || true
    cp -r "$BACKUP_DIR/configs/.profile" ~/ 2>/dev/null || true
    cp -r "$BACKUP_DIR/configs/.gitconfig" ~/ 2>/dev/null || true
    cp -r "$BACKUP_DIR/configs/daemon.json" /etc/docker/ 2>/dev/null || true
fi

# 사용자 데이터 복구
echo "5. 사용자 데이터 복구 중..."
if [ -d "$BACKUP_DIR/user_data" ]; then
    cp -r "$BACKUP_DIR/user_data"/* ~/ 2>/dev/null || true
fi

# 권한 수정
chown -R $SUDO_USER:$SUDO_USER /home/$SUDO_USER/

echo "=== 복구 완료 ==="
echo "Docker 재시작 중..."
systemctl restart docker

echo "복구가 완료되었습니다!"
echo "시스템을 재부팅하는 것을 권장합니다: sudo reboot"
RESTORE_EOF

chmod +x "$BACKUP_PATH/restore_gpu_environment.sh"

# 백업 정보 파일 생성
cat > "$BACKUP_PATH/backup_info.txt" << 'INFO_EOF'
=== GPU 서버 백업 정보 ===

백업 날짜: $(date)
백업 경로: $BACKUP_PATH
백업 내용:
- 시스템 정보 (system_info.txt, nvidia_info.txt, docker_info.txt)
- 설치된 패키지 목록 (installed_packages.txt, apt_packages.txt)
- Docker 이미지 및 컨테이너 정보
- 주요 Docker 이미지 파일 (docker_images/)
- Python 패키지 목록 (pip_packages.txt)
- 설정 파일 (configs/)
- 사용자 데이터 (user_data/)

복구 방법:
1. 백업 폴더를 새 서버로 복사
2. sudo bash restore_gpu_environment.sh 실행
3. 시스템 재부팅

주의사항:
- 복구 전 새 서버에 기본 GPU 드라이버와 Docker가 설치되어 있어야 함
- 복구 과정에서 기존 설정이 덮어쓰여질 수 있음
- 중요한 데이터는 별도 백업 권장
INFO_EOF

echo "백업 완료!"
echo "백업 위치: $BACKUP_PATH"
echo "복구 방법: sudo bash $BACKUP_PATH/restore_gpu_environment.sh"

# 백업 압축 옵션
read -p "백업을 압축하시겠습니까? (y/n): " compress
if [[ $compress == "y" || $compress == "Y" ]]; then
    echo "백업 압축 중..."
    tar -czf "${BACKUP_PATH}.tar.gz" -C "$HOME" "$BACKUP_DIR"
    rm -rf "$BACKUP_PATH"
    echo "압축 완료: ${BACKUP_PATH}.tar.gz"
fi
EOF

chmod +x /home/$SUDO_USER/backup_gpu_environment.sh
chown $SUDO_USER:$SUDO_USER /home/$SUDO_USER/backup_gpu_environment.sh

# 빠른 백업 스크립트 생성 (Docker 이미지 제외)
cat > /home/$SUDO_USER/quick_backup.sh << 'EOF'
#!/bin/bash

# 빠른 백업 스크립트 (Docker 이미지 제외)
# 실행 방법: ./quick_backup.sh

set -e

BACKUP_DIR="quick_backup_$(date +%Y%m%d_%H%M%S)"
BACKUP_PATH="$HOME/$BACKUP_DIR"

echo "=== 빠른 백업 시작 ==="
mkdir -p "$BACKUP_PATH"

# 시스템 정보 및 패키지 목록만 백업
lsb_release -a > "$BACKUP_PATH/system_info.txt"
nvidia-smi > "$BACKUP_PATH/nvidia_info.txt"
docker --version > "$BACKUP_PATH/docker_info.txt"
dpkg -l > "$BACKUP_PATH/installed_packages.txt"
docker images --format "table {{.Repository}}\t{{.Tag}}" > "$BACKUP_PATH/docker_images_list.txt"

if command -v python3 &> /dev/null; then
    python3 -m pip list > "$BACKUP_PATH/pip_packages.txt"
fi

# 설정 파일 백업
mkdir -p "$BACKUP_PATH/configs"
cp ~/.bashrc "$BACKUP_PATH/configs/" 2>/dev/null || true
cp ~/.profile "$BACKUP_PATH/configs/" 2>/dev/null || true

# 압축
tar -czf "${BACKUP_PATH}.tar.gz" -C "$HOME" "$BACKUP_DIR"
rm -rf "$BACKUP_PATH"

echo "빠른 백업 완료: ${BACKUP_PATH}.tar.gz"
EOF

chmod +x /home/$SUDO_USER/quick_backup.sh
chown $SUDO_USER:$SUDO_USER /home/$SUDO_USER/quick_backup.sh

echo "설정이 완료되었습니다!"