#!/usr/libexec/flua
local USAGE = [[
Patch or upgrade one pkgbase host by hand:
- create rollback boot environment,
- do pkg upgrades.
Equivalent to the patch_freebsd role, just with no safety guards.
Run on the host:
Usage: bump-freebsd-manually [-n] [-b] [-r] [-f] [-u X.Y]
  -n      report only: no boot environment, pkg upgrade -n
          (pkg itself is still upgraded, so the report is accurate)
  -b      base system only, leave other packages alone
  -r      reboot when done (default: leave the reboot to the operator)
  -f      official FreeBSD repos only (FreeBSD-base, FreeBSD-ports),
          ignoring rxv_devops
  -u X.Y  release upgrade to X.Y (e.g. 15.1), same major only; a
          -RELEASE or -pN suffix is ignored. Without it the host is
          patched on the release it already runs
The host OS is always upgraded from FreeBSD-base, and packages always
from rxv_devops, unless -f is given, in which case from FreeBSD-ports.
]]

local BASE_REPO = "FreeBSD-base"
local PORTS_REPO = "rxv_devops"
local OFFICIAL_PORTS_REPO = "FreeBSD-ports"
local BE_PREFIX = "patch"
local MIN_PKG = "2.8.4"
local EFI_LOADER = "/boot/loader.efi"
local EFI_BOOT_DIR = "/boot/efi/efi/boot"

local dry, base_only, reboot, official_only, target = false, false, false, false, nil

local function usage(code)
    io.stderr:write(USAGE)
    os.exit(code)
end

local function die(msg)
    io.stderr:write(msg, "\n")
    os.exit(1)
end

-- Single-quotes a value for safe interpolation into a shell command
-- string, since os.execute/io.popen hand the whole line to /bin/sh.
local function shquote(s)
    return "'" .. tostring(s):gsub("'", "'\\''") .. "'"
end

local function run(cmd)
    io.write("+ ", cmd, "\n")
    local ok = os.execute(cmd)
    if not ok then die("command failed: " .. cmd) end
end

-- Like run(), for "pkg upgrade -n": pkg exits non-zero both when there
-- is work to do and when something is wrong (repo unreachable or
-- disabled, broken dependency), so the exit status alone cannot tell
-- the two apart. Only a non-zero exit after pkg got as far as printing
-- its job summary counts as success.
local function run_dry(cmd)
    io.write("+ ", cmd, "\n")
    local p = io.popen(cmd .. " 2>&1", "r")
    if not p then die("failed to run: " .. cmd) end
    local summary = false
    for line in p:lines() do
        io.write(line, "\n")
        if line:find("^Number of packages to be ") then summary = true end
    end
    if not p:close() and not summary then die("command failed: " .. cmd) end
end

-- Every caller treats this like shell command substitution: $(cmd)
-- under "set -eu" does NOT abort on the inner command's exit status,
-- only the io.popen() call itself (the shell failing to start) does.
local function capture(cmd)
    local p = io.popen(cmd, "r")
    if not p then die("failed to run: " .. cmd) end
    local out = (p:read("*a") or ""):gsub("%s+$", "")
    p:close()
    return out
end

local function file_exists(path)
    return os.execute("test -f " .. shquote(path)) and true or false
end

