#!/bin/bash

# x86-64 시스템용 인페인팅 서버 설치 스크립트
# RTX 3060 12GB, RTX 4080 등 x86 GPU 최적화
# Usage: ./setup_x86.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}[X86]${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 "🖥️  x86-64 인페인팅 서버 설치"
echo "======================================"
echo ""

# 1단계: 시스템 확인
log_step "1단계: x86-64 시스템 확인"

# x86-64 아키텍처 확인
ARCH=$(uname -m)
if [[ "$ARCH" != "x86_64" && "$ARCH" != "amd64" ]]; then
    log_error "이 스크립트는 x86-64용입니다. Jetson Xavier는 setup_jetson.sh를 사용하세요."
    exit 1
fi

log_success "x86-64 시스템 확인됨"

# 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단계: 시스템 패키지 업데이트 (Ubuntu/Debian)
log_step "3단계: 시스템 패키지 업데이트"
if command -v apt >/dev/null 2>&1; then
    sudo apt update && sudo apt upgrade -y
    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
elif command -v yum >/dev/null 2>&1; then
    sudo yum update -y
    sudo yum groupinstall -y "Development Tools"
    sudo yum install -y python3-devel python3-pip cmake git curl wget unzip
else
    log_warning "패키지 매니저를 찾을 수 없습니다. 수동으로 개발 도구를 설치해주세요."
fi

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

# 5단계: x86 ONNX Runtime GPU 설치
log_step "5단계: x86 ONNX Runtime GPU 설치"
log_info "x86-64용 ONNX Runtime GPU 설치 중..."

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

# x86-64 ONNX Runtime GPU 설치
pip install onnxruntime-gpu

log_success "x86 ONNX Runtime GPU 설치 완료"

# 6단계: PyTorch 설치 (x86 CUDA)
log_step "6단계: x86 CUDA PyTorch 설치"
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

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

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

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()}')
    for i in range(torch.cuda.device_count()):
        print(f'GPU {i}: {torch.cuda.get_device_name(i)}')
" || log_warning "PyTorch 확인 실패"

echo -e "\n=== NVIDIA GPU 확인 ==="
nvidia-smi 2>/dev/null || log_warning "nvidia-smi 정보 없음"

# 완료
echo ""
echo "======================================"
log_success "🎉 x86-64 설치 완료!"
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"