#!/usr/bin/env python3

# Copyright 2026 DeepMI Lab, German Center for Neurodegenerative Diseases (DZNE), Bonn
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Run a command and report the resources it used.

This replaces a bash frontend to ``/usr/bin/time -f``. ``-f`` is a GNU extension, and macOS ships
BSD time, so that version could not work there at all: ``functions.sh`` probes it, finds it broken
and silently drops per-command timings for every macOS user. Here the numbers come from
``os.wait4()``, i.e. from the kernel's own rusage for the child, so one implementation covers both
platforms and there is no external dependency left to be missing.

The output format is a contract. ``recon_surf/utils/extract_recon_surf_time_info.py`` parses it
positionally: field 1 is the timestamp, field 2 the command name, field 5 must be the literal ``e``
and field 6 the elapsed seconds. Keep those positions if you change anything here.
"""

import os
import signal
import sys
import time
from datetime import datetime

VERSION = "fs_time, part of FastSurfer, see 'run_fastsurfer.sh --version'"

# Fields the kernel does not populate on this platform, reported as "." rather than a misleading 0.
# macOS never counts block IO or swaps for a child; Linux does count block IO. ru_nswap is
# unmaintained on both, but on Linux we keep printing its 0 so the output stays identical to what
# /usr/bin/time produced there for years.
UNAVAILABLE = ("I", "O", "W") if sys.platform == "darwin" else ()

USAGE = """fs_time [options] command args
 options:
  -o outputfile : save resource info into outputfile instead
  -a : append to outputfile (see time)
  -k key : key to find timing info in output
  -l|--load : report on load averages as from uptime (default: $FSTIME_LOAD or on)
  --no-load : do not report on load averages
  --debug : make output more verbose
"""

HELP = """
This keeps track of the resources used by a process. The basic usage is like that of time, ie,

fs_time [options] command args

It creates a formatted output that allows for easy processing. See below.

If the env variable FSTIME_LOAD is set to 1 or not set at all, then the load averages of the system
are reported before and after each process (see below for output).

