#!/bin/bash

set -e  # Exit on error

# Default values
PORT=7860
ENV_FILE=".env"
CONTAINER_NAME="llm-inference-server"
BUILD_ONLY=false
PLATFORM=""

# Function to display usage
usage() {
    echo "Usage: $0 [-p <port>] [-e <env_file>] [-n <container_name>] [-b] [-h]"
    echo "  -p : Port to expose (default: 7860)"
    echo "  -e : Environment file path (default: .env)"
    echo "  -n : Container name (default: llm-inference-server)"
    echo "  -b : Build only (don't run container)"
    echo "  -h : Display this help message"
    exit 1
}

# Function to handle errors
handle_error() {
    echo "Error: $1"
    exit 1
}

# Parse command line arguments
while getopts "p:e:n:bh" opt; do
    case $opt in
        p) PORT=$OPTARG ;;
        e) ENV_FILE=$OPTARG ;;
        n) CONTAINER_NAME=$OPTARG ;;
        b) BUILD_ONLY=true ;;
        h) usage ;;
        ?) usage ;;
    esac
done

# Detect platform and set appropriate options
if [[ "$(uname -s)" == "Darwin" ]]; then
    echo "Detected MacOS platform"
    PLATFORM="linux/arm64"
    if [[ "$(uname -m)" == "x86_64" ]]; then
        PLATFORM="linux/amd64"
    fi
    PLATFORM_ARG="--platform=${PLATFORM}"
else
    echo "Detected Linux platform"
    # Check if podman is available (common in RHEL environments)
    if command -v podman &> /dev/null; then
        echo "Podman detected, using podman instead of docker"
        function docker { podman "${@}"; }
    fi
    PLATFORM_ARG=""
fi

# Check if docker is installed
if ! command -v docker &> /dev/null; then
    handle_error "Docker is not installed"
fi

# Check if .env file exists
if [ ! -f "$ENV_FILE" ]; then
    echo "Warning: $ENV_FILE file not found"
    echo "Creating sample $ENV_FILE file..."
    cat > "$ENV_FILE" << EOL
InfAPITokenWrite=your_huggingface_token_here
EOL
    echo "Please edit $ENV_FILE with your actual configuration"
fi

# Build the Docker image
echo "Building Docker image..."
docker build $PLATFORM_ARG -t $CONTAINER_NAME . || handle_error "Failed to build Docker image"

# Exit if build only
if [ "$BUILD_ONLY" = true ]; then
    echo "Build completed. Exiting as requested."
    exit 0
fi

# Check if container is already running
if docker ps -q -f name=$CONTAINER_NAME | grep -q .; then
    echo "Stopping existing container..."
    docker stop $CONTAINER_NAME || handle_error "Failed to stop existing container"
fi

# Remove existing container
if docker ps -aq -f name=$CONTAINER_NAME | grep -q .; then
    echo "Removing existing container..."
    docker rm $CONTAINER_NAME || handle_error "Failed to remove existing container"
fi

# Run the container
echo "Starting container on port $PORT..."
if ! docker run -d \
    --name $CONTAINER_NAME \
    --env-file $ENV_FILE \
    -p $PORT:7860 \
    -v $HOME/.cache/huggingface:/app/.cache/huggingface \
    $CONTAINER_NAME; then
    handle_error "Failed to start container"
fi

# Verify container is running
if ! docker ps | grep -q $CONTAINER_NAME; then
    handle_error "Container failed to start. Check logs with: docker logs $CONTAINER_NAME"
fi

echo "Container started successfully!"
echo "API should be available at http://localhost:$PORT"
echo "To view logs, run: docker logs -f $CONTAINER_NAME"