#!/bin/bash
set -e

# Logging configuration
LOG_DIR="/var/log/whisper_asr_setup"
LOG_FILE="$LOG_DIR/install_$(date +%Y%m%d_%H%M%S).log"

# Create log directory
mkdir -p "$LOG_DIR"
chmod 755 "$LOG_DIR"

# Logging function
log() {
    local timestamp=$(date "+%Y-%m-%d %H:%M:%S")
    echo "[$timestamp] $1" | tee -a "$LOG_FILE"
}

# Command execution with error checking
execute() {
    log "Executing: $1"
    eval "$1" >> "$LOG_FILE" 2>&1
    if [ $? -ne 0 ]; then
        log "ERROR: Command failed - $1"
        exit 1
    fi
    log "SUCCESS: Command completed - $1"
}

log "Starting Whisper ASR installation process..."

# Update package lists
execute "apt update -y"

# Install Docker
execute "apt install docker.io -y"

# Add NVIDIA container toolkit repository
execute "curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
execute "curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list"

# Update and install NVIDIA toolkit
execute "apt-get update -y"
execute "apt-get install -y nvidia-container-toolkit"

# Configure Docker daemon
log "Configuring Docker daemon..."
cat <<EOF | sudo tee /etc/docker/daemon.json >> "$LOG_FILE" 2>&1
{
    "exec-opts": [
        "native.cgroupdriver=systemd"
    ],
    "log-driver": "json-file",
    "log-opts": {
        "max-size": "100m"
    },
    "default-runtime": "nvidia",
    "runtimes": {
        "nvidia": {
            "args": [],
            "path": "nvidia-container-runtime"
        }
    },
    "storage-driver": "overlay2"
}
EOF
[ $? -eq 0 ] && log "SUCCESS: Docker daemon configured" || log "ERROR: Docker configuration failed"

# Restart Docker
execute "systemctl restart docker"
log "Waiting for Docker to stabilize..."
sleep 2

# Run Whisper ASR container
log "start pull images . Log file: $LOG_FILE"
execute "docker run -d --gpus all -p 9000:9000 -e ASR_MODEL=turbo -e ASR_ENGINE=openai_whisper onerahmet/openai-whisper-asr-webservice:latest-gpu"

log "Installation completed successfully. Log file: $LOG_FILE"

MAX_WAIT_SECONDS=600
CHECK_INTERVAL=5

start_time=$(date +%s)
end_time=$((start_time + MAX_WAIT_SECONDS))

log "Start checking whether port 9000 is listening, wait for up to 10 minutes..."

while [ $(date +%s) -lt $end_time ]; do
    if ss -tuln | grep -q ':9000\b'; then
        log "$(date): 9000 port already listen"
        exit 0
    fi

    log "$(date): Port 9000 is not listening, waiting ${CHECK_INTERVAL} seconds and trying again..."
    sleep $CHECK_INTERVAL
done

log "Error: No listening on port 9000 detected within 10 minutes"
exit 1
