#!/bin/bash

echo2()  { echo -e "$@" >&2; }
WARN()   { echo2 "==> $progname: warning: $1"; }
DIE()    { Cleanup; echo2 "==> $progname: error: $1"; exit 1; }

ASSERT_DIE()  { "$@" || DIE "'$*' failed."; }
ASSERT_WARN() { "$@" && return 0; WARN "'$*' failed."; return 1; }

FetchMirrors() {
    [ $has_internet_connection = no ] && return 1
    local -r active_mirrors_url=https://archlinux.org/mirrorlist/all     # mirrors active at the time listed in the file, both https and http
    local -r url=$active_mirrors_url/
    local -r file="$mirrordata"
    local -r fetched=$file.tmp
    local operation=update

    if curl -Lsm $TIMEOUT_ACTIVE_MIRRORS -o"$fetched" "$url" ; then
        if [ -e "$file" ] ; then
            ASSERT_WARN rm -f "$file".bak      || return 1
            ASSERT_WARN mv "$file" "$file".bak || return 1
        else
            operation=create
        fi
        ASSERT_WARN mv "$fetched" "$file"      || return 1
        echo2 "==> ${operation}d $file."
        rm -f "$file".bak
    else
        rm -f "$fetched"
        if [ ! -e "$file" ] ; then
            WARN "could not $operation '$file'."
            return 1
        fi
    fi
    return 0
}

FetchCountries() {
    [ $has_internet_connection = no ] && return 1
    local info codes=() countries=()
    local reflector_command=(
        /bin/reflector
        --list-countries
        --connection-timeout=$TIMEOUT_REFLECTOR_CONNECTION
        --download-timeout=$TIMEOUT_REFLECTOR_DOWNLOAD
    )

    info=$("${reflector_command[@]}" 2>/dev/null | /bin/sed -E '/^Country[ ]+Code/,/^-----/d')
    if [ -z "$info" ] ; then
        WARN "could not update '$dbfile'"
        return 1
    fi
    # shellcheck disable=SC2207
    codes=(WW $(echo "$info" | sed -E 's|(.*[a-z])[ ]+([A-Z][A-Z])[ ]+[0-9]+|\2|'))
    # shellcheck disable=SC1090,SC2046
    readarray -t countries <<< $(echo "Worldwide"; echo "$info" | sed -E 's|(.*[a-z])[ ]+([A-Z][A-Z])[ ]+[0-9]+|\1|')

    local code country ix count=${#codes[*]}

    echo -e "#!/bin/bash\nreflector_countries=(" > "$dbfile"

    for ((ix=0; ix < count; ix++ )) ; do
        code=${codes[$ix]}
        country=${countries[$ix]}
        echo "    [${code,,}]='$country'" >> "$dbfile"
    done
    echo ")" >> "$dbfile"
    echo2 "==> created $dbfile."
    return 0
}

CC2name() {
    echo "${reflector_countries[$1]}"
}

ExtractCountryMirrors() {
    local countryname="$1" mirrors_in_country
    mirrors_in_country=$(cat "$mirrordata" | sed -n "/^## ${countryname}$/,/^$/p" | sed -E 's|^#(Server = )|\1|')
    # include wanted protocols
    for pp in "${protocols[@]}" ; do
        echo "$mirrors_in_country" | grep "$pp://"
    done
}

AddCountry() {
    local cc="$1"
    [ "$cc" ] || return
    local cn=${reflector_countries[$cc]}
    if [ -z "$cn" ] ; then
        WARN "'$cc': no active Arch mirrors found for this country"
        return
    fi
    if printf "%s\n" "${countries_handled[@]}" | grep "^$cc$" >/dev/null ; then
        return
    else
        echo2 "==> Including mirrors from $cn"
        countries_handled+=("$cc")
        ExtractCountryMirrors "$cn" >> "$tmplist"
    fi
}

Cleanup() {
    [ "$cleanup_files" ] && rm -f "${cleanup_files[@]}"
}

ShowLocationInfo() {
    local data url
    for url in https://ipinfo.io https://ipapi.co ; do
        data=$(curl -Lsm 10 -O- $url)
        if [ "$data" ] ; then
            data=$(echo "$data" | grep '"country"' | sed -E 's|.*"([A-Z][A-Z])",$|\1|')
            echo "${data,,}"
            return
        fi
    done
    WARN "cannot get current location"
}

UserCountriesFromFile() {
    if [ -r "$user_fav_countries" ] ; then
        local item item2 items
        items=$(cat "$user_fav_countries")
        for item in $items ; do
            item2="$item"
            item=${item,,}
            case "$item" in
                [a-z][a-z]) printf "%s\n" "${countries_given[@]}" | grep "^$item$" >/dev/null || countries_given+=("$item") ;;
                *)          WARN "$user_fav_countries: country code '$item2' ignored" ;;
            esac
        done
    else
        WARN "cannot read file '$user_fav_countries'"
    fi
}

