#!/bin/bash

echo2()   { echo -e "$@" >&2; }
printf2() { printf  "$@" >&2; }
INFO()    { eos-color info 2; echo2 "==> $progname: info: $1"; eos-color reset 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; }
DIE()     {
    Cleanup
    eos-color error 2
    local prompt="$progname: error"
    printf2 "$prompt: %s\n" "$1"
    shift
    if [ "$1" ] ; then
        printf2 "${prompt//?/ }  %s\n" "$@"
    fi
    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 '$cc_vs_name_file'"
        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\ncc_to_cname=(" > "$cc_vs_name_file"
    for ((ix=0; ix < count; ix++ )) ; do
        code=${codes[$ix]}
        country=${countries[$ix]}
        echo "    [${code,,}]='$country'" >> "$cc_vs_name_file"
    done
    echo ")" >> "$cc_vs_name_file"

    echo -e "\ncname_to_cc=(" >> "$cc_vs_name_file"
    for ((ix=0; ix < count; ix++ )) ; do
        code=${codes[$ix]}
        country=${countries[$ix]}
        echo "    [$country]=${code,,}" >> "$cc_vs_name_file"
    done
    echo ")" >> "$cc_vs_name_file"

    # echo2 "==> Created $(realpath "$cc_vs_name_file")."
    return 0
}

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

UserPrefs() {
    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 "^Server = .*$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
}

ShowRankResults() {
    local -r header="MIRROR-URL UPDATED-AT FETCH-TIME"
    local -r underline="${header//[A-Z-]/\~}"
    {
        echo "# $header"
        echo "# $underline"
        echo -e "$just_times\n"
    } | column -t
    echo ""
}

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=(--max-time $TIMEOUT_MIRROR)
        [ $verbose  = yes ]    && rankmirrors_opt+=(-v)
        [ $show_failed = yes ] && rankmirrors_opt+=(--show-failed)
        [ $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
            if [ $verbose = yes ] ; then
                local just_times just_mirrors
                just_times=$(echo "$ranked_out" | grep -E "^# ([^ ]+)[ ]+([^ ]+)")      # Skip failed mirrors.
                just_times=$(echo "$just_times" | LC_ALL=C sort -k3rb,3 -k4V,4)         # Sort: age first, then speed.
                just_mirrors=$(echo "$ranked_out" | grep -v "^# ")                      # Show mirrors.
                ShowRankResults
                UserPrefs
                echo "$just_mirrors"
            else
                UserPrefs
                echo "$ranked_out"
            fi
        } > "$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) ;;
        ar)                                                  additional_mirror_countries+=(ww us) ;;
        '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 "$cc_vs_name_file"                                               # gets array cc_to_cname
    # shellcheck disable=SC2046
    echo $(printf "%s\n" "${!cc_to_cname[@]}" | sort)
}

# CC2name() { echo "${cc_to_cname[$1]}" ; }
# Name2cc() { echo ${cname_to_cc[$1]} ; }

