#!/bin/bash
# KORMo MoE vLLM Support Installation Script
# This script automatically installs the necessary files for vLLM support

set -e

echo "=========================================="
echo "KORMo MoE vLLM Support Installer"
echo "=========================================="
echo ""

# 1. vLLM 설치 확인
echo "Step 1: Checking vLLM installation..."
if ! python3 -c "import vllm" 2>/dev/null; then
    echo "❌ vLLM is not installed."
    echo "Installing vLLM..."
    pip install vllm
    echo "✅ vLLM installed successfully"
else
    echo "✅ vLLM is already installed"
fi
echo ""

# 2. vLLM 설치 경로 찾기
echo "Step 2: Finding vLLM installation path..."
VLLM_PATH=$(python3 -c "import vllm; import os; print(os.path.dirname(vllm.__file__))")
echo "vLLM path: $VLLM_PATH"
echo ""

# 3. kormo_moe_vllm.py 다운로드
echo "Step 3: Downloading KORMo MoE vLLM implementation..."
if [ -f "kormo_moe_vllm.py" ]; then
    echo "✅ kormo_moe_vllm.py found locally"
else
    echo "Downloading from HuggingFace (dev7halo/KORMo-10B-sft-moe)..."
    wget https://huggingface.co/dev7halo/KORMo-10B-sft-moe/resolve/main/kormo_moe_vllm.py -O kormo_moe_vllm.py 2>/dev/null || \
    curl -L https://huggingface.co/dev7halo/KORMo-10B-sft-moe/resolve/main/kormo_moe_vllm.py -o kormo_moe_vllm.py 2>/dev/null || {
        echo "❌ Failed to download. Please ensure kormo_moe_vllm.py is in the current directory."
        echo "You can manually download it from:"
        echo "https://huggingface.co/dev7halo/KORMo-10B-sft-moe/blob/main/kormo_moe_vllm.py"
        exit 1
    }
    echo "✅ Downloaded successfully"
fi
echo ""

# 4. 파일 복사
echo "Step 4: Installing KORMo MoE model file..."
TARGET_PATH="$VLLM_PATH/model_executor/models/kormo_moe.py"
cp kormo_moe_vllm.py "$TARGET_PATH"
echo "✅ Copied to $TARGET_PATH"
echo ""

# 5. 레지스트리 수정
echo "Step 5: Registering KORMo MoE in vLLM..."
REGISTRY_PATH="$VLLM_PATH/model_executor/models/registry.py"

# 이미 등록되어 있는지 확인
if grep -q "KORMoMoeForCausalLM" "$REGISTRY_PATH"; then
    echo "✅ KORMo MoE is already registered"
else
    echo "Adding KORMo MoE to registry..."

    # 백업 생성
    cp "$REGISTRY_PATH" "$REGISTRY_PATH.backup"
    echo "Created backup: $REGISTRY_PATH.backup"

    # JambaForCausalLM 다음에 KORMo MoE 추가
    if grep -q "JambaForCausalLM" "$REGISTRY_PATH"; then
        sed -i '/\"JambaForCausalLM\"/a\    \"KORMoMoeForCausalLM\": (\"kormo_moe\", \"KORMoMoeForCausalLM\"),' "$REGISTRY_PATH"
        echo "✅ KORMo MoE registered successfully"
    else
        echo "⚠️  Could not find JambaForCausalLM in registry."
        echo "Please manually add the following line to $REGISTRY_PATH in _TEXT_GENERATION_MODELS:"
        echo '    "KORMoMoeForCausalLM": ("kormo_moe", "KORMoMoeForCausalLM"),'
    fi
fi
echo ""

# 6. 설치 확인
echo "Step 6: Verifying installation..."
python3 << EOF
try:
    from vllm.model_executor.models.registry import ModelRegistry
    if "KORMoMoeForCausalLM" in ModelRegistry.get_supported_archs():
        print("✅ Installation successful! KORMo MoE is now supported in vLLM")
    else:
        print("❌ Registration verification failed")
        exit(1)
except Exception as e:
    print(f"❌ Error during verification: {e}")
    exit(1)
EOF

echo ""
echo "=========================================="
echo "Installation Complete!"
echo "=========================================="
echo ""
echo "You can now use KORMo MoE with vLLM:"
echo ""
echo "Example usage:"
echo ""
echo "from vllm import LLM, SamplingParams"
echo ""
echo "# Load the model"
echo "llm = LLM(model='dev7halo/KORMo-10B-sft-moe', dtype='float16')"
echo ""
echo "# Generate text"
echo "prompts = ['안녕하세요']"
echo "outputs = llm.generate(prompts, SamplingParams(temperature=0.8, max_tokens=100))"
echo ""