#!/bin/bash
set -e

#====================================================================
# Qwen3.5 + Open WebUI (vllm-mlx) 원클릭 셋업
# 환경: Apple Silicon Mac (M1/M2/M3/M4) / Docker Desktop / Python 3.10+
#====================================================================

VENV_DIR="$HOME/mlx-env"
PROJECT_DIR="$(cd "$(dirname "$0")" && pwd)"
MODEL="mlx-community/Qwen3.5-35B-A3B-4bit"
PORT=8090
WEBUI_PORT=3000
MAX_TOKENS=8192
MAX_INPUT_TOKENS=8192
TEMPERATURE=0.7
TOP_P=0.9

echo "============================================"
echo " Qwen3.5 + Open WebUI (vllm-mlx) 셋업"
echo "============================================"
echo ""

#--------------------------------------------------------------------
# 1. 사전 요구사항 확인
#--------------------------------------------------------------------
echo "[1/6] 사전 요구사항 확인..."

# Python
if ! command -v python3 &>/dev/null; then
    echo "❌ python3가 설치되어 있지 않습니다."
    exit 1
fi
PYTHON_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
PYTHON_MAJOR=$(echo "$PYTHON_VERSION" | cut -d. -f1)
PYTHON_MINOR=$(echo "$PYTHON_VERSION" | cut -d. -f2)
if [ "$PYTHON_MAJOR" -lt 3 ] || { [ "$PYTHON_MAJOR" -eq 3 ] && [ "$PYTHON_MINOR" -lt 10 ]; }; then
    echo "❌ Python 3.10 이상이 필요합니다. (현재: $PYTHON_VERSION)"
    exit 1
fi
echo "  ✓ Python $PYTHON_VERSION"

# Docker
if ! command -v docker &>/dev/null; then
    echo "❌ Docker가 설치되어 있지 않습니다."
    exit 1
fi
if ! docker info &>/dev/null; then
    echo "❌ Docker Desktop이 실행 중이 아닙니다."
    exit 1
fi
echo "  ✓ Docker"

# Apple Silicon
if ! sysctl -n machdep.cpu.brand_string 2>/dev/null | grep -q "Apple"; then
    echo "❌ Apple Silicon이 아닙니다. MLX는 Apple Silicon에서만 동작합니다."
    exit 1
fi
echo "  ✓ Apple Silicon"

# RAM
TOTAL_RAM_GB=$(sysctl -n hw.memsize | awk '{printf "%.0f", $1/1024/1024/1024}')
echo "  ✓ RAM: ${TOTAL_RAM_GB}GB"
if [ "$TOTAL_RAM_GB" -lt 32 ]; then
    echo "  ⚠️  RAM이 32GB 미만입니다. 4bit 모델(~20GB)도 빡빡할 수 있습니다."
fi

echo ""

#--------------------------------------------------------------------
# 2. 가상환경 생성 및 패키지 설치
#--------------------------------------------------------------------
echo "[2/6] 가상환경 및 패키지 설치..."

# conda 감지 경고
if [ -n "$CONDA_DEFAULT_ENV" ]; then
    echo "  ⚠️  conda 환경 감지 ($CONDA_DEFAULT_ENV)."
    echo "     MPICH 충돌 방지를 위해 별도 venv를 사용합니다."
fi

if [ ! -d "$VENV_DIR" ]; then
    python3 -m venv "$VENV_DIR"
    echo "  ✓ 가상환경 생성: $VENV_DIR"
else
    echo "  ✓ 기존 가상환경 사용: $VENV_DIR"
fi

source "$VENV_DIR/bin/activate"

# vllm-mlx
if ! pip show vllm-mlx &>/dev/null; then
    echo "  vllm-mlx 설치 중... (시간이 걸릴 수 있습니다)"
    pip install -q git+https://github.com/waybarrios/vllm-mlx.git
    echo "  ✓ vllm-mlx 설치 완료"
else
    echo "  ✓ vllm-mlx 이미 설치됨"
fi

# torch, torchvision
if ! pip show torch &>/dev/null; then
    echo "  torch, torchvision 설치 중..."
    pip install -q torch torchvision
    echo "  ✓ torch 설치 완료"
else
    echo "  ✓ torch 이미 설치됨"
fi

echo ""

#--------------------------------------------------------------------
# 3. transformers fast image processor 패치
#    원인: transformers >= 5.x의 fast image processor가 PyTorch 텐서만
#          반환하지만, mlx_vlm은 numpy/MLX 배열을 기대하여 충돌 발생.
#    해결: fast 클래스를 slow 클래스로 리다이렉트.
#--------------------------------------------------------------------
echo "[3/6] 이미지 프로세서 호환성 패치..."

