#!/bin/sh
# minis-mirror — sandbox apt mirror selector with automatic fallback chain.
#
# Design goals (see docs/DESIGN_PHILOSOPHY.md: "验证优先于能力"):
#   * The sandbox must not die with its first mirror. TUNA hiccups, DNS
#     fails, a campus firewall blocks ports — apt should self-heal instead
#     of bubbling "Unable to fetch" to the user.
#   * Plain /bin/sh + curl only — no hard runtime deps. jq is used by
#     `status` when present and silently falls back to a raw JSON print
#     otherwise; the minbase rootfs does not ship jq by default.
#   * Writes only /etc/apt/sources.list.d/minis.sources (the app-managed
#     file). Never touches ubuntu.sources. The stock file is archived once
#     as minis.sources.stock so the chain can always be reset.
#
# Usage:
#   minis-mirror auto                 # probe chain, apply fastest working mirror
#   minis-mirror set <id> [id...]     # manually pin the chain order
#   minis-mirror reset                # restore stock (tuna first)
#   minis-mirror status               # JSON: current + chain + last probe
#   minis-mirror probe [--json]      # probe every mirror, print ranking
#
# Exit codes: 0 ok, 1 no mirror reachable, 2 usage error, 3 apt update failed
# after switching.

set -u

SOURCES_DIR="/etc/apt/sources.list.d"
SOURCES_FILE="$SOURCES_DIR/minis.sources"
STOCK_FILE="$SOURCES_DIR/minis.sources.stock"
STATE_DIR="/var/lib/minis-mirror"
STATE_FILE="$STATE_DIR/state.json"
LOG_FILE="$STATE_DIR/probe.log"
LOCK_FILE="$STATE_DIR/apt.lock"
APPLY_USER="${SUDO_USER:-$(id -un)}"
APPLY_GROUP="$(id -gn "$APPLY_USER" 2>/dev/null || echo "$APPLY_USER")"

# Ordered fallback chain. First entry = preferred (default: TUNA, matching
# the stock rootfs). Mirrors must serve /dists/noble/Release over HTTPS
# (HTTP is acceptable pre-CA-bootstrap but this selector always has CA trust
# available via the offline bootstrap debs, so we probe HTTPS directly).
MIRRORS="tuna:https://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ \
ustc:https://mirrors.ustc.edu.cn/ubuntu-ports/ \
aliyun:https://mirrors.aliyun.com/ubuntu-ports/ \
huawei:https://mirrors.huaweicloud.com/ubuntu-ports/ \
tencent:https://mirrors.cloud.tencent.com/ubuntu-ports/ \
nju:https://mirrors.nju.edu.cn/ubuntu-ports/ \
bfsu:https://mirrors.bfsu.edu.cn/ubuntu-ports/ \
official:http://ports.ubuntu.com/ubuntu-ports/"

PROBE_PATH="dists/noble/Release"
PROBE_TIMEOUT=5
UPDATE_TIMEOUT=120
APT_HELPER="/usr/local/lib/minis-mirror/apt-helper"

log() { printf '%s\n' "$*" >&2; }

die() { log "minis-mirror: $*"; exit "${2:-1}"; }

now_iso() { date -u +%Y-%m-%dT%H:%M:%SZ; }

chown_apply() {
    # Rootfs files are owned by the app-uid user; keep them that way after
    # we rewrite them so the app can still manage the file.
    chown "$APPLY_USER:$APPLY_GROUP" "$1" 2>/dev/null || true
}

write_state() {
    # $1 = json payload
    mkdir -p "$STATE_DIR"
    printf '%s\n' "$1" > "$STATE_FILE.tmp" && mv "$STATE_FILE.tmp" "$STATE_FILE"
    chown_apply "$STATE_FILE"
}

require_state_dir() {
    mkdir -p "$STATE_DIR"
    chown_apply "$STATE_DIR"
}

