#!/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

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 ""

#--------------------------------------------------------------------
# 6. 완료
#--------------------------------------------------------------------
echo "[6/6] 셋업 완료!"
echo ""
echo "============================================"
echo " 다음 단계: 서버를 시작하세요"
echo "============================================"
echo ""
echo "  # 텍스트 전용:"
echo "  source $VENV_DIR/bin/activate"
echo "  vllm-mlx serve $MODEL \\"
echo "    --port $PORT --max-tokens 8192 \\"
echo "    --default-temperature 0.7 --default-top-p 0.9"
echo ""
echo "  # 텍스트 + 이미지 (멀티모달):"
echo "  source $VENV_DIR/bin/activate"
echo "  vllm-mlx serve $MODEL \\"
echo "    --port $PORT --max-tokens 8192 \\"
echo "    --default-temperature 0.7 --default-top-p 0.9 --mllm"
echo ""
echo "  서버 시작 후 브라우저에서 http://localhost:${WEBUI_PORT} 접속"
echo "  (첫 접속 시 회원가입 → 첫 계정이 admin)"
echo ""
echo "  종료하려면: ./stop-mlx.sh"
echo "============================================"