#!/bin/bash
set -e

echo "==================================="
echo "Audio Separator API - Test Script"
echo "==================================="

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

PASS=0
FAIL=0

test_pass() {
    echo -e "${GREEN}[PASS]${NC} $1"
    ((PASS++))
}

test_fail() {
    echo -e "${RED}[FAIL]${NC} $1"
    ((FAIL++))
}

test_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

cd "$(dirname "$0")"

# Test 1: Check Python environment
echo ""
echo "Test 1: Python environment"
if [ -f ".venv/bin/python" ]; then
    PYTHON_VERSION=$(.venv/bin/python --version 2>&1)
    test_pass "Python found: $PYTHON_VERSION"
else
    test_fail "Python virtual environment not found. Run ./install.sh first"
    exit 1
fi

# Test 2: Check core dependencies
echo ""
echo "Test 2: Core dependencies"

.venv/bin/python << 'EOF'
import sys

deps = [
    ("fastapi", "FastAPI"),
    ("uvicorn", "Uvicorn"),
    ("torch", "PyTorch"),
    ("audio_separator", "Audio Separator"),
    ("pydub", "PyDub"),
]

for module, name in deps:
    try:
        __import__(module)
        print(f"  [OK] {name}")
    except ImportError as e:
        print(f"  [FAIL] {name}: {e}")
        sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    test_pass "All core dependencies available"
else
    test_fail "Missing dependencies"
fi

# Test 3: Check CUDA availability
echo ""
echo "Test 3: CUDA / GPU availability"

CUDA_RESULT=$(.venv/bin/python << 'EOF'
import torch
if torch.cuda.is_available():
    device_name = torch.cuda.get_device_name(0)
    memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    print(f"CUDA_AVAILABLE|{device_name}|{memory_gb:.1f}")
else:
    print("CUDA_NOT_AVAILABLE")
EOF
)

if [[ "$CUDA_RESULT" == CUDA_AVAILABLE* ]]; then
    IFS='|' read -r _ GPU_NAME GPU_MEM <<< "$CUDA_RESULT"
    test_pass "CUDA available: $GPU_NAME (${GPU_MEM}GB VRAM)"
else
    test_warn "CUDA not available - will use CPU (slower)"
fi

# Test 4: Check model access
echo ""
echo "Test 4: Model loading"

MODEL_RESULT=$(.venv/bin/python << 'EOF'
import os
os.makedirs('/tmp/audio-separator/models', exist_ok=True)

try:
    from audio_separator.separator import Separator
    s = Separator(model_file_dir='/tmp/audio-separator/models')
    s.load_model()
    print("MODEL_OK")
except Exception as e:
    print(f"MODEL_FAIL|{e}")
EOF
)

if [[ "$MODEL_RESULT" == "MODEL_OK" ]]; then
    test_pass "Model loads successfully"
else
    test_fail "Model loading failed: ${MODEL_RESULT#MODEL_FAIL|}"
fi

# Test 5: Check FFmpeg
echo ""
echo "Test 5: FFmpeg availability"

if command -v ffmpeg &> /dev/null; then
    FFMPEG_VERSION=$(ffmpeg -version 2>&1 | head -n1)
    test_pass "FFmpeg found: $FFMPEG_VERSION"
else
    test_fail "FFmpeg not found - required for audio processing"
fi

# Test 6: Test API endpoints (if server is running)
echo ""
echo "Test 6: API endpoints"

API_URL="http://localhost:8000"

if curl -s --connect-timeout 2 "$API_URL/health" > /dev/null 2>&1; then
    # Health endpoint
    HEALTH=$(curl -s "$API_URL/health")
    if echo "$HEALTH" | grep -q '"status":"healthy"'; then
        test_pass "Health endpoint responding"
    else
        test_fail "Health endpoint unhealthy"
    fi

    # Models endpoint
    MODELS=$(curl -s "$API_URL/models")
    if echo "$MODELS" | grep -q '"models"'; then
        test_pass "Models endpoint responding"
    else
        test_fail "Models endpoint failed"
    fi

    # Docs endpoint
    if curl -s --connect-timeout 2 "$API_URL/docs" | grep -q "swagger"; then
        test_pass "API docs available at $API_URL/docs"
    else
        test_warn "API docs may not be accessible"
    fi
else
    test_warn "API server not running - skipping endpoint tests"
    echo "        Start the server with: .venv/bin/uvicorn app:app --host 0.0.0.0 --port 8000"
fi

# Test 7: Disk space check
echo ""
echo "Test 7: Disk space"

AVAILABLE_GB=$(df -BG /tmp | tail -1 | awk '{print $4}' | tr -d 'G')
if [ "$AVAILABLE_GB" -gt 10 ]; then
    test_pass "Sufficient disk space: ${AVAILABLE_GB}GB available in /tmp"
else
    test_warn "Low disk space: ${AVAILABLE_GB}GB available (recommend 10GB+)"
fi

# Summary
echo ""
echo "==================================="
echo "Test Summary"
echo "==================================="
echo -e "Passed: ${GREEN}$PASS${NC}"
echo -e "Failed: ${RED}$FAIL${NC}"
echo ""

if [ $FAIL -gt 0 ]; then
    echo -e "${RED}Some tests failed. Please fix issues before running the API.${NC}"
    exit 1
else
    echo -e "${GREEN}All critical tests passed!${NC}"
    exit 0
fi