#!/bin/sh

set -e
set -x

OPENSSL=${OPENSSL:-openssl}

export TRIAGE_FINISHED="TRIAGE-FINISHED"
export TRIAGE_ABORTED="TRIAGE-ABORTED"
export TRIAGE_STARTED="TRIAGE-STARTED"

# Source common helper functions
# shellcheck disable=SC1091
. "$(dirname "$0")/rpi-sb-common.sh"
# shellcheck disable=SC1091
. /var/lib/rpi-sb-provisioner/state-recording

TARGET_DEVICE_SERIAL="${1}"
TARGET_DEVICE_SERIAL32=$(echo "${TARGET_DEVICE_SERIAL}" | cut -c $((${#TARGET_DEVICE_SERIAL}/2+1))-)
LOG_DIRECTORY="/var/log/rpi-sb-provisioner/${TARGET_DEVICE_SERIAL}"

log() {
    timestamp=$(date +"%Y-%m-%d %H:%M:%S.$(date +%N | cut -c1-3)")
    mkdir -p "${LOG_DIRECTORY}"
    echo "[${timestamp}] $*" >> "${LOG_DIRECTORY}"/triage.log
    printf "[%s] %s\n" "${timestamp}" "$*"
}

read_config
compute_image_summary

die() {
    # shellcheck disable=SC2086
    echo "$@" ${DEBUG}
    exit 1
}

CLEANUP_DONE=0

cleanup() {
    returnvalue=$?

    [ "$CLEANUP_DONE" -eq 1 ] && return
    CLEANUP_DONE=1

    set +e

    if [ "${returnvalue}" -ne 0 ]; then
        # Fastboot may appear briefly during bootstrap USB churn. A failed
        # early triage while bootstrap still holds its lock is not a user-
        # visible failure on programming rigs (keep the in-progress LED).
        if [ -z "${FASTBOOT_DEVICE_SPECIFIER:-}" ] && \
           [ -n "$(find /var/lock/rpi-sb-bootstrap -mindepth 1 -maxdepth 1 -type d -print -quit 2>/dev/null)" ]; then
            log "Skipping provision-failed: bootstrap in progress (expected USB re-enumeration)"
        else
            PROVISIONER_NAME=$(get_configured_provisioner_name)
            if [ -n "${FASTBOOT_DEVICE_SPECIFIER:-}" ]; then
                run_provision_failed_hook "${PROVISIONER_NAME}" "provisioning"
            else
                run_provision_failed_hook "${PROVISIONER_NAME}" "bootstrap"
            fi
        fi
    fi

    exit ${returnvalue}
}
trap cleanup EXIT INT TERM

timeout_nonfatal() {
    command="$*"
    set +e
    log "Running command with 10-second timeout: \"${command}\""
    # shellcheck disable=SC2086
    timeout 10 ${command}
    command_exit_status=$?
    
    # Handle different exit codes from the timeout command
    case ${command_exit_status} in
        0)
            # Command completed successfully within the time limit
            log "\"$command\" succeeded with exit code 0."
            ;;
        124)
            # Exit code 124 means the command timed out (TERM signal sent but command didn't exit)
            log "\"${command}\" FAILED: Timed out after 10 seconds (exit code 124)."
            ;;
        125)
            # Exit code 125 means the timeout command itself failed
            log "\"${command}\" FAILED: The timeout command itself failed (exit code 125)."
            ;;
        126)
            # Exit code 126 means the command was found but could not be executed
            log "\"${command}\" FAILED: Command found but could not be executed (exit code 126)."
            ;;
        127)
            # Exit code 127 means the command was not found
            log "\"${command}\" FAILED: Command not found (exit code 127)."
            ;;
        137)
            # Exit code 137 (128+9) means the command was killed by SIGKILL (kill -9)
            log "\"${command}\" FAILED: Command was killed by SIGKILL (exit code 137)."
            ;;
        *)
            # Any other non-zero exit code is a general failure
            log "\"${command}\" FAILED: Command returned exit code ${command_exit_status}."
            ;;
    esac
    set -e
    return ${command_exit_status}
}

timeout_fatal() {
    command="$*"
    set +e
    log "Running command with 30-second timeout: \"${command}\""
    # shellcheck disable=SC2086
    timeout 30 ${command}
    command_exit_status=$?
    
    # Handle different exit codes from the timeout command
    case ${command_exit_status} in
        0)
            # Command completed successfully within the time limit
            log "\"$command\" succeeded with exit code 0."
            ;;
        124)
            # Exit code 124 means the command timed out (TERM signal sent but command didn't exit)
            record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_ABORTED}" "${TARGET_USB_PATH}"
            die "\"${command}\" FAILED: Timed out after 30 seconds (exit code 124)."
            ;;
        125)
            # Exit code 125 means the timeout command itself failed
            record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_ABORTED}" "${TARGET_USB_PATH}"
            die "\"${command}\" FAILED: The timeout command itself failed (exit code 125)."
            ;;
        126)
            # Exit code 126 means the command was found but could not be executed
            record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_ABORTED}" "${TARGET_USB_PATH}"
            die "\"${command}\" FAILED: Command found but could not be executed (exit code 126)."
            ;;
        127)
            # Exit code 127 means the command was not found
            record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_ABORTED}" "${TARGET_USB_PATH}"
            die "\"${command}\" FAILED: Command not found (exit code 127)."
            ;;
        137)
            # Exit code 137 (128+9) means the command was killed by SIGKILL (kill -9)
            record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_ABORTED}" "${TARGET_USB_PATH}"
            die "\"${command}\" FAILED: Command was killed by SIGKILL (exit code 137)."
            ;;
        *)
            # Any other non-zero exit code is a general failure
            record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_ABORTED}" "${TARGET_USB_PATH}"
            die "\"${command}\" FAILED: Command returned exit code ${command_exit_status}."
            ;;
    esac
    set -e
}

