#!/bin/bash

set -euo pipefail

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, script will not function" > /dev/stderr
        exit 1
    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.ubuntu.com|g" {} +


    UPDATED=0
    if ! which jq > /dev/null; then
        echo "jq not found, installing..." > /dev/stderr
        if [[ "$UPDATED" -eq 0 ]]; then
            apt update
            UPDATED=1
        fi
        apt install -y jq
    fi

    if ! which curl > /dev/null; then
        echo "curl not found, installing..." > /dev/stderr
        if [ "$UPDATED" -eq 0 ]; then
            apt update
            UPDATED=1
        fi
        apt install -y curl
    fi

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

    IPINFO="$(curl -s ipinfo.io)"
    if [ -z "$IPINFO" ]; then
        echo "Failed to get IP info" > /dev/stderr
        exit 1
    fi
    COUNTRYCODE="$(jq .country -r <<< "$IPINFO")"
    echo "IP geolocated as country: $COUNTRYCODE" > /dev/stderr

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

    if [[ "$COUNTRYCODE" == "US" ]]; then
        echo "Mirror already set to US, nothing to do" > /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" {} +
    exit 0
}

main