#!/bin/bash

# Face Recognition Test Script
# Tests the face recognition functionality

set -e

echo "========================================="
echo "Testing Face Recognition Functionality"
echo "========================================="

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

# Configuration
TEST_VIDEO="test_video.mp4"
TEST_IMAGE="test_face.jpg"
FACE_DATABASE="face_database.json"
OUTPUT_DIR="/tmp/face_recognition_test"
API_URL="http://localhost:3002"

# Create test directory
mkdir -p "$OUTPUT_DIR"

# Function to print status
print_status() {
	if [ $? -eq 0 ]; then
		echo -e "${GREEN}[PASS]${NC} $1"
	else
		echo -e "${RED}[FAIL]${NC} $1"
		exit 1
	fi
}

# Function to check if service is running
check_service() {
	echo -e "${YELLOW}[INFO]${NC} Checking if Momentry Core is running..."
	curl -s "$API_URL/health" >/dev/null
	if [ $? -ne 0 ]; then
		echo -e "${RED}[ERROR]${NC} Momentry Core is not running. Please start it first."
		echo "Run: cargo run --bin momentry -- server"
		exit 1
	fi
	print_status "Momentry Core is running"
}

# Function to test Python dependencies
test_python_deps() {
	echo -e "${YELLOW}[INFO]${NC} Testing Python dependencies..."

	# Check Python version
	python3 --version
	print_status "Python is available"

	# Check OpenCV
	python3 -c "import cv2; print(f'OpenCV version: {cv2.__version__}')"
	print_status "OpenCV is installed"

	# Check numpy
	python3 -c "import numpy; print(f'NumPy version: {numpy.__version__}')"
	print_status "NumPy is installed"

	# Check ONNX Runtime for MPS support
	python3 -c "try:
    import onnxruntime as ort
    providers = ort.get_available_providers()
    print(f'ONNX Runtime available providers: {providers}')
    if 'CoreMLExecutionProvider' in providers:
        print('✓ MPS/CoreML acceleration available')
    elif 'CUDAExecutionProvider' in providers:
        print('✓ CUDA acceleration available')
    else:
        print('⚠ Only CPU available')
except ImportError:
    print('ONNX Runtime not installed (required for MPS acceleration)')"

	# Check InsightFace (optional)
	python3 -c "try:
    import insightface
    print(f'InsightFace version: {insightface.__version__}')
except ImportError:
    print('InsightFace not installed (optional)')"
	echo -e "${YELLOW}[INFO]${NC} InsightFace check completed"
}

# Function to test face processor
test_face_processor() {
	echo -e "${YELLOW}[INFO]${NC} Testing basic face processor..."

	# Check if test video exists
	if [ ! -f "$TEST_VIDEO" ]; then
		echo -e "${YELLOW}[WARN]${NC} Test video not found, creating dummy test..."
		# Create a simple test video with ffmpeg if available
		if command -v ffmpeg &>/dev/null; then
			ffmpeg -f lavfi -i testsrc=duration=5:size=640x480:rate=30 \
				-f lavfi -i sine=frequency=1000:duration=5 \
				-c:v libx264 -c:a aac "$TEST_VIDEO" -y >/dev/null 2>&1
			print_status "Created test video"
		else
			echo -e "${YELLOW}[SKIP]${NC} ffmpeg not available, skipping video test"
			return 0
		fi
	fi

	# Test basic face detection
	OUTPUT_FILE="$OUTPUT_DIR/face_detection.json"
	python3 scripts/face_processor.py "$TEST_VIDEO" "$OUTPUT_FILE" --uuid "test_face"

	if [ -f "$OUTPUT_FILE" ]; then
		echo -e "${YELLOW}[INFO]${NC} Face detection output:"
		jq '.frames | length' "$OUTPUT_FILE"
		print_status "Basic face processor works"
	else
		echo -e "${RED}[FAIL]${NC} Face processor did not create output file"
		return 1
	fi
}

# Function to test face recognition processor
test_face_recognition_processor() {
	echo -e "${YELLOW}[INFO]${NC} Testing face recognition processor..."

	# Check if test video exists
	if [ ! -f "$TEST_VIDEO" ]; then
		echo -e "${YELLOW}[SKIP]${NC} Test video not found, skipping recognition test"
		return 0
	fi

	# Test face recognition with CPU
	OUTPUT_FILE="$OUTPUT_DIR/face_recognition_cpu.json"
	echo -e "${YELLOW}[INFO]${NC} Testing with CPU..."
	python3 scripts/face_recognition_processor.py \
		"$TEST_VIDEO" \
		"$OUTPUT_FILE" \
		"1" "1" "1" \
		--uuid "test_recognition_cpu"

	if [ -f "$OUTPUT_FILE" ]; then
		echo -e "${YELLOW}[INFO]${NC} CPU Face recognition output:"
		jq '. | {frames: .frames | length, recognized_faces: .recognized_faces | length, clusters: .face_clusters | length}' "$OUTPUT_FILE"
		print_status "Face recognition processor works with CPU"
	else
		echo -e "${RED}[FAIL]${NC} Face recognition processor did not create output file"
		return 1
	fi

	# Test face recognition with MPS (if available)
	OUTPUT_FILE_MPS="$OUTPUT_DIR/face_recognition_mps.json"
	echo -e "${YELLOW}[INFO]${NC} Testing with MPS acceleration..."
	python3 scripts/face_recognition_processor.py \
		"$TEST_VIDEO" \
		"$OUTPUT_FILE_MPS" \
		"1" "1" "1" \
		--uuid "test_recognition_mps" \
		--use-mps

	if [ -f "$OUTPUT_FILE_MPS" ]; then
		echo -e "${YELLOW}[INFO]${NC} MPS Face recognition output:"
		jq '. | {frames: .frames | length, recognized_faces: .recognized_faces | length, clusters: .face_clusters | length}' "$OUTPUT_FILE_MPS"
		print_status "Face recognition processor works with MPS"
	else
		echo -e "${YELLOW}[WARN]${NC} MPS acceleration not available or failed, using CPU fallback"
	fi
}