FAST_FILE=$("$VENV_DIR/bin/python3" -c "
import transformers, os
print(os.path.join(os.path.dirname(transformers.__file__),
      'models/qwen2_vl/image_processing_qwen2_vl_fast.py'))
")

if grep -q "mlx_vlm 호환" "$FAST_FILE" 2>/dev/null; then
    echo "  ✓ 이미 패치 적용됨"
else
    cp "$FAST_FILE" "${FAST_FILE}.bak"

    cat > "$FAST_FILE" << 'PATCH'
"""
Fast Image processor class for Qwen2-VL.
Patched: mlx_vlm 호환을 위해 slow 버전으로 폴백합니다.
원본: image_processing_qwen2_vl_fast.py.bak

배경:
  transformers >= 5.x는 Fast Image Processor를 기본 로드합니다.
  Fast 버전은 PyTorch 텐서만 반환하지만, mlx_vlm은 numpy/MLX 배열을 기대합니다.
  이 패치는 fast 클래스를 slow 클래스로 대체하여 호환성을 확보합니다.

복원:
  cp image_processing_qwen2_vl_fast.py.bak image_processing_qwen2_vl_fast.py
"""

from .image_processing_qwen2_vl import Qwen2VLImageProcessor as Qwen2VLImageProcessorFast

__all__ = ["Qwen2VLImageProcessorFast"]
PATCH

    echo "  ✓ 패치 적용 완료 (원본 백업됨)"
fi

echo ""

#--------------------------------------------------------------------
# 4. Docker Compose 설정
#--------------------------------------------------------------------
echo "[4/6] Docker Compose 설정..."

if [ ! -f "$PROJECT_DIR/docker-compose.mlx.yml" ]; then
    cat > "$PROJECT_DIR/docker-compose.mlx.yml" << EOF
services:
  open-webui:
    image: ghcr.io/open-webui/open-webui:main
    container_name: open-webui-mlx
    ports:
      - "${WEBUI_PORT}:8080"
    environment:
      - OPENAI_API_BASE_URL=http://host.docker.internal:${PORT}/v1
      - OPENAI_API_KEY=none
      - OLLAMA_BASE_URL=
    volumes:
      - open-webui-mlx-data:/app/backend/data
    extra_hosts:
      - "host.docker.internal:host-gateway"
    restart: unless-stopped

volumes:
  open-webui-mlx-data:
EOF
    echo "  ✓ docker-compose.mlx.yml 생성"
else
    echo "  ✓ docker-compose.mlx.yml 이미 존재"
fi

echo ""

#--------------------------------------------------------------------
# 5. Open WebUI 실행
#--------------------------------------------------------------------
echo "[5/6] Open WebUI 실행..."

cd "$PROJECT_DIR"
docker compose -f docker-compose.mlx.yml up -d 2>&1 | grep -v "^$"
echo "  ✓ Open WebUI 실행 중 (http://localhost:${WEBUI_PORT})"

echo ""

MLLM_FLAG="--mllm"

#--------------------------------------------------------------------
# 6. vllm-mlx 서버 시작
#--------------------------------------------------------------------
LOG_FILE="$PROJECT_DIR/vllm-mlx.log"

echo "[6/6] vllm-mlx 서버 시작 (백그라운드)..."
echo ""
echo "  모델: $MODEL"
echo "  포트: $PORT"
echo "  max_tokens (출력): $MAX_TOKENS"
echo "  max_input (입력): $MAX_INPUT_TOKENS"
echo "  temperature: $TEMPERATURE"
echo "  top_p: $TOP_P"
echo "  모드: 멀티모달 (이미지 + 텍스트)"
echo ""

# 백그라운드 실행, 로그는 파일로
MAX_CACHE_BLOCKS=$((MAX_INPUT_TOKENS / 64))

nohup vllm-mlx serve "$MODEL" \
  --port "$PORT" \
  --max-tokens "$MAX_TOKENS" \
  --default-temperature "$TEMPERATURE" \
  --default-top-p "$TOP_P" \
  --kv-cache-quantization \
  --use-paged-cache \
  --max-cache-blocks "$MAX_CACHE_BLOCKS" \
  --timeout 600 \
  $MLLM_FLAG > "$LOG_FILE" 2>&1 &

SERVER_PID=$!
echo "  서버 PID: $SERVER_PID"
echo "  로그 파일: $LOG_FILE"
echo ""

# 서버가 뜰 때까지 대기 (최대 120초, 첫 실행 시 모델 다운로드 포함)
echo "  서버 준비 대기 중... (첫 실행 시 모델 다운로드 ~20GB)"
for i in $(seq 1 120); do
    if curl -s http://localhost:$PORT/v1/models > /dev/null 2>&1; then
        echo ""
        echo "  ✓ 서버 준비 완료!"
        break
    fi
    # 프로세스가 죽었는지 확인
    if ! kill -0 $SERVER_PID 2>/dev/null; then
        echo ""
        echo "  ❌ 서버 시작 실패. 로그를 확인하세요:"
        echo "     tail -50 $LOG_FILE"
        exit 1
    fi
    printf "."
    sleep 1
done

# 120초 후에도 안 되면 안내
if ! curl -s http://localhost:$PORT/v1/models > /dev/null 2>&1; then
    echo ""
    echo "  ⚠️  서버가 아직 준비 중입니다. (모델 다운로드 중일 수 있음)"
    echo "     로그 확인: tail -f $LOG_FILE"
    echo "     준비되면 http://localhost:${WEBUI_PORT} 접속"
else
    echo ""
fi

echo ""
echo "============================================"
echo " 셋업 완료!"
echo "============================================"
echo ""
echo "  브라우저: http://localhost:${WEBUI_PORT}"
echo "  (첫 접속 시 회원가입 → 첫 계정이 admin)"
echo ""
echo "  로그 확인: tail -f $LOG_FILE"
echo "  종료:     ./stop-mlx.sh"
echo "============================================"