#!/bin/bash

set -e  # Exit on error

# Default values
CONFIG_FILE="main/resources/local_config.yaml"
ENV_FILE=".env"
CONTAINER_NAME="llm-inference-server"
BUILD_ONLY=false
PLATFORM=""

# Function to display usage
usage() {
    echo "Usage: $0 [-c <config_file>] [-e <env_file>] [-n <container_name>] [-b] [-h]"
    echo "  -c : Config file path (default: resources/local_config.yaml)"
    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
}

# Function to read yaml using grep and sed
parse_yaml() {
    local file=$1
    # Get port - look for port under server section and extract the number
    PORT=$(grep -A5 "^server:" "$file" | grep "port:" | sed 's/[^0-9]*//g')
    # Get host - look for host under server section and extract the value in quotes
    HOST=$(grep -A5 "^server:" "$file" | grep "host:" | sed 's/.*"//;s/".*//')
    echo "$PORT $HOST"
}

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

# Check if config file exists
if [ ! -f "$CONFIG_FILE" ]; then
    handle_error "Config file not found: $CONFIG_FILE"
fi

# Read port and host from config
read PORT HOST < <(parse_yaml "$CONFIG_FILE")
echo "Using configuration - Port: $PORT, Host: $HOST"

# 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 in interactive mode
echo "Starting container on port $PORT..."
echo "Press Ctrl+C to stop the server"
docker run --rm -it \
    --name $CONTAINER_NAME \
    --env-file $ENV_FILE \
    -p $PORT:$PORT \
    -v $HOME/.cache/huggingface:/app/.cache/huggingface \
    -v "$(pwd)/$CONFIG_FILE:/app/$CONFIG_FILE" \
    $CONTAINER_NAME

echo "Container stopped"