# ── stock archive ────────────────────────────────────────────────────────────
ensure_stock() {
    [ -f "$STOCK_FILE" ] && return 0
    if [ -f "$SOURCES_FILE" ]; then
        cp "$SOURCES_FILE" "$STOCK_FILE"
    else
        # Synthetic stock from the default chain head (tuna) — matches what
        # the app's MirrorSettings writes and what the rootfs ships.
        gen_sources "https://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/" > "$STOCK_FILE"
    fi
    chown_apply "$STOCK_FILE"
}

# ── template writer ──────────────────────────────────────────────────────────
gen_sources() {
    # $1 = mirror base URI (no trailing newline issues; callers append)
    cat <<EOF
Types: deb
URIs: $1
Suites: noble noble-updates noble-backports
Components: main universe restricted multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg

Types: deb
URIs: $1
Suites: noble-security
Components: main universe restricted multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
EOF
}

# ── state ────────────────────────────────────────────────────────────────────
current_uri() {
    # Current base URI from the live sources file, "" when unreadable.
    [ -f "$SOURCES_FILE" ] || { echo ""; return; }
    awk '/^URIs:/{print $2; exit}' "$SOURCES_FILE"
}

id_for_uri() {
    case "$1" in
        *tuna.tsinghua*)     echo "tuna" ;;
        *ustc.edu.cn*)       echo "ustc" ;;
        *aliyun.com*)        echo "aliyun" ;;
        *huaweicloud.com*)   echo "huawei" ;;
        *cloud.tencent.com*) echo "tencent" ;;
        *nju.edu.cn*)        echo "nju" ;;
        *bfsu.edu.cn*)       echo "bfsu" ;;
        *ports.ubuntu.com*)  echo "official" ;;
        *)                   echo "unknown" ;;
    esac
}

id_position() {
    _pos=1
    for _entry in $MIRRORS; do
        [ "${_entry%%:*}" = "$1" ] && { echo "$_pos"; return; }
        _pos=$((_pos + 1))
    done
    echo 999
}

# ── probing ──────────────────────────────────────────────────────────────────
probe_one() {
    # $1 = base url; prints "<ms> <http_code>" or "999999 000" on failure.
    _url="$1$PROBE_PATH"
    _out="$(curl -sL -o /dev/null --max-time "$PROBE_TIMEOUT" \
        -w '%{time_total} %{http_code}' "$_url" 2>/dev/null)" || true
    [ -n "$_out" ] || _out="999999 000"
    echo "$_out"
}

mirror_ok() {
    # $1 = base url; 0 when probe returns HTTP 200 within the timeout.
    _r="$(probe_one "$1")"
    _code="${_r#* }"
    [ "$_code" = "200" ]
}

cmd_probe() {
    json_out=0
    for a in "$@"; do
        [ "$a" = "--json" ] && json_out=1
    done
    require_state_dir
    : > "$LOG_FILE"
    chown_apply "$LOG_FILE"
    if [ "$json_out" = "1" ]; then printf '['; _first=1; fi
    for entry in $MIRRORS; do
        _id="${entry%%:*}"
        _url="${entry#*:}"
        _r="$(probe_one "$_url")"
        _ms="$(echo "$_r" | cut -d' ' -f1)"
        _code="$(echo "$_r" | cut -d' ' -f2)"
        case "$_code" in 200) ok=true ;; *) ok=false ;; esac
        printf '%s %8s ms  HTTP %s  %s\n' "$_id" "$_ms" "$_code" "$_url" >&2
        printf '%s|%s|%s|%s\n' "$_id" "$_ms" "$_code" "$(now_iso)" >> "$LOG_FILE"
        if [ "$json_out" = "1" ]; then
            [ $_first = 0 ] && printf ','
            printf '{"id":"%s","ms":%s,"http":%s,"url":"%s"}' \
                "$_id" "$_ms" "$_code" "$_url"
            _first=0
        fi
    done
    if [ "$json_out" = "1" ]; then printf ']\n'; fi
    return 0
}

