#!/bin/bash
# ═══════════════════════════════════════════════════════════
#  Video Blur AI - Automated Setup Script
#  Handles NTFS/exFAT drives (no symlink support)
# ═══════════════════════════════════════════════════════════
set -e

echo "╔══════════════════════════════════════════════════╗"
echo "║       🎬 Video Blur AI - Setup                   ║"
echo "╚══════════════════════════════════════════════════╝"
echo ""

PROJECT_DIR="$(cd "$(dirname "$0")" && pwd)"
echo "📂 Project: $PROJECT_DIR"

# ── Check Python ──
echo "🐍 Checking Python..."
if ! command -v python3 &> /dev/null; then
    echo "❌ Python 3 not found! Install Python 3.10+"
    exit 1
fi
PYTHON_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
echo "   Python $PYTHON_VERSION ✅"

# ── Check FFmpeg ──
echo "📹 Checking FFmpeg..."
if ! command -v ffmpeg &> /dev/null; then
    echo "⚠️  FFmpeg not found. Installing..."
    if command -v apt &> /dev/null; then
        sudo apt update && sudo apt install -y ffmpeg
    elif command -v brew &> /dev/null; then
        brew install ffmpeg
    else
        echo "❌ Install FFmpeg manually: https://ffmpeg.org/download.html"
        exit 1
    fi
fi
echo "   FFmpeg ✅"

# ── Check Git ──
echo "📦 Checking Git..."
if ! command -v git &> /dev/null; then
    echo "⚠️  Git not found. Installing..."
    if command -v apt &> /dev/null; then
        sudo apt install -y git
    else
        echo "❌ Install Git: https://git-scm.com"
        exit 1
    fi
fi
echo "   Git ✅"

# ── Check CUDA ──
echo "🎮 Checking GPU..."
if command -v nvidia-smi &> /dev/null; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
    GPU_MEM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1)
    echo "   GPU: $GPU_NAME ($GPU_MEM) ✅"
else
    echo "   ⚠️  No NVIDIA GPU. Will use CPU (slow)."
fi

# ══════════════════════════════════════════════════════════
#  Detect filesystem — NTFS/exFAT can't create symlinks
# ══════════════════════════════════════════════════════════
CAN_SYMLINK=true
TEST_LINK="$PROJECT_DIR/.symtest_link_$$"
TEST_FILE="$PROJECT_DIR/.symtest_file_$$"
touch "$TEST_FILE" 2>/dev/null || true
if ! ln -s "$TEST_FILE" "$TEST_LINK" 2>/dev/null; then
    CAN_SYMLINK=false
fi
rm -f "$TEST_LINK" "$TEST_FILE" 2>/dev/null

VENV_DIR="$PROJECT_DIR/venv"

if [ "$CAN_SYMLINK" = false ]; then
    echo ""
    echo "════════════════════════════════════════════════════"
    echo "⚠️  NTFS/exFAT drive detected (no symlink support)"
    echo "   venv will be created in your home directory"
    echo "════════════════════════════════════════════════════"
    
    VENV_DIR="$HOME/.venvs/video-blur-ai"
    mkdir -p "$(dirname "$VENV_DIR")"
    
    # Create helper scripts
    cat > "$PROJECT_DIR/run.sh" << EOF
#!/bin/bash
source "$VENV_DIR/bin/activate"
cd "$PROJECT_DIR"
python app.py "\$@"
EOF
    chmod +x "$PROJECT_DIR/run.sh"

    cat > "$PROJECT_DIR/run_cli.sh" << EOF
#!/bin/bash
source "$VENV_DIR/bin/activate"
cd "$PROJECT_DIR"
python cli.py "\$@"
EOF
    chmod +x "$PROJECT_DIR/run_cli.sh"

    cat > "$PROJECT_DIR/activate.sh" << EOF
#!/bin/bash
source "$VENV_DIR/bin/activate"
cd "$PROJECT_DIR"
echo "✅ Activated. venv: $VENV_DIR"
EOF
    chmod +x "$PROJECT_DIR/activate.sh"
    
    echo "   Created run.sh, run_cli.sh, activate.sh ✅"
fi

echo "   venv location: $VENV_DIR"

# ── Create Virtual Environment ──
echo ""
echo "📦 Creating virtual environment..."
if [ ! -d "$VENV_DIR" ]; then
    python3 -m venv "$VENV_DIR"
    echo "   Created ✅"
else
    echo "   Already exists ✅"
fi
source "$VENV_DIR/bin/activate"
pip install --upgrade pip -q
echo "   Activated ✅"

# ── Install PyTorch ──
echo ""
echo "🔥 Installing PyTorch (this may take a while)..."
if command -v nvidia-smi &> /dev/null; then
    pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 -q