DumpOptions() {
    if [ "$OPTS" ] ; then
        local o=${OPTS//:/}                   # remove every ':'
        [ "${o::1}" = "$sep" ] || o="--$o"    # add leading '--' if first option is long
        o=${o//,/ --}                         # manage long options
        o=${o//$sep/ -}                       # manage short options
        echo "$o"
    fi
}

Header() {
    echo -e "### Program: $progname version $(expac %v $pkgname)"
    echo -e "### Mirror list generated at: $(date -u "+%x %X") UTC"
    echo -n "### Command: $progname"
    if [ "${orig_args[0]}" ] ; then
        printf " '%s'" "${orig_args[@]}"
    fi
    echo -e "\n"

    if [ "$prefslist" ] ; then
        echo -e "####### >>>> Mirrors by user preference"
        for item in $prefslist ; do
            item=$(grep "$item" $mirrordata | awk '{print $NF}')
            echo "Server = $item"
        done
        echo -e "####### <<<< Mirrors by user preference\n"
    fi
}

Main2() {
    local tmplist mirrorlist
    tmplist=$(mktemp)                           # collect mirrors that user wanted here
    mirrorlist=$(mktemp)

    chmod go-rwx "$tmplist" "$mirrorlist"
    cleanup_files+=("$tmplist" "$mirrorlist")

    # if user wanted, add the current country into $tmplist
    if [ $local_country_wanted = yes ] ; then   # && [ $has_internet_connection = yes ] ; then
        AddCountry "$country_code"
    fi

    # add user given countries into $tmplist
    local cc
    for cc in "${countries_given[@]}" ; do
        case "$cc" in
            [a-z][a-z]) AddCountry "$cc" ;;
            *)          WARN "country code '$cc' is not supported" ;;
        esac
    done

    if [ $has_internet_connection = yes ] ; then
        local rankmirrors_opt=()
        [ $verbose  = yes ] && rankmirrors_opt+=(-v)
        [ $parallel = yes ] && rankmirrors_opt+=(-p)

        echo2 "==> Ranking mirrors."
        {
            Header
            # shellcheck disable=SC2016
            create-ml-rankmirrors-arch "${rankmirrors_opt[@]}" "$tmplist" | column -t -s'|' | sed 's|/lastupdate|/$repo/os/$arch|'
        } > "$mirrorlist"
        echo2 "==> Mirrors ranked."
    else
        local msg_not_ranked="NOTE: this mirrorlist was NOT ranked due to unavailable internet connection."
        local msg_not_ranked2="We strongly recommend to rank it soon."
        {
            Header
            echo -e "### $msg_not_ranked"
            echo -e "### $msg_not_ranked2\n"
        } > "$mirrorlist"
        cat "$tmplist" >> "$mirrorlist"
    fi

    if [ $save = yes ] ; then
        if [ "$target" = $target_def ] ; then
            echo2 "==> Updating $target:"
            sudo rm -f "$target.bak"
            sudo mv "$target" "$target.bak"
            sudo cp -iv "$mirrorlist" "$target"
            sudo chmod go+r "$target"
        else
            if [ -w "${target%/*}" ] ; then
                echo2 "==> Updating $target."
                cp "$mirrorlist" "$target"
            else
                echo2 "==> Updating $target:"
                sudo cp "$mirrorlist" "$target"
            fi
        fi
        if [ $has_internet_connection = no ] ; then
            echo2 "==> $msg_not_ranked"
            echo2 "==> $msg_not_ranked2"
        fi
    else
        echo2 ""
        cat "$mirrorlist" >&2
        echo2 "\n==> Tip: use option --save to update $target."
    fi
    Cleanup
}

AddRecommendedCountries() {
    # Adding recommended countries for ranking.
    # Note that the current country will be added later if user wants it and has Arch mirrors.
    case "$country_code" in
        de|us|cn)
            : ;;
        ca)
            countries_given+=(us) ;;
        au)
            countries_given+=(ww) ;;
        br|cl|co|mx)
            countries_given+=(ww us); protocols+=(http) ;;
        'fi')
            countries_given+=(ww se) ;;
        al|at|be|cz|dk|ee|es|fr|gb|gr|hr|it|nl|no|pl|pt|se)
            countries_given+=(ww de fr) ;;
        *)
            countries_given+=(ww de us); protocols+=(http) ;;
    esac
}

