#!/bin/bash

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

ASSERT_DIE()  { "$@" || DIE "'$*' failed."; }
ASSERT_WARN() { "$@" && return 0; WARN "'$*' failed."; return 1; }
IssueTest()   { eos-color warning 2; echo2 "==>" "$@"; eos-color reset 2; }

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 $(realpath "$file")."
        rm -f "$file".bak
    else
        rm -f "$fetched"
        if [ ! -e "$file" ] ; then
            WARN "could not $operation '$(realpath "$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 $(realpath "$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" "${additional_mirror_countries[@]}" | grep "^$item$" >/dev/null || additional_mirror_countries+=("$item") ;;
                *)          WARN "$user_fav_countries: country code '$item2' ignored" ;;
            esac
        done
    else
        WARN "cannot read file '$user_fav_countries'"
    fi
}
UserCountriesFromCommand() {
    local list="$1" cc                  # list items can separated by spaces ( ), commas (,), or pipes (|)
    for cc in ${list//[,|]/ } ; do
        additional_mirror_countries+=("${cc,,}")
    done
}

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"
    [ "${orig_args[0]}" ] && printf " '%s'" "${orig_args[@]}"
    echo -e "\n"

    if [ "$prefslist" ] ; then
        echo -e "####### Mirrors by user preference >>>>"
        local tmp
        for item in $prefslist ; do
            if [ "$ranked_out" ] ; then
                tmp=$(echo "$ranked_out" | grep "$item" | awk '{print $NF}')
                if [ -z "$tmp" ] ; then
                    WARN "--prefs arguments: no match for '$item' within ranked mirrors"
                    continue
                fi
                item="$tmp"
            else
                item=$(grep "$item" $mirrordata | awk '{print $NF}')
            fi
            echo "Server = $item"
        done
        echo -e "####### Mirrors by user preference <<<<\n"
    fi
}

Main2() {
    local tmplist mirrorlist ranked_out=""
    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 "${additional_mirror_countries[@]}" ; 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."
        {
            # shellcheck disable=SC2016
            ranked_out=$(create-ml-rankmirrors-arch "${rankmirrors_opt[@]}" "$tmplist" | column -t -s'|' | sed 's|/lastupdate|/$repo/os/$arch|')
            Header
            echo "$ranked_out"
        } > "$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
        echo2 "==> Updating $target"
        if [ "$target" = $target_def ] ; then
            sudo rm -f "$target.bak"
            sudo mv "$target" "$target.bak"
            sudo cp -i "$mirrorlist" "$target"
            sudo chmod go+r "$target"
        else
            if [ -w "${target%/*}" ] ; then
                cp "$mirrorlist" "$target"
            else
                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.
    # In offline mode we fall back to the last branch of the case..esac below because the $current_country
    # is unknown.
    case "$country_code" in
        al|at|be|cz|dk|ee|es|fr|gb|gr|hr|it|nl|no|pl|pt|se)  additional_mirror_countries+=(ww de fr) ;;
        'fi')                                                additional_mirror_countries+=(ww se) ;;
        de|us)                                               ;;
        cn|ru)                                               protocols+=(http) ;;
        ca)                                                  additional_mirror_countries+=(us) ;;
        au)                                                  additional_mirror_countries+=(ww) ;;
        br|cl|co|mx)                                         additional_mirror_countries+=(ww us); protocols+=(http) ;;
        tw)                                                  additional_mirror_countries+=(ww sg kr) ;;
        kr)                                                  additional_mirror_countries+=(ww sg tw) ;;
        sg)                                                  additional_mirror_countries+=(ww kr tw) ;;
        'in')                                                additional_mirror_countries+=(ww de us) ;;
        *)                                                   additional_mirror_countries+=(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|(${sep}[a-zA-Z][:]*)||g")"
    sopts="$(echo "$OPTS" | sed -E "s|[^$sep]*$sep([a-zA-Z][:]*)[^$sep]*|\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 ;;
            -c | --countries)    UserCountriesFromCommand "$2"; shift ;;
            --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 regexps, space separated, use 'single quotes'
            --issue-test)        issue_test=yes ;;          # for internal testing about potential problems
            -h | --help)
                cat <<EOF >&2
Usage:   $progname [options]

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).
         --sequential                 Rank mirrors sequentially (slower) instead of in parallel (faster).
         --verbose, -v                Show more ranking details.
         --http                       Include the http:// mirrors.
         --countries, -c              Give a list of country codes that will be used as additional mirror countries.
                                      The list items can be separated by commas or spaces.
                                      Examples:
                                           -c ca,fr,tw
                                           -c 'ca fr tw'    # note: quotes required here
         --user-countries             Use file $user_fav_countries to give country codes.
                                      It can contain a list of country codes separated by white spaces.
         --no-recommended-countries   Don't use recommended countries.
                                      Instead you may give one or more country codes
                                      with option --countries or option --user-countries.
         --fake-country=*             Set a fake current country code (advanced).
         --savefile=*                 File path to save the mirrorlist (advanced).
         --prefs='regexp-list'        List of preferred mirrors as regexps for grep (advanced).
                                      Use single quotes around the list. Separate items by a space.

Notes:   * Country codes are the two-letter codes as listed by command 'reflector --list-countries'.
           Note: a special code 'ww' can be used too. It is a compact list of 'Worldwide' mirrors.
         * By default only https:// mirrors are included.
         * 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
#         --rsync                      Include the rsync:// mirrors (experimental, not fully implemented).
                exit 0
                ;;
        esac
        shift
    done
}

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

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

    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 additional_mirror_countries=()                                   # 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=""
    local issue_test=no

    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

    [ -r "$recommended_countries" ] && source "$recommended_countries"
    AddRecommendedCountries

    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 "$@"
