#!/usr/bin/env bash
#
# aptly-mirror.sh — Create and maintain aptly mirrors for Debian and Ubuntu repositories.
#
# Mirrors:
#   Debian:  11 (bullseye), 12 (bookworm), 13 (trixie)
#   Ubuntu:  20.04 (focal), 22.04 (jammy), 24.04 (noble), 26.04 (roaring)
#
# Usage:
#   ./aptly-mirror.sh create      — Create all mirrors (first-time setup)
#   ./aptly-mirror.sh update      — Update all mirrors, snapshot, merge, and publish
#   ./aptly-mirror.sh publish     — Snapshot current state and publish/switch
#   ./aptly-mirror.sh cleanup     — Drop old snapshots and run db cleanup
#   ./aptly-mirror.sh list        — List all mirrors, snapshots, and published repos
#
set -euo pipefail

# ---------------------------------------------------------------------------
# Load configuration
# ---------------------------------------------------------------------------

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Config file search order:
#   1. APTLY_MIRROR_CONF environment variable
#   2. ./aptly-mirror.conf (next to the script)
#   3. /etc/aptly-mirror.conf

find_config() {
    if [[ -n "${APTLY_MIRROR_CONF:-}" && -r "$APTLY_MIRROR_CONF" ]]; then
        echo "$APTLY_MIRROR_CONF"
    elif [[ -r "${SCRIPT_DIR}/aptly-mirror.conf" ]]; then
        echo "${SCRIPT_DIR}/aptly-mirror.conf"
    elif [[ -r "/etc/aptly-mirror.conf" ]]; then
        echo "/etc/aptly-mirror.conf"
    else
        echo ""
    fi
}

CONFIG_FILE=""
CONFIG_FILE=$(find_config)

if [[ -z "$CONFIG_FILE" ]]; then
    echo "Error: No config file found." >&2
    echo "Searched: \$APTLY_MIRROR_CONF, ${SCRIPT_DIR}/aptly-mirror.conf, /etc/aptly-mirror.conf" >&2
    echo "Create one from aptly-mirror.conf.example or set APTLY_MIRROR_CONF." >&2
    exit 1
else
    # shellcheck source=aptly-mirror.conf
    source "$CONFIG_FILE"
fi

DATE=$(date +%Y%m%d%H%M%S)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

log() {
    local msg="$*"
    echo "$msg"
    logger -t aptly-mirror "$msg"
}

die() {
    log "FATAL: $*"
    exit 1
}

run_aptly() {
    log "  -> aptly $*"
    if ! aptly "$@" 2>&1; then
        log "  !! aptly $1 failed (exit $?)"
        return 1
    fi
}

mirror_exists() {
    aptly mirror show "$1" &>/dev/null
}

snapshot_exists() {
    aptly snapshot show "$1" &>/dev/null
}

published_exists() {
    # $1 = distribution, $2 = prefix (default ".")
    local dist="$1" prefix="${2:-.}"
    aptly publish list -raw 2>/dev/null | grep -qE "^${prefix} ${dist}$"
}

common_mirror_flags() {
    local flags="-architectures=${ARCHITECTURES}"
    if [[ -n "$GPG_KEYRING" ]]; then
        flags+=" -keyring=${GPG_KEYRING}"
    fi
    if (( IGNORE_SIGNATURES )); then
        flags+=" -ignore-signatures"
    fi
    echo "$flags"
}

common_update_flags() {
    local flags=""
    if (( DOWNLOAD_LIMIT > 0 )); then
        flags+=" -download-limit=${DOWNLOAD_LIMIT}"
    fi
    if (( IGNORE_SIGNATURES )); then
        flags+=" -ignore-signatures"
    fi
    echo "$flags"
}

publish_flags() {
    local flags=""
    if (( SKIP_SIGNING )); then
        flags+=" -skip-signing"
    fi
    echo "$flags"
}

# ---------------------------------------------------------------------------
# get_debian_components <codename>
# ---------------------------------------------------------------------------
get_debian_components() {
    local codename="$1"
    case "$codename" in
        bullseye) echo "$DEBIAN_COMPONENTS_LEGACY" ;;
        *)        echo "$DEBIAN_COMPONENTS_MODERN" ;;
    esac
}

get_debian_security_components() {
    local codename="$1"
    case "$codename" in
        bullseye) echo "$DEBIAN_SECURITY_COMPONENTS_LEGACY" ;;
        *)        echo "$DEBIAN_SECURITY_COMPONENTS_MODERN" ;;
    esac
}