else
    pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu -q
fi
echo "   PyTorch ✅"

# ── Install Dependencies ──
echo ""
echo "📚 Installing dependencies..."
pip install -q \
    "opencv-python>=4.8.0" \
    "numpy>=1.24.0" \
    "Pillow>=10.0.0" \
    "ffmpeg-python>=0.2.0" \
    "gradio>=4.20.0" \
    "tqdm>=4.65.0" \
    "supervision>=0.18.0" \
    "huggingface-hub>=0.20.0" \
    "transformers>=4.35.0"
echo "   Dependencies ✅"

# ── Install SAM 2 ──
echo ""
echo "🎯 Installing SAM 2 from source (official method)..."

SAM2_DIR="$HOME/.sam2_source"
if [ -d "$SAM2_DIR" ] && python3 -c "import sam2" 2>/dev/null; then
    echo "   SAM 2 already installed ✅"
else
    rm -rf "$SAM2_DIR"
    git clone --depth 1 https://github.com/facebookresearch/sam2.git "$SAM2_DIR"
    cd "$SAM2_DIR"
    
    # Skip CUDA extension build (works fine without it)
    SAM2_BUILD_CUDA=0 pip install -e "." -q
    
    cd "$PROJECT_DIR"
    echo "   SAM 2 installed ✅"
fi

# ── Install Grounding DINO ──
# Grounding DINO is used via HuggingFace Transformers (AutoModelForZeroShotObjectDetection)
# No separate package needed! transformers handles it.
# If you need the native package as fallback, uncomment below:
echo ""
echo "🔍 Grounding DINO will load via HuggingFace Transformers ✅"
# pip install -q "groundingdino @ git+https://github.com/IDEA-Research/GroundingDINO.git"

# ── Download Models ──
echo ""
echo "⬇️  Downloading model weights (first time only, cached after)..."
python3 << 'PYEOF'
from huggingface_hub import hf_hub_download
import os

cache_dir = os.path.expanduser("~/.cache/huggingface")
print(f"   Cache dir: {cache_dir}")

print("   Downloading Grounding DINO Base (~999MB)...")
try:
    from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
    AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-base")
    AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-base")
    print("   Grounding DINO Base ✅")
except Exception as e:
    print(f"   ⚠️  Will download on first run: {e}")

print("   Downloading SAM 2.1 Base+ (~350MB)...")
try:
    hf_hub_download(repo_id="facebook/sam2.1-hiera-base-plus", filename="sam2.1_hiera_base_plus.pt")
    print("   SAM 2.1 ✅")
except Exception as e:
    print(f"   ⚠️  Will download on first run: {e}")

print("   Models are cached in ~/.cache/huggingface — next launch will be instant.")
PYEOF

# ── Verify ──
echo ""
echo "🔧 Verifying..."
python3 -c "
import torch, cv2, gradio, numpy
print(f'   PyTorch {torch.__version__}', '(CUDA)' if torch.cuda.is_available() else '(CPU)')
print(f'   OpenCV {cv2.__version__} | Gradio {gradio.__version__} | NumPy {numpy.__version__}')
try:
    import sam2
    print(f'   SAM 2 ✅')
except ImportError:
    print('   ⚠️  SAM 2 import failed - check installation')
try:
    from transformers import AutoModelForZeroShotObjectDetection
    print(f'   Grounding DINO (via transformers) ✅')
except ImportError:
    print('   ⚠️  Transformers import failed')
print('   All OK ✅')
"

# ── Done ──
echo ""
echo "╔══════════════════════════════════════════════════╗"
echo "║            ✅ Setup Complete!                     ║"
echo "╠══════════════════════════════════════════════════╣"

if [ "$CAN_SYMLINK" = false ]; then
echo "║                                                  ║"
echo "║  🚀 Start the app:                               ║"
echo "║     ./run.sh                                     ║"
echo "║                                                  ║"
echo "║  🖥️  CLI mode:                                    ║"
echo "║     ./run_cli.sh video.mp4 -p \"face.\"            ║"
echo "║                                                  ║"
echo "║  Or manually:                                    ║"
echo "║     source activate.sh                           ║"
echo "║     python app.py                                ║"
else
echo "║                                                  ║"
echo "║  🚀 Start the app:                               ║"
echo "║     source venv/bin/activate                     ║"
echo "║     python app.py                                ║"
fi

echo "║                                                  ║"
echo "║  🌐 Open: http://localhost:7860                  ║"
echo "║  📡 Public: python app.py --share                ║"
echo "║                                                  ║"
echo "╚══════════════════════════════════════════════════╝"