DumpCCs() {
    # shellcheck disable=SC1090
    source "$dbfile"                                               # gets array reflector_countries
    # shellcheck disable=SC2046
    echo $(printf "%s\n" "${!reflector_countries[@]}" | sort)
}

Parameters() {
    local lopts sopts
    lopts="$(echo "$OPTS" | sed -E 's|(/.[:]*)||g')"
    sopts="$(echo "$OPTS" | sed -E 's|[^/]*/(.[:]*)[^/]*|\1|g')"

    local opts

    opts="$(/bin/getopt -o="$sopts" --longoptions "$lopts" --name "$progname" -- "$@")" || exit 1
    eval set -- "$opts"

    while [ "$1" ] ; do
        case "$1" in
            --)                  shift; break ;;

            --user-countries)    UserCountriesFromFile ;;
            --fake-country)      fake_country="$2"; shift ;;
            --offline)           has_internet_connection=no ;;
            --nolocal)           local_country_wanted=no ;;
            --http | --rsync)    protocols+=("${1:2}") ;;
            --sequential)        parallel=no ;;
            --save)              save=yes; target=$target_def ;;
            --savefile)          save=yes; target="$2"; shift ;;
            -v | --verbose)      verbose=yes ;;
            --no-recommended-countries) use_recommended_countries=no ;;
            --dump-options)      DumpOptions; exit 0 ;;
            --dump-ccs)          DumpCCs; exit 0 ;;
            --update-supports)   update_supports=yes ;;
            --prefs)             prefslist="$2"; shift ;;   # list: mirror urls space separates, use 'single quotes'
            -h | --help)
                cat <<EOF >&2
Usage:   $progname [options] [country-code [...]]

Options: --help, -h                   This help.
         --nolocal                    Do not include mirrors from the current country.
         --offline                    Don't use internet even when a connection is available.
         --save                       Save mirrorlist to $target_def (see also option --savefile).
         --savefile=*                 File path to save the mirrorlist.
         --sequential                 Rank mirrors sequentially (slower) instead of in parallel (faster).
         --verbose, -v                Show more ranking details.
         --http                       Include the http:// mirrors.
         --rsync                      Include the rsync:// mirrors.
         --no-recommended-countries   Don't use recommended countries.
                                      Instead give one or more country codes
                                      as pararameters or use option --user-countries.
         --user-countries             Use file $user_fav_countries to give country codes.
                                      It can contain a list of codes separated by white spaces.
         --fake-country=*             Set a fake current country code (advanced).

Notes:   * Only https:// mirrors are included by default.
         * Country codes are the two-letter codes listed by command 'reflector --list-countries'.
         * Use option --save to change the existing $target_def.
         * Use option --no-recommended-countries to rank mirrors without adding recommended countries.
         * If available and option --offline is not used, internet connection is used for:
               1) fetching a list of active mirrors from the Arch web site
               2) fetching country code and name mappings from the Arch web site
               3) ranking mirrors
           Without a connection the ranking result is very likely suboptimal.
EOF
                exit 0
                ;;
        esac
        shift
    done

    while [ "$1" ] ; do
        countries_given+=("${1,,}")
        shift
    done
}