create_release_mirrors() {
    local family="$1" codename="$2" version="$3"
    local main_mirror="$4" security_mirror="$5"
    local components="$6" security_components="$7"
    local suite name suite_name suite_components source_url

    log "Creating ${family^} ${version} (${codename}) mirrors..."

    for suite in main updates security; do
        name="${family}-${codename}-${suite}"
        if mirror_exists "$name"; then
            log "  Mirror $name already exists, skipping."
            continue
        fi

        case "$suite" in
            main)
                source_url="$main_mirror"
                suite_name="$codename"
                suite_components="$components"
                ;;
            updates)
                source_url="$main_mirror"
                suite_name="${codename}-updates"
                suite_components="$components"
                ;;
            security)
                source_url="$security_mirror"
                suite_name="${codename}-security"
                suite_components="$security_components"
                ;;
        esac

        run_aptly mirror create $(common_mirror_flags) \
            "$name" "$source_url" "$suite_name" $suite_components
    done
}

snapshot_and_publish_release() {
    local family="$1" codename="$2" version="$3"
    local pub_flags="$4"
    local suite snap_name snap_merged
    local -a snaps=()

    log "Snapshotting ${family^} ${version} (${codename})..."

    for suite in main updates security; do
        snap_name="${family}-${codename}-${suite}-${DATE}"
        run_aptly snapshot create "$snap_name" from mirror "${family}-${codename}-${suite}"
        snaps+=("$snap_name")
    done

    snap_merged="${family}-${codename}-merged-${DATE}"
    log "Merging snapshots for ${family^} ${codename}..."
    run_aptly snapshot merge -latest "$snap_merged" "${snaps[@]}"

    if published_exists "$codename" "$family"; then
        log "Switching published ${family^} ${codename} to ${snap_merged}..."
        run_aptly publish switch $pub_flags "$codename" "$family" "$snap_merged"
    else
        log "Publishing ${family^} ${codename} for the first time..."
        run_aptly publish snapshot $pub_flags \
            -distribution="$codename" -architectures="$ARCHITECTURES" \
            "$snap_merged" "$family"
    fi
}

# ---------------------------------------------------------------------------
# CREATE — set up all mirrors
# ---------------------------------------------------------------------------

create_debian_mirrors() {
    local codename version components security_components
    for codename in "${!DEBIAN_RELEASES[@]}"; do
        version="${DEBIAN_RELEASES[$codename]}"
        components=$(get_debian_components "$codename")
        security_components=$(get_debian_security_components "$codename")

        create_release_mirrors \
            "debian" "$codename" "$version" \
            "$DEBIAN_MIRROR" "$DEBIAN_SECURITY_MIRROR" \
            "$components" "$security_components"
    done
}

create_ubuntu_mirrors() {
    local codename version
    for codename in "${!UBUNTU_RELEASES[@]}"; do
        version="${UBUNTU_RELEASES[$codename]}"

        create_release_mirrors \
            "ubuntu" "$codename" "$version" \
            "$UBUNTU_MIRROR" "$UBUNTU_SECURITY_MIRROR" \
            "$UBUNTU_COMPONENTS" "$UBUNTU_COMPONENTS"
    done
}

do_create() {
    log "=== Creating all mirrors ==="
    create_debian_mirrors
    create_ubuntu_mirrors
    log "=== Mirror creation complete ==="
}

# ---------------------------------------------------------------------------
# UPDATE — fetch latest package metadata and files
# ---------------------------------------------------------------------------

update_all_mirrors() {
    local mirror_name
    log "=== Updating all mirrors ==="

    while IFS= read -r mirror_name; do
        [[ -z "$mirror_name" ]] && continue
        log "Updating mirror: ${mirror_name}"
        run_aptly mirror update $(common_update_flags) "$mirror_name" || \
            log "  !! Warning: failed to update ${mirror_name}, continuing..."
    done < <(aptly mirror list -raw 2>/dev/null)

    log "=== Mirror updates complete ==="
}

# ---------------------------------------------------------------------------
# SNAPSHOT & PUBLISH — create dated snapshots, merge, and publish/switch
# ---------------------------------------------------------------------------

snapshot_and_publish_debian() {
    local codename version
    local pub_flags
    pub_flags=$(publish_flags)

    for codename in "${!DEBIAN_RELEASES[@]}"; do
        version="${DEBIAN_RELEASES[$codename]}"

        snapshot_and_publish_release "debian" "$codename" "$version" "$pub_flags"
    done
}

snapshot_and_publish_ubuntu() {
    local codename version
    local pub_flags
    pub_flags=$(publish_flags)

    for codename in "${!UBUNTU_RELEASES[@]}"; do
        version="${UBUNTU_RELEASES[$codename]}"

        snapshot_and_publish_release "ubuntu" "$codename" "$version" "$pub_flags"
    done
}

do_publish() {
    log "=== Snapshot and publish ==="
    snapshot_and_publish_debian
    snapshot_and_publish_ubuntu
    log "=== Publish complete ==="
    log ""
    log "Repositories are published under ~/.aptly/public/"
    log ""
    log "Client sources.list entries:"
    for codename in "${!DEBIAN_RELEASES[@]}"; do
        log "  deb http://<server>/debian ${codename} $(get_debian_components "$codename")"
    done
    for codename in "${!UBUNTU_RELEASES[@]}"; do
        log "  deb http://<server>/ubuntu ${codename} ${UBUNTU_COMPONENTS}"
    done
}

