#!/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. -c <path> command-line flag
#   2. APTLY_MIRROR_CONF environment variable
#   3. ./aptly-mirror.conf (next to the script)
#   4. /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=""

# Parse flags before subcommand
while [[ "${1:-}" == -* ]]; do
    case "$1" in
        -c) CONFIG_FILE="$2"; shift 2 ;;
        -h|--help) CONFIG_FILE="__help__"; break ;;
        *) break ;;
    esac
done

if [[ -z "$CONFIG_FILE" ]]; then
    CONFIG_FILE=$(find_config)
fi

if [[ "$CONFIG_FILE" == "__help__" ]]; then
    # Defer to main() which will show usage
    :
elif [[ -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 pass -c <path>." >&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 — 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")

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

        local name="debian-${codename}-main"
        if ! mirror_exists "$name"; then
            run_aptly mirror create $(common_mirror_flags) \
                "$name" "$DEBIAN_MIRROR" "$codename" $components
        else
            log "  Mirror $name already exists, skipping."
        fi

        name="debian-${codename}-updates"
        if ! mirror_exists "$name"; then
            run_aptly mirror create $(common_mirror_flags) \
                "$name" "$DEBIAN_MIRROR" "${codename}-updates" $components
        else
            log "  Mirror $name already exists, skipping."
        fi

        name="debian-${codename}-security"
        if ! mirror_exists "$name"; then
            run_aptly mirror create $(common_mirror_flags) \
                "$name" "$DEBIAN_SECURITY_MIRROR" "${codename}-security" $security_components
        else
            log "  Mirror $name already exists, skipping."
        fi
    done
}

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

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

        local name="ubuntu-${codename}-main"
        if ! mirror_exists "$name"; then
            run_aptly mirror create $(common_mirror_flags) \
                "$name" "$UBUNTU_MIRROR" "$codename" $UBUNTU_COMPONENTS
        else
            log "  Mirror $name already exists, skipping."
        fi

        name="ubuntu-${codename}-updates"
        if ! mirror_exists "$name"; then
            run_aptly mirror create $(common_mirror_flags) \
                "$name" "$UBUNTU_MIRROR" "${codename}-updates" $UBUNTU_COMPONENTS
        else
            log "  Mirror $name already exists, skipping."
        fi

        name="ubuntu-${codename}-security"
        if ! mirror_exists "$name"; then
            run_aptly mirror create $(common_mirror_flags) \
                "$name" "$UBUNTU_SECURITY_MIRROR" "${codename}-security" $UBUNTU_COMPONENTS
        else
            log "  Mirror $name already exists, skipping."
        fi
    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 snap_main snap_updates snap_security snap_merged
    local pub_flags
    pub_flags=$(publish_flags)

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

        log "Snapshotting Debian ${version} (${codename})..."

        snap_main="debian-${codename}-main-${DATE}"
        snap_updates="debian-${codename}-updates-${DATE}"
        snap_security="debian-${codename}-security-${DATE}"
        snap_merged="debian-${codename}-merged-${DATE}"

        run_aptly snapshot create "$snap_main" from mirror "debian-${codename}-main"
        run_aptly snapshot create "$snap_updates" from mirror "debian-${codename}-updates"
        run_aptly snapshot create "$snap_security" from mirror "debian-${codename}-security"

        log "Merging snapshots for Debian ${codename}..."
        run_aptly snapshot merge -latest \
            "$snap_merged" "$snap_main" "$snap_updates" "$snap_security"

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

snapshot_and_publish_ubuntu() {
    local codename version snap_main snap_updates snap_security snap_merged
    local pub_flags
    pub_flags=$(publish_flags)

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

        log "Snapshotting Ubuntu ${version} (${codename})..."

        snap_main="ubuntu-${codename}-main-${DATE}"
        snap_updates="ubuntu-${codename}-updates-${DATE}"
        snap_security="ubuntu-${codename}-security-${DATE}"
        snap_merged="ubuntu-${codename}-merged-${DATE}"

        run_aptly snapshot create "$snap_main" from mirror "ubuntu-${codename}-main"
        run_aptly snapshot create "$snap_updates" from mirror "ubuntu-${codename}-updates"
        run_aptly snapshot create "$snap_security" from mirror "ubuntu-${codename}-security"

        log "Merging snapshots for Ubuntu ${codename}..."
        run_aptly snapshot merge -latest \
            "$snap_merged" "$snap_main" "$snap_updates" "$snap_security"

        local prefix="ubuntu"
        if published_exists "$codename" "$prefix"; then
            log "Switching published Ubuntu ${codename} to ${snap_merged}..."
            run_aptly publish switch $pub_flags "$codename" "$prefix" "$snap_merged"
        else
            log "Publishing Ubuntu ${codename} for the first time..."
            run_aptly publish snapshot $pub_flags \
                -distribution="$codename" -architectures="$ARCHITECTURES" \
                "$snap_merged" "$prefix"
        fi
    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 [-c <config>] <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

Options:
    -c <path>     Path to config file (default: auto-detected)

Configuration:
    The config file is searched in this order:
      1. -c <path> flag
      2. \$APTLY_MIRROR_CONF environment variable
      3. ${SCRIPT_DIR}/aptly-mirror.conf
      4. /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 "$@"