# Load config early so the image summary is populated before any
# record_state call -- otherwise the device-details page shows an empty
# image for the lifetime of triage and only fills in later, which looks
# inconsistent.
# Best-effort USB path purely for the initial TRIAGE-STARTED record;
# setup_fastboot_and_id_vars recomputes it authoritatively below. Query the
# full serial (the gadget's USB serial is the full value, not the lower half)
# and never let a miss abort triage: under `set -e`, a failed command
# substitution in an assignment exits the script -- which previously killed
# triage before it could even record TRIAGE-STARTED or write any log.
TARGET_USB_PATH=$(get_usb_path_for_serial "${TARGET_DEVICE_SERIAL}") || TARGET_USB_PATH=""

# Record state changes atomically
record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_STARTED}" "${TARGET_USB_PATH}"

if [ -d "/var/log/rpi-sb-provisioner/${TARGET_DEVICE_SERIAL32}" ]; then
    cp -rf "/var/log/rpi-sb-provisioner/${TARGET_DEVICE_SERIAL32}/." "${LOG_DIRECTORY}/"
    rm -rf "/var/log/rpi-sb-provisioner/${TARGET_DEVICE_SERIAL32}"
fi

setup_fastboot_and_id_vars "${TARGET_DEVICE_SERIAL}"

# Ensure the device has a firmware crypto ECDSA key. Every provisioned
# device gets one regardless of style -- it is needed for HMAC-based LUKS
# key derivation and is useful for device identity.
#
# `oem fwcrypto init` is idempotent on the gadget side: OKAY whether the
# slot was empty (just generated) or already populated (no-op). Any
# non-zero exit here is a real firmware/transport failure, not "key
# already exists", so we can rely on the exit status alone.
log "Ensuring device firmware crypto key is provisioned"
timeout_fatal fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" oem fwcrypto init

KEYPAIR_DIR="${LOG_DIRECTORY}/${TARGET_DEVICE_SERIAL}"/keypair
if [ -d "${RPI_DEVICE_RETRIEVE_KEYPAIR}" ]; then
    KEYPAIR_DIR="${RPI_DEVICE_RETRIEVE_KEYPAIR}"
fi
mkdir -p "${KEYPAIR_DIR}"
log "Capturing device keypair to ${KEYPAIR_DIR}"
N_ALREADY_PROVISIONED=0
PRIVATE_KEY=$(get_variable private-key) || N_ALREADY_PROVISIONED=$?
if [ 0 -ne "$N_ALREADY_PROVISIONED" ]; then
    log "Warning: Unable to retrieve device private key; already provisioned"
else
    echo "${PRIVATE_KEY}" > "${KEYPAIR_DIR}/${TARGET_DEVICE_SERIAL}.der"
    PRIVATE_KEY=""
fi
get_variable public-key > "${KEYPAIR_DIR}/${TARGET_DEVICE_SERIAL}.pub"

# Based on the image type and provisioning style, we can determine which
# systemd unit to trigger.  All units are parameterised with the device serial.
echo "${TRIAGE_STARTED}" >> "${LOG_DIRECTORY}"/triage.log

# Select the target unit first, but DON'T start it yet.  `systemctl start` is
# synchronous: it blocks until the provisioner unit is active, by which point
# the provisioner has already written PROVISIONER-STARTED.  If we recorded
# TRIAGE-FINISHED after the start, it would land in the same (second-resolution)
# timestamp as PROVISIONER-STARTED but sort after it -- the state race seen in
# the UI.  Triage's job is done once it has selected the provisioner, so record
# TRIAGE-FINISHED first and hand off last.
if [ -d "${GOLD_MASTER_OS_FILE}" ]; then
    # GOLD_MASTER_OS_FILE is a directory -- this is an IDP artefact.
    # Route to the IDP provisioner regardless of PROVISIONING_STYLE,
    # since the IDP provisioner handles the image format natively.
    # (Secure boot is handled by the bootstrap phase, not the provisioner.)
    log "GOLD_MASTER_OS_FILE is a directory; selecting IDP Provisioner"
    TARGET_PROVISIONER_UNIT="rpi-idp-provisioner@${TARGET_DEVICE_SERIAL}.service"
else
    # Traditional .img file -- use the PROVISIONING_STYLE switch as before
    case ${PROVISIONING_STYLE} in
        "secure-boot")
            log "Selecting Secure Boot Provisioner"
            TARGET_PROVISIONER_UNIT="rpi-sb-provisioner@${TARGET_DEVICE_SERIAL}.service"
        ;;
        "fde-only")
            log "Selecting Full-Disk Encryption Provisioner"
            TARGET_PROVISIONER_UNIT="rpi-fde-provisioner@${TARGET_DEVICE_SERIAL}.service"
        ;;
        "naked")
            log "Selecting Naked Provisioner"
            TARGET_PROVISIONER_UNIT="rpi-naked-provisioner@${TARGET_DEVICE_SERIAL}.service"
        ;;
        *)
            log "Fatal: Unknown provisioning style: ${PROVISIONING_STYLE}"
            exit 1
        ;;
    esac
fi

# Mark triage finished before dispatching, so the state sequence reads
# TRIAGE-STARTED -> TRIAGE-FINISHED -> PROVISIONER-STARTED in causal order.
log "${TRIAGE_FINISHED}" >> "${LOG_DIRECTORY}"/triage.log
record_state "${TARGET_DEVICE_SERIAL}" "${TRIAGE_FINISHED}" "${TARGET_USB_PATH}"

# Hand off to the selected provisioner.
systemctl start "${TARGET_PROVISIONER_UNIT}"
