#!/usr/bin/env bash

set -Eeuo pipefail

if [[ $# -lt 1 ]]; then
    echo "Usage: $0 WORK_DIR [SUMKA_ARGUMENT ...]" >&2
    exit 64
fi

work_dir=$1
shift

if [[ ! -d "$work_dir" ]]; then
    echo "Work directory does not exist: $work_dir" >&2
    exit 66
fi

input_found=false
for extension in mp4 mkv avi mp3 m4a wav; do
    if [[ -f "$work_dir/input.$extension" ]]; then
        input_found=true
        break
    fi
done
if [[ "$input_found" != true ]]; then
    echo "No supported input file found in: $work_dir" >&2
    exit 66
fi

: "${AI_API_KEY:?AI_API_KEY is required}"

sumka_image=${SUMKA_IMAGE:-sumka:amd}
container_name=${SUMKA_CONTAINER_NAME:-matrix-sumka-job}
cache_volume=${SUMKA_CACHE_VOLUME:-whisper-cache}
tmpfs_size=${SUMKA_TMPFS_SIZE:-8g}

docker_args=(
    run
    --rm
    --name "$container_name"
    --env AI_API_KEY
    --mount "type=volume,src=${cache_volume},dst=/tmp/sumka-cache/whisper"
    --tmpfs "/tmp:rw,size=${tmpfs_size}"
)

gpu_kind=${SUMKA_GPU:-auto}
if [[ "$gpu_kind" == auto ]]; then
    case "$sumka_image" in
        *nvidia*) gpu_kind=nvidia ;;
        *amd*) gpu_kind=amd ;;
        *) gpu_kind=amd ;;
    esac
fi

case "$gpu_kind" in
    nvidia)
        docker_args+=(--gpus all)
        ;;
    amd)
        docker_args+=(
            --device=/dev/kfd
            --device=/dev/dri
            --group-add video
            --security-opt seccomp=unconfined
        )
        ;;
    *)
        echo "Unsupported SUMKA_GPU value: $gpu_kind (expected amd or nvidia)" >&2
        exit 64
        ;;
esac

volumes_from=${SUMKA_VOLUMES_FROM:-}
if [[ -z "$volumes_from" && -f /.dockerenv ]]; then
    # Docker uses the short container ID as the default hostname. Unlike a
    # Compose service name, it also works for `docker compose run` containers.
    volumes_from=${HOSTNAME:-}
fi

if [[ -n "$volumes_from" ]]; then
    # The sibling sumka container sees the bot's /runtime bind mount at the
    # same path, independently of how the bot container itself was named.
    docker_args+=(
        --volumes-from "${volumes_from}:rw"
        --workdir "$work_dir"
    )
else
    # Direct host/.venv launch: mount only this job's directory.
    work_dir=$(realpath "$work_dir")
    docker_args+=(
        --mount "type=bind,src=${work_dir},dst=/work"
        --workdir /work
    )
fi

signal_exit_code=0

stop_container() {
    docker stop --time 10 "$container_name" >/dev/null 2>&1 || true
}

on_term() {
    signal_exit_code=143
    stop_container
}

on_int() {
    signal_exit_code=130
    stop_container
}

on_hup() {
    signal_exit_code=129
    stop_container
}

trap on_term TERM
trap on_int INT
trap on_hup HUP

set +e
docker "${docker_args[@]}" "$sumka_image" "$@" &
docker_pid=$!
wait "$docker_pid"
exit_code=$?
set -e

if [[ $signal_exit_code -ne 0 ]]; then
    exit "$signal_exit_code"
fi
exit "$exit_code"