#!/bin/bash

# Jetson Xavier 전용 인페인팅 서버 설치 스크립트
# ARM64 아키텍처 + NVIDIA Jetson Xavier + 32GB 통합 메모리 최적화
# Usage: ./setup_jetson.sh

set -e

# 색상 코드
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 로그 함수들
log_info() {
    echo -e "${BLUE}[JETSON]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${CYAN}[STEP]${NC} $1"
}

# 스크립트 시작
clear
echo "======================================"
echo "🚀 Jetson Xavier 인페인팅 서버 설치"
echo "======================================"
echo ""

# 1단계: 시스템 확인
log_step "1단계: Jetson Xavier 시스템 확인"

# ARM64 아키텍처 확인
if [ "$(uname -m)" != "aarch64" ]; then
    log_error "이 스크립트는 Jetson Xavier (ARM64)용입니다. x86 시스템은 setup_x86.sh를 사용하세요."
    exit 1
fi

# Jetson 확인
if [ ! -f "/etc/nv_tegra_release" ]; then
    log_warning "Jetson 시스템이 아닌 것 같습니다. 계속 진행하시겠습니까? (y/N)"
    read -r response
    if [[ ! "$response" =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

log_success "Jetson Xavier ARM64 시스템 확인됨"

# 2단계: 가상환경 설정
log_step "2단계: Python 가상환경 설정"

# 프로젝트 루트로 이동
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
cd "$PROJECT_ROOT"

# 가상환경 생성 함수
create_venv() {
    local venv_path="$1"
    if [ ! -d "$venv_path" ]; then
        log_info "가상환경 생성 중: $venv_path"
        python3 -m venv "$venv_path"
        log_success "가상환경 생성 완료: $venv_path"
    else
        log_info "기존 가상환경 발견: $venv_path"
    fi
}

# venv 또는 현재 디렉토리에 가상환경 확인/생성
if [ -d "venv" ]; then
    VENV_PATH="venv"
    log_info "기존 venv 가상환경 사용"
elif [ -f "pyvenv.cfg" ]; then
    VENV_PATH="."
    log_info "프로젝트 자체가 가상환경으로 설정됨"
else
    VENV_PATH="venv"
    log_info "새로운 venv 가상환경 생성"
    create_venv "$VENV_PATH"
fi

# 가상환경 활성화
if [ "$VENV_PATH" = "." ]; then
    # 현재 디렉토리가 가상환경인 경우
    if [ -z "$VIRTUAL_ENV" ]; then
        log_info "프로젝트 디렉토리 가상환경 활성화"
        source bin/activate 2>/dev/null || {
            log_error "가상환경 활성화 실패. venv를 새로 생성합니다."
            VENV_PATH="venv"
            create_venv "$VENV_PATH"
            source "$VENV_PATH/bin/activate"
        }
    fi
else
    source "$VENV_PATH/bin/activate"
fi

log_success "가상환경 활성화: $VIRTUAL_ENV"

# 3단계: 시스템 패키지 업데이트
log_step "3단계: 시스템 패키지 업데이트"
sudo apt update && sudo apt upgrade -y

# 4단계: 필수 도구 설치
log_step "4단계: Jetson 개발 도구 설치"
sudo apt install -y \
    python3-dev \
    python3-pip \
    python3-venv \
    build-essential \
    cmake \
    git \
    curl \
    wget \
    unzip \
    pkg-config \
    libssl-dev \
    libffi-dev \
    libjpeg-dev \
    libpng-dev \
    libtiff-dev \
    libopencv-dev \
    nvidia-jetpack

# 5단계: GCC 11 설치 (ONNX Runtime GPU 호환성)
log_step "5단계: GCC 11 설치 (ONNX Runtime GPU 필수)"
sudo apt install -y gcc-11 g++-11
sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100
sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-11 100

# 6단계: Python 패키지 업그레이드
log_step "6단계: Python 패키지 매니저 업그레이드"
pip install --upgrade pip setuptools wheel

# 7단계: Jetson 특화 ONNX Runtime GPU 설치
log_step "7단계: Jetson ONNX Runtime GPU 설치"
log_info "Jetson Xavier용 ONNX Runtime GPU 다운로드 중..."

# 기존 onnxruntime 제거
pip uninstall -y onnxruntime onnxruntime-gpu 2>/dev/null || true

# Jetson Xavier 전용 ONNX Runtime GPU 설치
ONNX_WHEEL_URL="https://nvidia.box.com/shared/static/zostg6agm00fb6t5uisw51qi6kpcuwzd.whl"
ONNX_WHEEL_FILE="onnxruntime_gpu-1.17.0-cp38-cp38-linux_aarch64.whl"

wget -O "$ONNX_WHEEL_FILE" "$ONNX_WHEEL_URL"
pip install --force-reinstall "$ONNX_WHEEL_FILE"
rm -f "$ONNX_WHEEL_FILE"

log_success "Jetson ONNX Runtime GPU 설치 완료"

# 8단계: PyTorch 설치 (Jetson 최적화)
log_step "8단계: Jetson 최적화 PyTorch 설치"
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# 9단계: 프로젝트 의존성 설치
log_step "9단계: 프로젝트 의존성 설치"
pip install -r requirements.txt

# 10단계: numpy 버전 조정 (호환성)
log_step "10단계: 호환성 패키지 조정"
pip install numpy==1.23.5

# 11단계: 설치 검증
log_step "11단계: 설치 검증"

echo "=== Python 환경 ==="
python --version
echo "가상환경: $VIRTUAL_ENV"

echo -e "\n=== ONNX Runtime 확인 ==="
python -c "
import onnxruntime as ort
print(f'ONNX Runtime 버전: {ort.__version__}')
providers = ort.get_available_providers()
print(f'사용 가능한 providers: {providers}')
if 'CUDAExecutionProvider' in providers:
    print('✅ CUDA 지원 확인됨')
if 'TensorrtExecutionProvider' in providers:
    print('✅ TensorRT 지원 확인됨')
" || log_warning "ONNX Runtime 확인 실패"

echo -e "\n=== PyTorch CUDA 확인 ==="
python -c "
import torch
print(f'PyTorch 버전: {torch.__version__}')
print(f'CUDA 사용 가능: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA 버전: {torch.version.cuda}')
    print(f'GPU 개수: {torch.cuda.device_count()}')
" || log_warning "PyTorch 확인 실패"

echo -e "\n=== Jetson 시스템 정보 ==="
tegrastats --verbose 2>/dev/null | head -1 || echo "tegrastats 정보 없음"

# 완료
echo ""
echo "======================================"
log_success "🎉 Jetson Xavier 설치 완료!"
echo "======================================"
echo ""
echo "다음 명령으로 서버를 시작하세요:"
echo "  bash scripts/start_server.sh"
echo ""
echo "모니터링 대시보드:"
echo "  http://localhost:8009"
echo ""
echo "API 테스트:"
echo "  curl http://localhost:8008/api/v1/health"
echo ""