#!/bin/bash -e
# Usage: mkrel [uv]              - build for PLATFORM or native, optionally runs
#                                  `fossil uv add`
#        mkrel farm              - build all platforms, extract, fallback, fossil uv
#        mkrel container         - build and push container image only
#        mkrel platform PLAT     - rebuild one linux/* platform (container + dist)
#                                  e.g. TAG=0.28.0 ./tools/mkrel platform linux/arm64
#
# Env: PLATFORM (e.g. darwin/amd64), TAG (or VERSION for backward compat)
#      PUSH=1 on platform mode: amend manifest and push to registry
#      UV=1 on platform mode: fossil uv add/sync changed artifacts (default)
#      AMD64_CONNECTION: podman system connection for remote amd64 builds (default amd64)
#      FARM_DIR: OpenSSH control socket and Podman socket forward (default
#                ~/.local/run/farm). The amd64 connection is that local unix
#                socket. Podman's Go SSH client shuffles keys and trips sshd's
#                auth-retry limit, so OpenSSH pins ~/.ssh/farm.pub and forwards
#                the remote socket instead.
set -e

mode=$1
TAG="${TAG:-$VERSION}"
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
DIST="$ROOT/dist"
AMD64_CONNECTION="${AMD64_CONNECTION:-amd64}"
FARM_DIR="${FARM_DIR:-$HOME/.local/run/farm}"
FARM_HOST=amd64

image="quay.io/tangent/rejig"

expected_artifacts=(
    rejig-linux-amd64.tar.xz
    rejig-linux-arm64.tar.xz
    rejig-macos-amd64.tar.xz
    rejig-macos-arm64.tar.xz
    rejig-windows-amd64.zip
    rejig-windows-arm64.zip
)

binary_version() {
    local bin=$1 ver
    ver=$(strings "$bin" 2>/dev/null | grep -E '^[0-9]+\.[0-9]+\.[0-9]+$' | head -1)
    if [ -z "$ver" ]; then
        ver=$(strings "$bin" 2>/dev/null |
            grep -oE 'mode\.Version=[0-9]+\.[0-9]+\.[0-9]+' | head -1 | cut -d= -f2)
    fi
    echo "$ver"
}

artifact_os() {
    case "$(basename "$1")" in
    rejig-linux-*)   echo linux ;;
    rejig-macos-*)   echo darwin ;;
    rejig-windows-*) echo windows ;;
    *)               echo unknown ;;
    esac
}

artifact_arch() {
    case "$(basename "$1")" in
    *-amd64.*) echo amd64 ;;
    *-arm64.*) echo arm64 ;;
    *)         echo unknown ;;
    esac
}

binary_arch() {
    local bin=$1 os=$2 info
    info=$(file -b "$bin")
    case "$os" in
    linux)
        case "$info" in
        *x86-64*)                 echo amd64 ;;
        *ARM\ aarch64*|*aarch64*) echo arm64 ;;
        esac
        ;;
    darwin)
        case "$info" in
        *x86_64*) echo amd64 ;;
        *arm64*)  echo arm64 ;;
        esac
        ;;
    windows)
        case "$info" in
        *x86-64*)          echo amd64 ;;
        *Aarch64*|*ARM64*) echo arm64 ;;
        esac
        ;;
    esac
}

extract_artifact_binary() {
    local art=$1 dest=$2
    if [[ $art == *.zip ]]; then
        unzip -p "$art" rejig.exe >"$dest" 2>/dev/null
    else
        tar xOf "$art" rejig >"$dest" 2>/dev/null
    fi
}

verify_binary() {
    local bin=$1 os=$2 want_arch=$3
    local ver got_arch

    ver=$(binary_version "$bin")
    if [ "$ver" != "$TAG" ]; then
        echo "mkrel: binary has version ${ver:-?}, expected $TAG" >&2
        return 1
    fi
    got_arch=$(binary_arch "$bin" "$os")
    if [ "$got_arch" != "$want_arch" ]; then
        echo "mkrel: binary is ${got_arch:-?}, expected $want_arch ($(file -b "$bin"))" >&2
        return 1
    fi
}