# ---------------------------------------------------------------------------
# UPDATE (full) — update mirrors then snapshot & publish
# ---------------------------------------------------------------------------

do_update() {
    update_all_mirrors
    do_publish
}

# ---------------------------------------------------------------------------
# CLEANUP — drop old snapshots, run db cleanup
# ---------------------------------------------------------------------------

do_cleanup() {
    log "=== Cleanup ==="

    # Collect all snapshot names grouped by prefix (everything before the timestamp)
    declare -A snap_groups
    while IFS= read -r snap; do
        [[ -z "$snap" ]] && continue
        # Snapshot names follow: <prefix>-YYYYMMDDHHMMSS
        # Strip the 14-digit timestamp suffix to get the group key
        local prefix="${snap%-[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9]}"
        if [[ "$prefix" == "$snap" ]]; then
            # No timestamp suffix — skip (manually created snapshot)
            continue
        fi
        snap_groups["$prefix"]+="${snap}"$'\n'
    done < <(aptly snapshot list -raw -sort=name 2>/dev/null)

    local dropped=0
    for prefix in "${!snap_groups[@]}"; do
        # Sort snapshots by name (timestamp order) and keep the newest KEEP_SNAPSHOTS
        local snaps
        snaps=$(echo -n "${snap_groups[$prefix]}" | sort)
        local count
        count=$(echo "$snaps" | wc -l)
        local to_drop=$(( count - KEEP_SNAPSHOTS ))

        if (( to_drop > 0 )); then
            echo "$snaps" | head -n "$to_drop" | while IFS= read -r old_snap; do
                [[ -z "$old_snap" ]] && continue
                log "Dropping snapshot: ${old_snap}"
                run_aptly snapshot drop "$old_snap" || \
                    log "  !! Could not drop ${old_snap} (may be published or referenced)"
                (( dropped++ )) || true
            done
        fi
    done

    log "Running database cleanup..."
    run_aptly db cleanup

    log "=== Cleanup complete ==="
}

# ---------------------------------------------------------------------------
# LIST — show current state
# ---------------------------------------------------------------------------

do_list() {
    echo "==============================="
    echo " Mirrors"
    echo "==============================="
    aptly mirror list 2>/dev/null || true
    echo ""
    echo "==============================="
    echo " Snapshots"
    echo "==============================="
    aptly snapshot list -sort=time 2>/dev/null || true
    echo ""
    echo "==============================="
    echo " Published Repositories"
    echo "==============================="
    aptly publish list 2>/dev/null || true
}

# ---------------------------------------------------------------------------
# GPG key import helper
# ---------------------------------------------------------------------------

do_import_keys() {
    log "=== Importing repository signing keys ==="

    local keyserver="${GPG_KEYSERVER:-hkps://keyserver.ubuntu.com}"

    log "Importing Debian GPG keys..."
    for key in "${DEBIAN_GPG_KEYS[@]}"; do
        gpg --no-default-keyring --keyring trustedkeys.gpg \
            --keyserver "$keyserver" --recv-keys "$key" 2>&1 || \
            log "  !! Warning: could not import key ${key}"
    done

    log "Importing Ubuntu GPG keys..."
    for key in "${UBUNTU_GPG_KEYS[@]}"; do
        gpg --no-default-keyring --keyring trustedkeys.gpg \
            --keyserver "$keyserver" --recv-keys "$key" 2>&1 || \
            log "  !! Warning: could not import key ${key}"
    done

    log "=== Key import complete ==="
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

usage() {
    cat <<USAGE
Usage: aptly-mirror.sh <command>

Commands:
    create        Create all mirrors (first-time setup)
    update        Update mirrors, snapshot, merge, and publish
    publish       Snapshot current mirror state and publish/switch
    cleanup       Drop old snapshots and run aptly db cleanup
    list          List all mirrors, snapshots, and published repos
    import-keys   Import GPG keys for Debian and Ubuntu repositories

Configuration:
    The config file is searched in this order:
      1. \$APTLY_MIRROR_CONF environment variable
      2. ${SCRIPT_DIR}/aptly-mirror.conf
      3. /etc/aptly-mirror.conf
USAGE
}

main() {
    if (( $# < 1 )); then
        usage
        exit 1
    fi

    local cmd="$1"; shift
    case "$cmd" in
        create)      do_create ;;
        update)      do_update ;;
        publish)     do_publish ;;
        cleanup)     do_cleanup ;;
        list)        do_list ;;
        import-keys) do_import_keys ;;
        -h|--help|help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown command: $cmd" >&2
            usage
            exit 1
            ;;
    esac
}

main "$@"