Default fs_time Output:
1. Key (default is @#@FSTIME)
2. Time stamp at the onset of execution
3. Command name
4. N Number of arguments passed to command
5. e Elapsed real time in seconds. This is the total amount of time from when the
     command started to when it ended regardless of how much of the CPU it got.
6. S Total number of CPU-seconds that the process spent in kernel mode.
7. U Total number of CPU-seconds that the process spent in user mode.
8. P Percentage of the CPU that this job got, computed as (U+S)/e.
9. M Maximum resident set size of the process during its lifetime, in Kbytes.
10. F Number of major page faults that occurred while the process was running.
      These are faults where the page has to be read in from disk.
11. R Number of minor, or recoverable, page faults.  These are faults for pages that
      are not valid but which have not yet been claimed by other virtual pages.  Thus
      the data in the page is still valid but the system tables must be updated.
12. W Number of times the process was swapped out of main memory.
13. c Number of times the process was context-switched involuntarily (because the
      time slice expired).
14. w Number of waits: times that the program was context-switched voluntarily, for
      instance while  waiting  for an I/O operation to complete.
15. I Number of file system inputs by the process.
16. O Number of file system outputs by the process.
17. L L1 L5 L15 : load averages at 1, 5, and 15 min (with setenv FSTIME_LOAD 1)

A field is reported as "." where the operating system does not measure it: macOS counts neither
file system IO (I, O) nor swaps (W) for a child process.

Example:

fs_time mri_convert orig.mgz myfile.mgz
mri_convert orig.mgz myfile.mgz
reading from orig.mgz...
writing to myfile.mgz...
@#@FSTIME 2016:01:21:18:27:08 mri_convert N 2 e 2.20 S 0.05 U 1.64 P 77% M 23628 F 0 R 5504 W 0 c 7 w 3 I 0 O 20408

The above command runs the mri_convert command with two arguments and produces the information
about resources. In this case, the above is interpreted as:

@#@FSTIME : key for easy searching/grepping
mri_convert : command that was run
2016:01:21:18:27:08 : time stamp at the onset of execution year:month:day:hour:min:sec
N 2 : mri_convert was run with 2 arguments
e 2.20 : total elapsed time in seconds from start to end
S 0.05 : seconds spent in system mode
U 1.64 : seconds spent in user mode
P 77%  : percent of cpu that process used (S+U)/e
M 23628 : maximum memory size in Kbytes
F 0 : no major page faults
R 5504 : number of minor page faults
W 0 : process was not swapped out of memory
c 7 : seven involuntary context-switches
w 3 : three voluntary context-switches
I 0 : zero file system inputs by the process.
O 20408 : Number of file system outputs by the process.

If the env variable FSTIME_LOAD is set to 1, the output looks something like

@#@FSLOADPRE 2016:01:23:15:11:04 mri_convert N 2 L 0.00 0.03 0.06
@#@FSLOADPOST 2016:01:23:15:11:06 mri_convert N 2 L 0.01 0.03 0.06

One sample before the command and one after it. The 3 numbers following the L are the system load
averages for the past 1, 5, and 15 minutes.
"""


def arg1err(flag: str) -> None:
    print(f"ERROR: Flag {flag} requires one argument!")
    sys.exit(1)


def shorten_command(cmd: list[str]) -> tuple[str, int]:
    """
    Return the name to report for `cmd`, and how many leading entries it consumed.

    A python invocation is reported as the script it runs, not as "python3", and a path inside
    $FASTSURFER_HOME is reported relative to it, so that the command name stays a short, stable key.

    An `env VAR=value ...` prefix is skipped first. Callers use it to set a variable for one command,
    and reporting every such call as "env" would collapse them into one key in the timing file.
    """
    skipped = 0
    if cmd and os.path.basename(cmd[0]) == "env":
        skipped = 1
        # the assignments belong to env; the first entry that is not one starts the real command
        for entry in cmd[1:]:
            if "=" in entry and not entry.startswith("-"):
                skipped += 1
            else:
                break
        cmd = cmd[skipped:]
        if not cmd:
            return "env", 0
    command = cmd[0]
    npyargs = skipped
    basename = os.path.basename(command)
    if basename in ("python", "python3") or (
        basename.startswith("python3.") and basename[8:].isdigit()
    ):
        npyargs = skipped + 1
        for entry in cmd[1:]:
            if not entry.startswith("-"):
                command = entry
                break
            npyargs += 1
    fastsurfer_home = os.environ.get("FASTSURFER_HOME", "")
    if fastsurfer_home and command.startswith(fastsurfer_home):
        command = command[len(fastsurfer_home):].lstrip("/")
    return command, npyargs


def format_load() -> str:
    """The 1, 5 and 15 minute load averages, in the same shape `uptime` reported them."""
    return "L " + " ".join(f"{value:.2f}" for value in os.getloadavg())


def make_line(key: str, command: str, nargs: int, extra: str = "", when: datetime | None = None) -> str:
    """One output line: `<key> <timestamp> <command> N <nargs>[ <extra>]`.

    `when` defaults to now, but the main line has to pass the time the command *started*: the help
    calls field 2 the timestamp at the onset of execution, and extract_recon_surf_time_info.py adds
    the elapsed seconds to it to get the end time.
    """
    if extra and not extra.startswith(" "):
        extra = " " + extra
    stamp = (when or datetime.now()).strftime("%Y:%m:%d:%H:%M:%S")
    return f"{key} {stamp} {command} N {nargs}{extra}"


def field(name: str, value: int | str) -> str:
    """`name value`, or `name .` where this platform does not measure it."""
    return f"{name} {'.' if name in UNAVAILABLE else value}"


def main(argv: list[str]) -> int:
    outfile = ""
    key = "@#@FSTIME "
    append = False
    cmd: list[str] = []
    verbose = False
    # on unless FSTIME_LOAD says otherwise. Unset and empty both mean on, as in the bash version
    # (which assigned the variable and then tested -z) and as the help describes.
    fstime_load = os.environ.get("FSTIME_LOAD", "")
    print_load = not fstime_load or fstime_load == "1"

    index = 0
    while index < len(argv):
        flag = argv[index]
        index += 1
        if flag in ("-o", "-k") and index >= len(argv):
            arg1err(flag)
        if flag in ("-help", "--help"):
            print(USAGE + HELP)
            return 0
        elif flag in ("-version", "--version"):
            print(VERSION)
            return 0
        elif flag == "-o":
            outfile = argv[index]
            index += 1
        elif flag == "-k":
            key = argv[index]
            index += 1
        elif flag == "-a":
            append = True
        elif flag in ("-l", "-load", "--load"):
            print_load = True
        elif flag in ("-no-load", "--no-load"):
            print_load = False
        elif flag in ("-debug", "--debug"):
            verbose = True
        else:
            # the command to run, and everything after it belongs to the command
            cmd = argv[index - 1:]
            break

    if not cmd:
        print(USAGE)
        print("ERROR: No command passed to execute!")
        return 1

    command, npyargs = shorten_command(cmd)
    nargs = len(cmd) - 1 - npyargs

    if verbose:
        print("Parameters to fs_time:")
        if outfile:
            print(f"  -o {outfile}")
        if key != "@#@FSTIME ":
            print(f"  -k {key}")
        if append:
            print("  -a")
        print(f"  {'--load' if print_load else '--no-load'}")
        print("command:")
        print(f"  {' '.join(cmd)}")
        print("")

    # open -o before running anything, as /usr/bin/time does: it exits 125 without touching the
    # command when it cannot write its output. Opening afterwards would spend a whole command's
    # runtime and then throw the result away.
    handle = None
    if outfile:
        try:
            handle = open(outfile, "a" if append else "w")
        except OSError as e:
            print(f"fs_time: {outfile}: {e.strerror}", file=sys.stderr)
            return 125

    if print_load:
        print(make_line("@#@FSLOADPRE", command, nargs, format_load()))
        sys.stdout.flush()

    # posix_spawn rather than fork+exec: it reports a missing command in the parent, so there is no
    # child to fail inside, and it leaves the child's minor-fault count reflecting the command
    # rather than the cost of copying our own address space.
    # setsigdef because python starts life with SIGPIPE and SIGXFSZ ignored and a spawned child
    # inherits that. The command would then get EPIPE errors where every other caller gets killed
    # by SIGPIPE, so `cmd | head` inside a timed script would report a write failure instead of
    # finishing quietly. This is what subprocess's restore_signals does, and it restores the
    # behaviour the /usr/bin/time version had, since bash execs children with default dispositions.

    # Do not abandon the command on Ctrl-C. A terminal sends SIGINT to the whole foreground process
    # group, so the command has it too and is already exiting; python's default handler would raise
    # KeyboardInterrupt out of the wait below, losing the record for the command that was
    # interrupted and printing a traceback into the log instead. A handler that does nothing lets
    # the wait restart until the command is reaped, so it is still reported and its 128+signal
    # status is still forwarded. Unlike SIG_IGN a handler is reset on exec, so the command itself
    # still dies on Ctrl-C.
    signal.signal(signal.SIGINT, lambda _signum, _frame: None)

    onset = datetime.now()
    start = time.monotonic()
    try:
        pid = os.posix_spawnp(cmd[0], cmd, os.environ, setsigdef=[signal.SIGPIPE, signal.SIGXFSZ])
    except FileNotFoundError:
        print(f"fs_time: {cmd[0]}: command not found", file=sys.stderr)
        return 127
    except PermissionError:
        print(f"fs_time: {cmd[0]}: permission denied", file=sys.stderr)
        return 126
    except OSError as e:
        print(f"fs_time: {cmd[0]}: {e}", file=sys.stderr)
        return 126

    _, status, usage = os.wait4(pid, 0)
    elapsed = time.monotonic() - start

    cpu = usage.ru_utime + usage.ru_stime
    percent = round(cpu / elapsed * 100) if elapsed > 0 else 0
    # ru_maxrss is bytes on macOS but kilobytes on Linux; report kilobytes on both, so the number
    # means the same thing wherever the pipeline ran.
    # Note Linux does not reset the peak on exec, so this cannot read below our own footprint of a
    # few MB: measured against /usr/bin/time it is identical for a 500 MB command and too high for a
    # trivial one. That only affects commands too small for the figure to be interesting.
    maxrss = usage.ru_maxrss // 1024 if sys.platform == "darwin" else usage.ru_maxrss

    measurements = " ".join((
        f"e {elapsed:.2f}",
        f"S {usage.ru_stime:.2f}",
        f"U {usage.ru_utime:.2f}",
        f"P {percent}%",
        f"M {maxrss}",
        field("F", usage.ru_majflt),
        field("R", usage.ru_minflt),
        field("W", usage.ru_nswap),
        field("c", usage.ru_nivcsw),
        field("w", usage.ru_nvcsw),
        field("I", usage.ru_inblock),
        field("O", usage.ru_oublock),
    ))
    if print_load:
        measurements += " " + format_load()

    line = make_line(key, command, nargs, measurements, when=onset)
    if handle is not None:
        with handle:
            handle.write(line + "\n")
    else:
        # stderr, as /usr/bin/time does: run_it in functions.sh relies on `2>&1 | tee` to capture it
        print(line, file=sys.stderr)

    if print_load:
        # with the load: the whole point of a post-run sample, and what the help documents. The
        # bash version passed no load here, so the line carried nothing but a timestamp; the
        # original tcsh fs_time did include it.
        print(make_line("@#@FSLOADPOST", command, nargs, format_load()))

    if os.WIFSIGNALED(status):
        return 128 + os.WTERMSIG(status)
    return os.WEXITSTATUS(status)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