-- Last bectl entry (by creation order) whose name starts with prefix
-- and has no active/mounted marker, i.e. "tail -1" over awk's matches.
local function find_existing_be(prefix)
    local match
    for line in capture("bectl list -H"):gmatch("[^\n]+") do
        local name, active = line:match("^([^\t]*)\t([^\t]*)")
        if active == "-" and name and name:sub(1, #prefix) == prefix then
            match = name
        end
    end
    return match
end

-- Hand-rolled getopt: -n -b -r are flags, -u takes an argument
-- (attached or as the next word), and short flags may be combined
-- (e.g. -nb), same as the shell version's getopts.
do
    local i = 1
    while i <= #arg do
        local a = arg[i]
        if a == "--" then
            i = i + 1
            break
        elseif a:match("^%-.") then
            local j = 2
            while j <= #a do
                local c = a:sub(j, j)
                if c == "n" then dry = true
                elseif c == "b" then base_only = true
                elseif c == "r" then reboot = true
                elseif c == "f" then official_only = true
                elseif c == "u" then
                    local val = a:sub(j + 1)
                    if val == "" then
                        i = i + 1
                        val = arg[i]
                        if not val then usage(2) end
                    end
                    target = val
                    j = #a
                else
                    usage(2)
                end
                j = j + 1
            end
        else
            break
        end
        i = i + 1
    end
end

if official_only then PORTS_REPO = OFFICIAL_PORTS_REPO end

if capture("id -u") ~= "0" then die("run as root (doas " .. arg[0] .. " ...)") end
if not capture("pkg which -q /usr/bin/uname"):find("FreeBSD%-runtime") then
    die("not a pkgbase host")
end

local installed
do
    local f = io.open("/etc/os-release", "r")
    if f then
        for line in f:lines() do
            local v = line:match('^VERSION_ID="([^"]*)"')
            if v then installed = v; break end
        end
        f:close()
    end
end
local running = capture("freebsd-version -r"):gsub("%-.*", "")
if not installed or installed == "" then
    die("cannot read VERSION_ID from /etc/os-release")
end

if target then
    -- 15.1-RELEASE and 15.1-p3 both mean 15.1, as in the patch_freebsd role
    local given = target
    target = target:gsub("%-.*", "")
    if not target:match("^%d+%.%d+$") then
        die("-u expects a release like 15.1, got " .. given)
    end
    local target_major = target:match("^(%d+)%.")
    local installed_major = installed:match("^(%d+)%.")
    local running_major = running:match("^(%d+)%.")
    if not (target_major == installed_major and target_major == running_major) then
        die("only minor upgrades within the same major are supported")
    end
    if capture("pkg version -t " .. shquote(target) .. " " .. shquote(installed)) == "<" then
        die("target " .. target .. " is older than installed " .. installed)
    end
    print("=== Upgrading " .. installed .. " (running " .. running .. ") -> " .. target .. " ===")
else
    -- a patch would reboot onto the new kernel but skip the EFI loader
    -- and package steps, leaving the upgrade half done
    if installed ~= running then
        die(installed .. " installed but " .. running .. " running: finish with -u " .. installed)
    end
    print("=== Patching " .. installed .. " ===")
end
run("freebsd-version -kru")

-- One rollback point per day is enough: reuse a BE made earlier today.
local be
if not dry then
    local today = os.date("!%Y%m%d")
    be = find_existing_be(BE_PREFIX .. "-" .. today .. "-")
    if be then
        print("rollback boot environment already exists: " .. be)
    else
        be = BE_PREFIX .. "-" .. os.date("!%Y%m%d-%H%M")
        run("bectl create " .. shquote(be))
    end
end

-- pkg first, on its own, so everything below uses the pkg that does the
-- work. This happens even with -n: the release upgrade needs MIN_PKG,
-- and a report from an older pkg would not show what a real run does.
run("pkg upgrade -y -r " .. shquote(PORTS_REPO) .. " pkg")

-- PKG is interpolated unquoted into the command below, word-split on
-- purpose: the release upgrade adds options to it. The -n/-y flag has
-- to precede any package pattern.
local PKG = "pkg"
local function upgrade(args)
    local cmd = PKG .. " upgrade " .. (dry and "-n" or "-y") .. " " .. args
    if dry then run_dry(cmd) else run(cmd) end
end

if not target then
    upgrade("-r " .. shquote(BASE_REPO) .. " -g 'FreeBSD-*'")
    if not base_only then upgrade("-r " .. shquote(PORTS_REPO)) end
else
    local pkg_static_v = capture("pkg-static -v")
    if capture("pkg version -t " .. shquote(pkg_static_v) .. " " .. shquote(MIN_PKG)) == "<" then
        die("pkg " .. pkg_static_v .. " is older than " .. MIN_PKG .. ", which a release upgrade requires")
    end
    -- the repo URLs expand ${VERSION_MINOR}, so overriding it points pkg
    -- at the target release's packages
    local minor = target:match("%.(%d+)$")
    -- The target release's packages are newer than the running system,
    -- which pkg refuses ("wrong OS version") unless told to ignore it.
    PKG = "pkg-static -o VERSION_MINOR=" .. minor .. " -o IGNORE_OSVERSION=yes"
    upgrade("-r " .. shquote(BASE_REPO))

    -- The firmware boots the ESP copy of the loader, which pkg does not
    -- touch. Keep the previous one as lastgood.efi.
    local arch = capture("uname -m")
    local efi_file
    if arch == "amd64" then efi_file = "bootx64.efi"
    elseif arch == "arm64" then efi_file = "bootaa64.efi"
    else die("unknown architecture " .. arch) end
    local efi_boot = EFI_BOOT_DIR .. "/" .. efi_file

    if not dry then
        if not (file_exists(EFI_LOADER) and file_exists(efi_boot)) then
            die("missing " .. EFI_LOADER .. " or " .. efi_boot ..
                " (is the ESP mounted on /boot/efi?). The base upgrade has been applied but the host has not been rebooted.")
        end
        if capture("sha256 -q " .. shquote(EFI_LOADER)) ~= capture("sha256 -q " .. shquote(efi_boot)) then
            run("mv " .. shquote(efi_boot) .. " " .. shquote(EFI_BOOT_DIR .. "/lastgood.efi"))
            run("cp -p " .. shquote(EFI_LOADER) .. " " .. shquote(efi_boot))
        end
    end

    if not base_only then upgrade("-r " .. shquote(PORTS_REPO)) end
end

if dry then os.exit(0) end

print("=== Done ===")
run("freebsd-version -kru")
print("rollback: bectl activate " .. be .. " && shutdown -r now")
if reboot then
    run('shutdown -r now "Rebooting after FreeBSD pkgbase upgrade"')
else
    print("reboot to finish: shutdown -r now")
end