#!/bin/bash

# FastAPI Data Collector Container Runner
# Automated script to build and run the data collector in a container
# Supports both Podman and Docker with automatic detection

set -e

# Configuration
CONTAINER_NAME="fastapi-data-collector"
IMAGE_NAME="fastapi-data-collector:latest"
HOST_PORT="8000"
CONTAINER_PORT="8000"
DATA_DIR="./input"

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

# Function to print colored output
print_status() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

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

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_header() {
    echo -e "${BLUE}[HEADER]${NC} $1"
}

# Detect container runtime
detect_container_runtime() {
    if command -v podman &> /dev/null; then
        CONTAINER_CMD="podman"
        print_status "Using Podman as container runtime"
    elif command -v docker &> /dev/null; then
        CONTAINER_CMD="docker"
        print_status "Using Docker as container runtime"
    else
        print_error "Neither Podman nor Docker found. Please install one of them."
        exit 1
    fi
}

# Stop and remove existing container
cleanup_existing_container() {
    print_status "Checking for existing container..."
    
    if $CONTAINER_CMD ps -a --format "table {{.Names}}" | grep -q "^${CONTAINER_NAME}$"; then
        print_warning "Stopping existing container: ${CONTAINER_NAME}"
        $CONTAINER_CMD stop "${CONTAINER_NAME}" || true
        
        print_warning "Removing existing container: ${CONTAINER_NAME}"
        $CONTAINER_CMD rm "${CONTAINER_NAME}" || true
    else
        print_status "No existing container found"
    fi
}

# Build container image
build_image() {
    print_status "Building container image: ${IMAGE_NAME}"
    
    if ! $CONTAINER_CMD build -t "${IMAGE_NAME}" .; then
        print_error "Failed to build container image"
        exit 1
    fi
    
    print_status "Container image built successfully"
}

# Create data directory
setup_data_directory() {
    if [ ! -d "${DATA_DIR}" ]; then
        print_status "Creating data directory: ${DATA_DIR}"
        mkdir -p "${DATA_DIR}"
    else
        print_status "Data directory already exists: ${DATA_DIR}"
    fi
}

# Check for .env file
check_env_file() {
    if [ ! -f ".env" ]; then
        print_warning ".env file not found"
        print_status "Creating .env file from .env.example..."
        
        if [ -f ".env.example" ]; then
            cp .env.example .env
            print_warning "Please edit .env file with your actual API keys before running the container"
            print_status "You can edit .env now or the container will use example keys"
            
            # Ask user if they want to edit now
            read -p "Do you want to edit .env file now? (y/N): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                ${EDITOR:-nano} .env
            fi
        else
            print_error ".env.example file not found. Cannot create .env file."
            exit 1
        fi
    else
        print_status ".env file found"
    fi
}

# Run container
run_container() {
    print_status "Starting container: ${CONTAINER_NAME}"
    
    # Convert relative path to absolute path
    DATA_DIR_ABS=$(realpath "${DATA_DIR}")
    ENV_FILE_ABS=$(realpath ".env")
    
    print_status "Mounting data directory: ${DATA_DIR_ABS}"
    print_status "Mounting .env file: ${ENV_FILE_ABS}"
    
    $CONTAINER_CMD run -d \
        --name "${CONTAINER_NAME}" \
        --publish "${HOST_PORT}:${CONTAINER_PORT}" \
        --volume "${DATA_DIR_ABS}:/app/input:Z" \
        --volume "${ENV_FILE_ABS}:/app/.env:ro,Z" \
        --restart unless-stopped \
        --security-opt no-new-privileges \
        --cap-drop ALL \
	--cap-add=CHOWN \
        --cap-add=SETGID \
        --cap-add=SETUID \
        --cap-add=DAC_OVERRIDE \
        --read-only \
        --tmpfs /tmp \
        "${IMAGE_NAME}"
    
    if [ $? -eq 0 ]; then
        print_status "Container started successfully"
    else
        print_error "Failed to start container"
        exit 1
    fi
}

# Verify container is running
verify_container() {
    print_status "Verifying container startup..."
    
    # Wait a moment for the container to start
    sleep 3
    
    if $CONTAINER_CMD ps --format "table {{.Names}}" | grep -q "^${CONTAINER_NAME}$"; then
        print_status "Container is running"
        
        # Show container logs
        print_status "Container logs:"
        $CONTAINER_CMD logs "${CONTAINER_NAME}" | tail -10
        
        # Test health endpoint
        print_status "Testing health endpoint..."
        sleep 2
        if curl -s "http://localhost:${HOST_PORT}/health" > /dev/null; then
            print_status "Health check passed"
        else
            print_warning "Health check failed - service may still be starting"
        fi
    else
        print_error "Container is not running"
        print_error "Container logs:"
        $CONTAINER_CMD logs "${CONTAINER_NAME}"
        exit 1
    fi
}

# Show usage information
show_usage_info() {
    print_header "=== FastAPI Data Collector Container Started ==="
    echo
    print_status "Service URL: http://localhost:${HOST_PORT}"
    print_status "API Documentation: http://localhost:${HOST_PORT}/docs"
    print_status "Health Check: http://localhost:${HOST_PORT}/health"
    echo
    print_status "API Key Configuration:"
    echo "  - Keys are loaded from .env file"
    echo "  - Separate keys for INPUT and READ operations"
    echo "  - Multiple keys supported per operation type"
    echo
    print_status "Example Usage:"
    echo "  # Send data (use any key from INPUT_API_KEYS):"
    echo '  curl -X POST http://localhost:8000/api/run1/ \'
    echo '    -H "Authorization: Bearer <YOUR_INPUT_API_KEY>" \'
    echo '    -d "host_a,is ok"'
    echo
    echo "  # Read results (use any key from READ_API_KEYS):"
    echo '  curl -H "Authorization: Bearer <YOUR_READ_API_KEY>" \'
    echo '    http://localhost:8000/results/run1/'
    echo
    print_status "Data Directory: ${DATA_DIR_ABS}"
    echo
    print_status "Container Management:"
    echo "  - View logs: ${CONTAINER_CMD} logs ${CONTAINER_NAME}"
    echo "  - Stop container: ${CONTAINER_CMD} stop ${CONTAINER_NAME}"
    echo "  - Restart container: ${CONTAINER_CMD} restart ${CONTAINER_NAME}"
    echo "  - Remove container: ${CONTAINER_CMD} rm ${CONTAINER_NAME}"
    echo
    print_header "=============================================="
}

# Main execution
main() {
    print_header "FastAPI Data Collector Container Setup"
    echo
    
    detect_container_runtime
    cleanup_existing_container
    build_image
    setup_data_directory
    check_env_file
    run_container
    verify_container
    show_usage_info
}

# Run main function
main "$@"