verify_dist_artifact() {
    local art=$1 os arch tmp

    os=$(artifact_os "$art")
    arch=$(artifact_arch "$art")
    tmp=$(mktemp)
    if ! extract_artifact_binary "$art" "$tmp"; then
        rm -f "$tmp"
        echo "mkrel: $(basename "$art") is unreadable" >&2
        return 1
    fi
    if ! verify_binary "$tmp" "$os" "$arch"; then
        echo "mkrel: $(basename "$art") failed verification" >&2
        rm -f "$tmp"
        return 1
    fi
    rm -f "$tmp"
}

verify_all_dist() {
    local f missing=0
    for f in "${expected_artifacts[@]}"; do
        if [ ! -f "$DIST/$f" ]; then
            echo "mkrel: missing dist/$f" >&2
            missing=1
        elif ! verify_dist_artifact "$DIST/$f"; then
            missing=1
        fi
    done
    return $missing
}

arch_tag() {
    echo "${image}:${TAG}-${1}"
}

amd64_local_tag() {
    echo "localhost/rejig-amd64"
}

# Extract a file from a single-arch image (not a manifest list).
copy_from_image() {
    local ref=$1 member=$2 dest=$3
    local tmpdir layer

    tmpdir=$(mktemp -d)
    podman save -q "$ref" -o "$tmpdir/img.tar"
    tar -xf "$tmpdir/img.tar" -C "$tmpdir"
    for layer in "$tmpdir"/*/layer.tar; do
        [ -f "$layer" ] || continue
        if tar -tf "$layer" 2>/dev/null | grep -qx "$member"; then
            tar -xOf "$layer" "$member" >"$dest"
            rm -rf "$tmpdir"
            return 0
        fi
    done
    rm -rf "$tmpdir"
    return 1
}

verify_container_image() {
    local ref=$1 want_arch=$2
    local img_arch tmp

    img_arch=$(podman image inspect --format '{{.Architecture}}' "$ref")
    if [ "$img_arch" != "$want_arch" ]; then
        echo "mkrel: image $ref is $img_arch, expected $want_arch" >&2
        return 1
    fi

    tmp=$(mktemp)
    if ! copy_from_image "$ref" rejig "$tmp"; then
        echo "mkrel: image $ref has no /rejig" >&2
        rm -f "$tmp"
        return 1
    fi
    if ! verify_binary "$tmp" linux "$want_arch"; then
        echo "mkrel: image $ref failed verification" >&2
        rm -f "$tmp"
        return 1
    fi
    rm -f "$tmp"
}

free_image_tag() {
    local ref=$1
    podman untag "$ref" 2>/dev/null || true
    podman manifest rm "$ref" 2>/dev/null || true
}

extract_linux_artifacts() {
    local arch=$1 ref

    ref=${ref:-$(arch_tag "$arch")}

    if copy_from_image "$ref" rejig "$DIST/rejig"; then
        (cd "$DIST" && tar cJf "rejig-linux-$arch.tar.xz" rejig && rm rejig)
        if ! verify_dist_artifact "$DIST/rejig-linux-$arch.tar.xz"; then
            rm -f "$DIST/rejig-linux-$arch.tar.xz"
        fi
    fi
    if copy_from_image "$ref" rejig.exe "$DIST/rejig.exe" 2>/dev/null; then
        (cd "$DIST" && zip -q "rejig-windows-$arch.zip" rejig.exe && rm rejig.exe)
        if ! verify_dist_artifact "$DIST/rejig-windows-$arch.zip"; then
            rm -f "$DIST/rejig-windows-$arch.zip"
        fi
    fi
}

fetch_amd64_image() {
    local local_tag dest authfile id
    local_tag=$(amd64_local_tag)
    dest=$(arch_tag amd64)
    authfile="${CONTAINERS_AUTHFILE:-$HOME/.config/containers/auth.json}"

    if ! podman --connection "$AMD64_CONNECTION" save "$local_tag" 2>/dev/null |
        podman load -q 2>/dev/null
    then
        echo "mkrel: save|load from $AMD64_CONNECTION failed; relaying via quay.io" >&2
        podman --connection "$AMD64_CONNECTION" tag "$local_tag" "$dest"
        podman --connection "$AMD64_CONNECTION" push --authfile "$authfile" "$dest"
        podman pull "$dest"
        verify_container_image "$dest" amd64
        return 0
    fi

    id=$(podman images --filter "reference=${local_tag}" --format '{{.ID}}' | head -1)
    if [ -z "$id" ]; then
        echo "mkrel: loaded amd64 image not found after save|load" >&2
        return 1
    fi
    podman tag "$id" "$dest"
    verify_container_image "$dest" amd64
}

build_arm64_image() {
    echo "mkrel: building linux/arm64 container image" >&2
    podman build \
        --platform linux/arm64 \
        "$@" \
        --tag "$(arch_tag arm64)" \
        "$ROOT"
    verify_container_image "$(arch_tag arm64)" arm64
}

# Pin ~/.ssh/farm.pub. Per-host details stay in ~/.ssh/config.
farm_sshwrap() {
    ssh \
        -i "$HOME/.ssh/farm.pub" \
        -o IdentitiesOnly=yes \
        -o ControlPath="$FARM_DIR/%h.ctl" \
        "$@"
}

# Forward the amd64 node's Podman socket to $FARM_DIR/amd64.sock.
farm_ssh() {
    local host=$1 remote
    farm_sshwrap -O exit "$host" >/dev/null 2>&1 ||
        rm -f "$FARM_DIR/$host.ctl" "$FARM_DIR/$host.sock"
    remote=$(farm_sshwrap "$host" \
        "podman info --format '{{.Host.RemoteSocket.Path}}'")
    farm_sshwrap -f -N \
        -L "$FARM_DIR/$host.sock:${remote#unix://}" \
        -o ControlMaster=yes \
        -o ExitOnForwardFailure=yes \
        -o StreamLocalBindUnlink=yes \
        "$host"
}

# Reuse a live socket owned by another process. Only a master this process
# created is closed on exit, so a nested local mkrel cannot drop it.
ensure_amd64_forward() {
    if podman --connection "$AMD64_CONNECTION" info >/dev/null 2>&1; then
        return 0
    fi
    mkdir -p "$FARM_DIR"
    farm_ssh "$FARM_HOST"
    trap 'farm_sshwrap -O exit "$FARM_HOST" >/dev/null 2>&1 || true' EXIT
}

build_amd64_image() {
    local local_tag
    local_tag=$(amd64_local_tag)

    ensure_amd64_forward

    echo "mkrel: building linux/amd64 container image on $AMD64_CONNECTION" >&2
    podman --connection "$AMD64_CONNECTION" build \
        --platform linux/amd64 \
        "$@" \
        --tag "$local_tag" \
        "$ROOT"
    fetch_amd64_image
}

# Fail unless ref is an image index containing linux/amd64 and linux/arm64.
# ref is a local list name or a registry name. manifest inspect rejects a
# docker:// transport, so strip one if a caller included it. A local list
# of the same name hides the registry tag; drop the local name first when
# the check is meant to see what was pushed.
require_multiarch_index() {
    local ref=$1 json err
    ref=${ref#docker://}
    err=$(mktemp)
    if ! json=$(podman manifest inspect "$ref" 2>"$err"); then
        echo "mkrel: $ref is not a multi-arch index" >&2
        # A single-image tag makes podman dump the whole manifest into the error.
        if ! grep -q 'Treating single images as manifest lists is not implemented' "$err"; then
            cat "$err" >&2
        fi
        rm -f "$err"
        return 1
    fi
    rm -f "$err"
    printf '%s\n' "$json" | python3 -c '
import json, sys
doc = json.load(sys.stdin)
plats = set()
for m in doc.get("manifests") or []:
    p = m.get("platform") or {}
    osname, arch = p.get("os"), p.get("architecture")
    if osname and arch:
        plats.add("%s/%s" % (osname, arch))
missing = [p for p in ("linux/amd64", "linux/arm64") if p not in plats]
if missing:
    have = ", ".join(sorted(plats)) or "none"
    sys.stderr.write("mkrel: %s missing %s (have %s)\n" % (
        sys.argv[1], ", ".join(missing), have))
    sys.exit(1)
' "$ref"
}

# Image IDs, not registry names — named refs make podman pull from the registry
# and can select the host architecture. --arch records the platform on the
# index; arm64 is variant v8, matching indexes Docker already selects correctly.
add_image_to_index() {
    local list=$1 id=$2 arch=$3
    if [ "$arch" = arm64 ]; then
        podman manifest add --os linux --arch arm64 --variant v8 "$list" "$id"
    else
        podman manifest add --os linux --arch "$arch" "$list" "$id"
    fi
}

# podman push of a list stores only the native architecture.
# After the push, drop a local tag of the destination so the follow-up
# inspect reads the registry copy instead of the local list.
push_container_index() {
    local list=$1 dest=$2
    require_multiarch_index "$list"
    podman manifest push --all "$list" "docker://${dest}"
    free_image_tag "$dest"
    require_multiarch_index "$dest"
}

assemble_container_manifest() {
    local amd_id arm_id list

    amd_id=$(podman image inspect --format '{{.Id}}' "$(arch_tag amd64)")
    arm_id=$(podman image inspect --format '{{.Id}}' "$(arch_tag arm64)")

    list="$image:latest"
    free_image_tag "$list"
    free_image_tag "$image:${TAG}"

    podman manifest create "$list"
    add_image_to_index "$list" "$amd_id" amd64
    add_image_to_index "$list" "$arm_id" arm64
    require_multiarch_index "$list"
}

container() {
    local style=${1:-release}
    local build_args=()

    podman rmi "$image:latest" "$image:${TAG}" 2>/dev/null || true
    podman rmi "$(amd64_local_tag)" 2>/dev/null || true

    if [ "$style" = "prerelease" ]; then
        build_args=(--build-arg "BRANCH=${BRANCH}" --build-arg "COMMIT=${TAG}" --build-arg "PRIOR=${PRIOR}" --build-arg "COMMIT_TIME=${COMMIT_TIME}")
    else
        build_args=(--build-arg "VERSION=${TAG}")
    fi

    build_arm64_image "${build_args[@]}"
    build_amd64_image "${build_args[@]}"
    assemble_container_manifest

    # One local list, two registry tags. Tagging the list and then using
    # podman push collapses it to the host architecture.
    push_container_index "$image:latest" "$image:${TAG}"
    push_container_index "$image:latest" "$image:latest"
}

platform_build() {
    local plat=$1
    local a=${plat#linux/}
    local tag other_a other_tag other_arch arch_id other_id

    case "$plat" in
    linux/amd64 | linux/arm64) ;;
    *)
        echo "mkrel platform: unsupported platform $plat (want linux/amd64 or linux/arm64)" >&2
        exit 1
        ;;
    esac

    [ -n "$TAG" ] || {
        echo "mkrel platform: TAG (or VERSION) must be set" >&2
        exit 1
    }

    tag=$(arch_tag "$a")
    if [ "$a" = arm64 ]; then
        build_arm64_image --build-arg "VERSION=${TAG}"
    else
        build_amd64_image --build-arg "VERSION=${TAG}"
    fi

    mkdir -p "$DIST"
    rm -f "$DIST/rejig-linux-$a.tar.xz" "$DIST/rejig-windows-$a.zip"
    extract_linux_artifacts "$a" "$tag"

    if [ "${PUSH:-1}" != "0" ]; then
        other_a=$([ "$a" = arm64 ] && echo amd64 || echo arm64)
        other_tag=$(arch_tag "$other_a")

        podman push "$tag"
        podman pull "$other_tag" 2>/dev/null || true
        if ! other_arch=$(podman image inspect --format '{{.Architecture}}' "$other_tag" 2>/dev/null); then
            echo "mkrel: $other_tag is not available; refusing to publish a one-arch index over ${image}:${TAG}" >&2
            exit 1
        fi
        if [ "$other_arch" != "$other_a" ]; then
            echo "mkrel: $other_tag is $other_arch, expected $other_a" >&2
            exit 1
        fi
        free_image_tag "$image:${TAG}"
        arch_id=$(podman image inspect --format '{{.Id}}' "$tag")
        other_id=$(podman image inspect --format '{{.Id}}' "$other_tag")
        podman manifest create "$image:${TAG}"
        add_image_to_index "$image:${TAG}" "$arch_id" "$a"
        add_image_to_index "$image:${TAG}" "$other_id" "$other_a"
        push_container_index "$image:${TAG}" "$image:${TAG}"
    fi

    if [ "${UV:-1}" != "0" ]; then
        ( cd "$DIST" &&
            for f in "rejig-linux-$a.tar.xz" "rejig-windows-$a.zip"; do
                [ -f "$f" ] && fossil uv add "$f"
            done &&
            fossil uv sync
        )
    fi
}

farm() {
    container release

    mkdir -p "$DIST"
    for a in amd64 arm64; do
        extract_linux_artifacts "$a"
    done

    for o in darwin linux windows; do
        for a in amd64 arm64; do
            plat="$o/$a"
            case "$o" in darwin) name="macos" ;; *) name="$o" ;; esac
            ext=$([ "$o" = windows ] && echo zip || echo tar.xz)
            art="$DIST/rejig-$name-$a.$ext"
            if [ ! -f "$art" ] || ! verify_dist_artifact "$art"; then
                rm -f "$art"
                (cd "$ROOT" && PLATFORM="$plat" TAG="$TAG" ./tools/mkrel)
            fi
        done
    done
    rm -f "$ROOT/rejig" "$ROOT/rejig.exe"

    verify_all_dist

    ( cd "$DIST" &&
        for f in rejig-*.tar.xz rejig-*.zip; do
            [ -f "$f" ] && fossil uv add "$f"
        done &&
        fossil uv sync
    )
}

if [ "$mode" = "container" ]; then
    container prerelease
elif [ "$mode" = "farm" ]; then
    farm
elif [ "$mode" = "platform" ]; then
    platform_build "$2"
else
    uv=$mode
    PLATFORM="${PLATFORM:-native}"
    if [ "$PLATFORM" = "native" ]; then
        export GOOS=$(go env GOOS)
        export GOARCH=$(go env GOARCH)
    else
        export GOOS="${PLATFORM%/*}"
        export GOARCH="${PLATFORM#*/}"
    fi
    OS=$GOOS
    test "$GOOS" = "darwin" && OS=macos

    bn="rejig-$OS-$GOARCH"
    if [ "$GOOS" = "windows" ]; then
        pkg="$bn.zip"
    else
        pkg="$bn.tar.xz"
    fi

    case "$GOOS" in
    darwin)  export CGO_LDFLAGS="-Wl,-dead_strip" ;;
    linux)   export CGO_LDFLAGS="-Wl,--gc-sections" ;;
    *)       export CGO_LDFLAGS="" ;;
    esac
    if [ -n "$TAG" ]; then
        go build -buildvcs=false -ldflags "-s -w -X rejig/mode.Version=$TAG" ./cmd/rejig
    else
        go build -buildvcs=false -ldflags "-s -w" ./cmd/rejig
    fi

    if [ "$GOOS" = "windows" ]; then
        zip -q "dist/$pkg" rejig.exe
        rm rejig.exe
    else
        tar cJf "dist/$pkg" rejig
    fi

    if [ -n "$TAG" ] && ! verify_dist_artifact "dist/$pkg"; then
        rm -f "dist/$pkg"
        exit 1
    fi

    if [ -n "$uv" ]; then
        fossil uv add "$pkg"
    fi
fi
