#!/bin/bash

set -euo pipefail

function apt_update_if_stale () {
    if ! find /var/lib/apt/lists -type f -mmin -30 | grep -q .; then
        apt-get update
    else
        echo "Skipping apt update, metadata is fresh." > /dev/stderr
    fi
}

function main () {
    # check for ubuntu only
    if ! grep -qi '^ID=ubuntu' /etc/os-release; then
        echo "only for ubuntu" > /dev/stderr
        exit 1
    fi

    # check for root
    if [ "$EUID" -ne 0 ]; then
        echo "Please run as root" > /dev/stderr
        exit 1
    fi

    if ! grep -qE 'http(s)?://archive\.ubuntu\.com/' /etc/apt/sources.list.d/*; then
        echo "Not using default ubuntu repo, nothing for us to change. Exiting." > /dev/stderr
        exit 1
    fi

    apt_update_if_stale
    if ! which jq > /dev/null; then
        apt-get install -y jq
    fi

    if ! which curl > /dev/null; then
        apt-get install -y curl
    fi

    echo "Changing mirror country to US as a first approximation..." > /dev/stderr
    find /etc/apt/sources.list.d/ -type f -exec \
      sed -i "s|archive\.ubuntu\.com|us.archive.ubuntu.com|g" {} +

    echo "Detecting actual country..." > /dev/stderr

    IPINFO="$(curl -s ipinfo.io)"

    if [ -z "$IPINFO" ]; then
        echo "Failed to geolocate (fetch).  Exiting." > /dev/stderr
        exit 1
    fi

    COUNTRYCODE="$(jq .country -r <<< "$IPINFO")"

    if [[ -z "$COUNTRYCODE" ]]; then
        echo "Failed to geolocate (parse).  Exiting." > /dev/stderr
        exit 1
    fi

    echo "IP geolocated as country: $COUNTRYCODE" > /dev/stderr

    if [[ "$COUNTRYCODE" == "US" ]]; then
        echo "Mirror was already set to US, nothing to do.  Exiting." > /dev/stderr
        exit 0
    fi

    MIRRORHOST="$(echo "$COUNTRYCODE" | tr 'A-Z' 'a-z').archive.ubuntu.com"

    echo "Changing mirror country to $COUNTRYCODE..." > /dev/stderr

    find /etc/apt/sources.list.d/ -type f -exec \
      sed -i "s|us\.archive\.ubuntu\.com|$MIRRORHOST|g" {} +

    # update is actually mandatory after changing the mirror:
    apt-get update

    echo "Finished." > /dev/stderr

    exit 0
}

main