FakeCountry2cc() {
    local -r country="$1"
    local ctmp
    case "$country" in
        [a-z][a-z])
            ctmp="${cc_to_cname[$country]}"
            if [ "$ctmp" ] ; then
                fake_country="$country"
            else
                DIE "country code '$country' is not recognized." "See /etc/create-ml-country-mapping-arch.conf about supported countries."
            fi
            ;;
        *)
            ctmp="${cname_to_cc[$country]}"
            if [ "$ctmp" ] ; then
                fake_country="$ctmp"
            else
                DIE "country name '$country' is not recognized." "See /etc/create-ml-country-mapping-arch.conf about supported countries."
            fi
            ;;
    esac
}

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)      FakeCountry2cc "$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 ;;
            --show-failed)       show_failed=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
            --use-saved-cc)      use_saved_cc=yes ;;
            --timeout-mirror)    TIMEOUT_MIRROR="$2"; shift ;;
            -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.
         --show-failed                Show info about failed mirrors.
         --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).
         --use-saved-cc               The detected local country code is saved to a file $stored_country_file.
                                      This option allows using it when local country detection fails (advanced).
         --savefile=*                 File path to save the mirrorlist (advanced).
         --timeout-mirror='seconds'   Max time to rank before timing out a mirror. Default: $timeout_mirror_def.
         --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 internet connection is available and option --offline is not used, the 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 can be suboptimal or even problematic.
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 stored_country_file="/etc/${progname}-saved-cc.conf"

    local -r sep="&"
    local OPTS=""
    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"
    OPTS+=",use-saved-cc,timeout-mirror:,countries:${sep}c:,show-failed"

    local -r user_fav_countries=$HOME/user-countries.txt
    local target_msg=""
    local target_def=/etc/pacman.d/mirrorlist
    if [ -L "$target_def" ] ; then
        target_def=$(/bin/ls -l "$target_def" | awk '{print $NF}')         # /etc/pacman.d/mirrorlist is symlink, so update the target, not the link name
        [ "$target_def" ] || DIE "symlink target of /etc/pacman.d/mirrorlist was not found."
        target_msg="default mirrorlist file is now '$target_def'."
    fi
    local target="$target_def"
    if [ -w . ] ; then
        local mirrordata="${progname}-active-mirrors-arch.conf"                # support file: active Arch mirrors in various countries
        local cc_vs_name_file="${progname}-country-mapping-arch.conf"          # support file: mappings between country-code and country-name
    else
        local mirrordata="$HOME/.${progname}-active-mirrors-arch.conf"                # support file: active Arch mirrors in various countries
        local cc_vs_name_file="$HOME/.${progname}-country-mapping-arch.conf"          # support file: mappings between country-code and country-name
    fi
    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 show_failed=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
    local use_saved_cc=no
    declare -A cc_to_cname                                         # country code -> country name
    declare -A cname_to_cc                                         # country name -> country code
    local update_supports=no
    local -r timeout_mirror_def=10
    local TIMEOUT_MIRROR=$timeout_mirror_def                       # timeout for each mirror, see --timeout-mirror
    # 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 CREATE_ML_OPTIONS=()                                     # user can include options into this array in $configfile
    local prefslist=""
    local issue_test=no

    GetCountryMappings                                             # creates cc_to_cname

    if ! source "$configfile" "$@" ; then
        [ $has_internet_connection = yes ] && INFO "file '$configfile' not found"
    fi

    Parameters "${CREATE_ML_OPTIONS[@]}" "$@"

    if [ "$target_msg" ] ; then
        INFO "$target_msg"
    fi

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

    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

    if [ "$use_recommended_countries" = "yes" ] ; then
        [ -r "$recommended_countries" ] && source "$recommended_countries"
        AddRecommendedCountries
    fi

    Main2
}

SaveCountryCode() {
    # Save detected country code to a file only if it is not a
    # - fake country
    # - previously saved country
    local cc="$1"
    [ "$cc" = "$fake_country" ] && return
    local -r country_stored="$(cat "$stored_country_file" 2>/dev/null)"
    if [ "$cc" != "$country_stored" ] ; then
        INFO "saving country code '$cc' to file $stored_country_file."
        echo "$cc" | sudo tee "$stored_country_file" > /dev/null
    fi
}
GetCountryCode() {
    [ "$fake_country" ] && { country_code="$fake_country"; return 0; }      # has fake country, use it
    if [ $has_internet_connection = no ] ; then
        country_code=ww
        return 1     # cannot determine, so use ww
    fi

    local code
    code=$(show-location-info country 2>/dev/null)
    if [ $? = 0 ] && [ "$code" ] ; then
        country_code="${code,,}"
        SaveCountryCode "$country_code"
        return 0
    else
        local -r country_stored="$(cat "$stored_country_file" 2>/dev/null)"
        if [ $use_saved_cc = yes ] && [ "$country_stored" ] ; then
            INFO "using the previously saved country code '$country_stored'"
            country_code="$country_stored"
            return 1
        fi
        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+=("$cc_vs_name_file")
    else
        cc_vs_name_file="/etc/${cc_vs_name_file##*/}"
    fi
    # shellcheck disable=SC1090
    source "$cc_vs_name_file"                                               # gets array country mappings into cc_to_cname
}

Main "$@"