# ── apt serialization (mirrors the app's apt lock discipline) ────────────────
acquire_lock() {
    # No flock in the minbase rootfs — use a mkdir lock (atomic on POSIX)
    # with a stale-lock protocol keyed on the recorded PID.
    require_state_dir
    if mkdir "$LOCK_FILE.d" 2>/dev/null; then
        echo $$ > "$LOCK_FILE.d/pid"
        return 0
    fi
    _pid="$(cat "$LOCK_FILE.d/pid" 2>/dev/null)"
    if [ -n "$_pid" ] && ! kill -0 "$_pid" 2>/dev/null; then
        rm -rf "$LOCK_FILE.d"
        mkdir "$LOCK_FILE.d" 2>/dev/null && { echo $$ > "$LOCK_FILE.d/pid"; return 0; }
    fi
    die "another apt/mirror operation appears to hold the lock" 1
}

release_lock() {
    rm -rf "$LOCK_FILE.d" 2>/dev/null
    return 0
}

# ── apply ────────────────────────────────────────────────────────────────────
apply_uri() {
    # $1 = uri; rewrite minis.sources atomically.
    _tmp="$SOURCES_FILE.tmp.$$"
    gen_sources "$1" > "$_tmp"
    chown_apply "$_tmp"
    mv "$_tmp" "$SOURCES_FILE"
    chown_apply "$SOURCES_FILE"
}

apt_update_ok() {
    # Run apt update with a hard timeout. Uses the app's serialized-apt
    # helper when present (it owns the dpkg/fetch locks); falls back to a
    # plain apt-get otherwise.
    #
    # apt's exit code alone is NOT a truth signal: a mirror that fails to
    # download only produces `W:`/`Err:` lines and apt still exits 0
    # (observed on noble). Treat error lines in the output as failure too.
    _out="$(mktemp)"
    if [ -x "$APT_HELPER" ]; then
        timeout "$UPDATE_TIMEOUT" "$APT_HELPER" update > "$_out" 2>&1
    else
        timeout "$UPDATE_TIMEOUT" apt-get update > "$_out" 2>&1
    fi
    _rc=$?
    if [ $_rc = 0 ] && ! grep -qE '^Err:|^E: |^W: Some index files failed' "$_out"; then
        rm -f "$_out"
        return 0
    fi
    # Keep the tail for diagnostics; apt errors are the reason we switched.
    log "apt update failed (rc=$_rc), tail of output:"
    tail -5 "$_out" >&2
    rm -f "$_out"
    return 1
}

switch_to() {
    # $1 = id, $2 = uri. Writes sources, runs apt update, records state.
    log "switching apt mirror -> $1 ($2)"
    apply_uri "$2"
    if apt_update_ok; then
        write_state "{\"current\":\"$1\",\"uri\":\"$2\",\"updated\":\"$(now_iso)\",\"auto\":true,\"ok\":true}"
        return 0
    fi
    write_state "{\"current\":\"$1\",\"uri\":\"$2\",\"updated\":\"$(now_iso)\",\"auto\":true,\"ok\":false}"
    return 1
}

cmd_auto() {
    ensure_stock
    require_state_dir
    cur="$(current_uri)"
    [ -n "$cur" ] || cur="https://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/"

    # Current mirror healthy? Probe; if 200 within timeout we are done.
    if mirror_ok "$cur"; then
        log "current mirror OK ($(id_for_uri "$cur")) — no switch"
        write_state "{\"current\":\"$(id_for_uri "$cur")\",\"uri\":\"$cur\",\"probed\":\"$(now_iso)\",\"auto\":true,\"ok\":true}"
        return 0
    fi
    log "current mirror ($(id_for_uri "$cur")) unreachable — probing fallback chain"

    for entry in $MIRRORS; do
        _id="${entry%%:*}"
        _url="${entry#*:}"
        [ "$_url" = "$cur" ] && continue
        if mirror_ok "$_url"; then
            acquire_lock
            if switch_to "$_id" "$_url"; then
                release_lock
                log "fallback complete: now on $_id"
                return 0
            fi
            release_lock
            log "$_id probe passed but apt update failed — trying next"
        fi
    done
    write_state "{\"current\":\"none\",\"uri\":\"\",\"updated\":\"$(now_iso)\",\"auto\":true,\"ok\":false}"
    die "no mirror in the chain is reachable" 1
}