# Function to test face registration
test_face_registration() {
	echo -e "${YELLOW}[INFO]${NC} Testing face registration..."

	# Check if test image exists
	if [ ! -f "$TEST_IMAGE" ]; then
		echo -e "${YELLOW}[WARN]${NC} Test image not found, creating dummy image..."
		# Create a simple test image with ImageMagick if available
		if command -v convert &>/dev/null; then
			convert -size 640x480 xc:gray -pointsize 72 -fill white -draw "text 100,240 'Test Face'" "$TEST_IMAGE"
			print_status "Created test image"
		else
			echo -e "${YELLOW}[SKIP]${NC} ImageMagick not available, skipping registration test"
			return 0
		fi
	fi

	# Test face registration
	OUTPUT_FILE="$OUTPUT_DIR/face_registration.json"
	python3 scripts/face_registration.py \
		"$TEST_IMAGE" \
		"$OUTPUT_FILE" \
		"Test Person" \
		--database "$OUTPUT_DIR/$FACE_DATABASE"

	if [ -f "$OUTPUT_FILE" ]; then
		echo -e "${YELLOW}[INFO]${NC} Face registration output:"
		jq '.' "$OUTPUT_FILE"
		print_status "Face registration works"
	else
		echo -e "${RED}[FAIL]${NC} Face registration did not create output file"
		return 1
	fi
}

# Function to test API endpoints
test_api_endpoints() {
	echo -e "${YELLOW}[INFO]${NC} Testing API endpoints..."

	# Note: These tests require the API to be running and a valid API key
	# For now, we'll just check if the endpoints are defined in the code

	echo -e "${YELLOW}[INFO]${NC} API endpoints defined:"
	echo "  POST /api/v1/face/recognize"
	echo "  POST /api/v1/face/register"
	echo "  POST /api/v1/face/search"
	echo "  GET  /api/v1/face/list"
	echo "  GET  /api/v1/face/{face_id}"
	echo "  DELETE /api/v1/face/{face_id}"
	echo "  GET  /api/v1/face/results/{video_uuid}"

	print_status "API endpoints are defined"
}

# Function to test database migration
test_database_migration() {
	echo -e "${YELLOW}[INFO]${NC} Testing database migration..."

	# Check if migration file exists
	MIGRATION_FILE="migrations/006_face_recognition_tables.sql"
	if [ -f "$MIGRATION_FILE" ]; then
		echo -e "${YELLOW}[INFO]${NC} Migration file content check:"
		grep -c "CREATE TABLE" "$MIGRATION_FILE"
		grep -c "face_identities" "$MIGRATION_FILE"
		grep -c "face_detections" "$MIGRATION_FILE"
		grep -c "face_clusters" "$MIGRATION_FILE"
		print_status "Database migration file is valid"
	else
		echo -e "${RED}[FAIL]${NC} Migration file not found: $MIGRATION_FILE"
		return 1
	fi
}

# Function to run Rust tests
test_rust_code() {
	echo -e "${YELLOW}[INFO]${NC} Testing Rust code..."

	# Check if face_recognition module exists
	if [ -f "src/core/processor/face_recognition.rs" ]; then
		echo -e "${YELLOW}[INFO]${NC} Face recognition module exists"

		# Run cargo check
		cargo check --lib
		print_status "Rust code compiles"

		# Run specific tests
		cargo test --lib face_recognition -- --nocapture
		print_status "Face recognition tests pass"
	else
		echo -e "${RED}[FAIL]${NC} Face recognition module not found"
		return 1
	fi
}

# Main test sequence
main() {
	echo "Starting face recognition tests..."
	echo ""

	# Run tests
	test_python_deps
	echo ""

	test_face_processor
	echo ""

	test_face_recognition_processor
	echo ""

	test_face_registration
	echo ""

	test_api_endpoints
	echo ""

	test_database_migration
	echo ""

	test_rust_code
	echo ""

	echo "========================================="
	echo -e "${GREEN}All tests completed successfully!${NC}"
	echo "========================================="
	echo ""
	echo "Next steps:"
	echo "1. Install InsightFace: pip install insightface"
	echo "2. Run database migration: psql -d momentry -f migrations/006_face_recognition_tables.sql"
	echo "3. Start Momentry Core: cargo run --bin momentry -- server"
	echo "4. Test API endpoints with curl or Postman"
	echo ""
}

# Run main function
main

# Cleanup
echo -e "${YELLOW}[INFO]${NC} Cleaning up test files..."
rm -rf "$OUTPUT_DIR"
if [ -f "$TEST_VIDEO" ] && [ ! -f "test_video.mp4" ]; then
	rm "$TEST_VIDEO"
fi
if [ -f "$TEST_IMAGE" ] && [ ! -f "test_face.jpg" ]; then
	rm "$TEST_IMAGE"
fi

echo -e "${GREEN}Test completed!${NC}"