Main() {
    local -r progname=${0##*/}
    local -r pkgname=iso-create-ml
    local -r configfile=/etc/$progname.conf

    local -r sep="|"
    local OPTS="help${sep}h,http,nolocal,offline,save,savefile:,recommended-countries${sep}r,rsync,sequential"
          OPTS+=",user-countries,fake-country:,verbose${sep}v,dump-options,dump-ccs,update-supports,prefs:"

    local -r user_fav_countries=$HOME/user-countries.txt
    local -r target_def=/etc/pacman.d/mirrorlist
    local target=$target_def
    local mirrordata="${progname}-active-mirrors-arch.conf"                # support file: active Arch mirrors in various countries
    local dbfile="${progname}-country-mapping-arch.conf"                   # support file: mappings between country-code and country-name
    local has_internet_connection=yes
    local country_code=""
    local fake_country=""
    local local_country_wanted=yes                                         # include local country in the mirrorlist or not?
    local verbose=no
    local save=no
    local parallel=yes                                                     # ranking mirrors in pararallel or not?
    local protocols=(https)                                                # list of supported protocols
    local countries_given=()                                               # list of user given countries which have mirrors to include
    local countries_handled=()
    local cleanup_files=()
    local orig_args=("$@")
    local args=()
    local has_countries_option=no
    local use_recommended_countries=yes
    declare -A reflector_countries
    local update_supports=no
    # Default timeouts in seconds:
    local TIMEOUT_ACTIVE_MIRRORS=8                                 # for fetching the list of active mirrors
    local TIMEOUT_REFLECTOR_CONNECTION=5                           # --connection-timeout in reflector
    local TIMEOUT_REFLECTOR_DOWNLOAD=5                             # --download-timeout in reflector
    local TIMEOUT_COUNTRY_CODE=5                                   # for fetching the current country code (not used!)
    local prefslist=""

    Parameters "$@"

    if [ $has_internet_connection = yes ] ; then
        eos-connection-checker || has_internet_connection=no       # make sure if we are online or offline
    fi

    if ! source "$configfile" ; then
        [ $has_internet_connection = yes ] && WARN "file '$configfile' not found, using default timeout values"
    fi

    GetCountryMappings                                             # creates reflector_countries
    GetMirrorsFile                                                 # creates a list of all active mirrors

    if [ $update_supports = yes ] ; then
        echo2 "==> support files updated."
        exit 0                              # we're done here
    fi

    # # handle some special options here
    # local arg ix
    # for ((ix=0; ix < ${#orig_args[@]}; ix++)) ; do
    #     arg="${orig_args[$ix]}"
    #     case "$arg" in
    #         # handle these options only here:
    #         --fake-country)                Parameters "$arg" "${orig_args[$((ix+1))]}"; ((ix++)) ;;
    #         --fake-country=*)              Parameters "--fake-country" "${arg#*=}" ;;
    #         --user-countries)              Parameters "$arg" ;;
    #         # these options already handled and no more needed:
    #         --offline | --update-supports | --dump-options | --dump-ccs) ;;
    #         # these options will be handled in Parameters:
    #         -r | --recommended-countries)  args+=("$arg"); has_countries_option=yes ;;
    #         *)                             args+=("$arg") ;;
    #     esac
    # done

    GetCountryCode

    AddRecommendedCountries   # TODO !!!

    Main2
}

GetCountryCode() {
    [ "$fake_country" ] && { country_code="$fake_country"; return 0; }      # has fake country, use it
    [ $has_internet_connection = no ] && { country_code=ww; return 1; }     # cannot determine, so use ww

    local code
    code=$(show-location-info country 2>/dev/null)
    if [ $? = 0 ] && [ "$code" ] ; then
        country_code="${code,,}"
        return 0
    else
        WARN "fetching the current country code failed."
        country_code=ww
        return 1
    fi
}

NeedsCleanupOfSupportFiles() {
    [ -z "$(grep "^pkgname=$pkgname$" PKGBUILD 2>/dev/null)" ]   # using, not developing
}

GetMirrorsFile() {
    if FetchMirrors ; then
        NeedsCleanupOfSupportFiles && cleanup_files+=("$mirrordata")
    else
        mirrordata="/etc/${mirrordata##*/}"
    fi
}
GetCountryMappings() {
    if FetchCountries ; then
        NeedsCleanupOfSupportFiles && cleanup_files+=("$dbfile")
    else
        dbfile="/etc/${dbfile##*/}"
    fi
    # shellcheck disable=SC1090
    source "$dbfile"                                               # gets array country mappings into reflector_countries
}

Main "$@"