cmd_set() {
    [ $# -ge 1 ] || die "usage: minis-mirror set <id> [id...]" 2
    ensure_stock
    require_state_dir
    _chain=""
    _first_uri=""
    for want in "$@"; do
        found=""
        for entry in $MIRRORS; do
            [ "${entry%%:*}" = "$want" ] && { found="$entry"; break; }
        done
        [ -n "$found" ] || die "unknown mirror id: $want (see 'minis-mirror probe')" 2
        _chain="$_chain $want"
        [ -n "$_first_uri" ] || _first_uri="${found#*:}"
    done
    acquire_lock
    log "setting chain:$_chain (head: $_first_uri)"
    _prev_uri="$(current_uri)"
    apply_uri "$_first_uri"
    if apt_update_ok; then
        write_state "{\"current\":\"$1\",\"uri\":\"$_first_uri\",\"updated\":\"$(now_iso)\",\"auto\":false,\"ok\":true,\"chain\":\"$_chain\"}"
        release_lock
        log "active mirror: $1"
        return 0
    fi
    # Roll back to the previous (known-working) sources so a bad `set`
    # never leaves the sandbox with a dead mirror.
    if [ -n "$_prev_uri" ] && [ "$_prev_uri" != "$_first_uri" ]; then
        log "rolling back to previous mirror: $_prev_uri"
        apply_uri "$_prev_uri"
        apt_update_ok || true
    fi
    release_lock
    write_state "{\"current\":\"$1\",\"uri\":\"$_first_uri\",\"updated\":\"$(now_iso)\",\"auto\":false,\"ok\":false,\"chain\":\"$_chain\"}"
    die "apt update failed on $1 — run 'minis-mirror auto' to fall back" 3
}

cmd_reset() {
    ensure_stock
    require_state_dir
    acquire_lock
    cp "$STOCK_FILE" "$SOURCES_FILE"
    chown_apply "$SOURCES_FILE"
    release_lock
    write_state "{\"current\":\"stock\",\"uri\":\"$(current_uri)\",\"updated\":\"$(now_iso)\",\"auto\":false,\"ok\":true}"
    log "restored stock sources ($(current_uri))"
}

cmd_status() {
    cur="$(current_uri)"
    if [ -f "$STATE_FILE" ]; then
        state="$(cat "$STATE_FILE")"
    else
        state="{\"current\":\"$(id_for_uri "$cur")\",\"uri\":\"$cur\",\"probed\":null,\"auto\":false,\"ok\":null}"
    fi
    if command -v jq >/dev/null 2>&1; then
        jq -n --arg uri "$cur" --argjson state "$state" \
            '{uri:$uri, id:($state.current // "unknown"), state:$state}'
    else
        printf '{"uri":"%s","state":%s}\n' "$cur" "$state"
    fi
}

usage() {
    cat <<'USAGE'
minis-mirror — apt mirror selector with automatic fallback chain

Usage:
  minis-mirror auto                 probe chain, switch to first healthy mirror (apt-verified)
  minis-mirror set <id> [id...]     pin the chain, e.g. `minis-mirror set ustc tuna official`
  minis-mirror reset                restore the stock sources file (tuna first)
  minis-mirror status               print current mirror + last probe state (JSON)
  minis-mirror probe [--json]       probe every mirror and rank by latency

Mirrors: tuna ustc aliyun huawei tencent nju bfsu official
USAGE
}

[ $# -ge 1 ] || { usage; exit 0; }
case "$1" in
    auto)   shift; cmd_auto "$@" ;;
    set)    shift; cmd_set "$@" ;;
    reset)  shift; cmd_reset "$@" ;;
    status) shift; cmd_status "$@" ;;
    probe)  shift; cmd_probe "$@" ;;
    -h|--help|help) usage ;;
    *) usage; die "unknown command: $1" 2 ;;
esac
