#!/bin/env python3

# sysrapl.py
#
# Copyright 2023 Ondřej Míchal
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later

import argparse
import bisect
import contextlib
import ctypes
import datetime
import itertools
import json
import platform
import signal
import sys
import tempfile
import time
import traceback
from typing import Any, IO


# TODO: The values can differ across architectures and CPU generations. Read
# them from powercap sysfs.
RAPL_MAX_VALUE = 262143328850
RAPL_SCALE = (
    # FIXME: This is... just wrong. But it's the easiest as well.
    # The original value on my Skylake CPU is 2.3283064365386962890625e-10. The
    # adjusted value then is the way to scale Joules to nanojoules. To get
    # joules again, multiply by 10^9.
    0.23283064365386962890625
)
KNOWN_SYSCALLS: dict[int, str] = {}


def _diff_rapl_value(prev: int, next: int) -> int:
    if prev > next:
        padding = RAPL_MAX_VALUE - prev
        return next + padding
    else:
        return next - prev


def _sigint_empty_handler(signum, frame):
    pass


def _get_process_name(pid: int) -> str:
    with open(f"/proc/{pid}/comm", "r") as f:
        return f.read().strip()

    return "<Unknown>"


def _get_processor_model() -> str:
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                if not line.startswith("model name"):
                    continue

                return line.split(":", 1)[1].strip()
    except Exception:
        pass

    return "<Unknown>"


def _get_syscall_id_name(syscall_id: int) -> str:
    if syscall_id in KNOWN_SYSCALLS:
        return KNOWN_SYSCALLS[syscall_id]

    try:
        import system_calls

        syscalls = system_calls.syscalls()

        for name in syscalls.names():
            try:
                if syscall_id == syscalls.get(name):
                    KNOWN_SYSCALLS[syscall_id] = name
                    return name
            except system_calls.NotSupportedSystemCall:
                pass
            except system_calls.NoSuchSystemCall:
                pass
    except ImportError:
        pass

    return ""


class SysraplProfileOpts(ctypes.Structure):
    _fields_ = [
        ("output_file", ctypes.c_char_p),
        ("filter_cgroup", ctypes.c_bool),
        ("delay", ctypes.c_int),
        ("filter_pid", ctypes.c_int),
        ("frequency", ctypes.c_int),
    ]


class RaplData(dict):
    def __init__(
        self,
        time: int = 0,
        cores: int = 0,
        pkg: int = 0,
        ram: int = 0,
        gpu: int = 0,
        psys: int = 0,
    ):
        super().__init__(
            time=time,
            cores=cores,
            pkg=pkg,
            ram=ram,
            gpu=gpu,
            psys=psys,
        )

    def diff(prev, next):
        new = RaplData(0, 0, 0, 0, 0, 0)

        new["time"] = prev["time"]  # TODO: Why?
        new["cores"] = _diff_rapl_value(prev["cores"], next["cores"])
        new["pkg"] = _diff_rapl_value(prev["pkg"], next["pkg"])
        new["ram"] = _diff_rapl_value(prev["ram"], next["ram"])
        new["gpu"] = _diff_rapl_value(prev["gpu"], next["gpu"])
        new["psys"] = _diff_rapl_value(prev["psys"], next["psys"])

        return new

    def add(self, other):
        self["cores"] += other["cores"]
        self["pkg"] += other["pkg"]
        self["ram"] += other["ram"]
        self["gpu"] += other["gpu"]
        self["psys"] += other["psys"]

    def ratio(self, ratio):
        self["cores"] *= ratio
        self["pkg"] *= ratio
        self["ram"] *= ratio
        self["gpu"] *= ratio
        self["psys"] *= ratio


class SyscallData(dict):
    def __init__(
        self,
        id: int,
        name: str,
        cpu: int,
        tid: int,
        entry_time: int,
        exit_time: int,
    ):
        super().__init__(
            id=id,
            name=name,
            cpu=cpu,
            tid=tid,
            entry_time=entry_time,
            exit_time=exit_time,
        )


def get_rapl_window(
    rapl_events: list, entry_time: int, exit_time: int, helper_data: dict
):
    """
    Finds and returns a list of RAPL readings matching the entry and exit times
    of a system call.

    Parameters:
    ===========
    rapl_events - list of all RAPL readings
    entry_time - entry time (in nanoseconds) of a system call
    exit_time - exit time (in nanoseconds) of a system call
    helper_data - dictionary hodling context data
    """

    # There are no RAPL data that we can use to score the system call
    if entry_time < rapl_events[0]["time"] or exit_time < rapl_events[0]["time"]:
        return

    # Set up a sliding window to speed up the lookup of relevant rapl events.
    # The events are ordered in time, so we can do this.
    left_index = helper_data["left_index"] if "left_index" in helper_data else 0
    right_index = helper_data["right_index"] if "right_index" in helper_data else 0

    left_index = bisect.bisect_left(
        rapl_events, entry_time, key=lambda i: i["time"])
    right_index = bisect.bisect_right(
        rapl_events, exit_time, key=lambda i: i["time"])

    # At least one sample should be used
    right_index += 1

    helper_data["left_index"] = left_index
    helper_data["right_index"] = right_index

    return rapl_events[left_index:right_index]


def score_syscall(rapl_data_window: list, syscall_data: dict, rapl_sample_width: int):
    """
    Algorithm for giving an energy score to a system call.

    Parameters:
    ===========
    rapl_data_window - RAPL time window during which the system call ran
    syscall_data - system call data
    rapl_sample_width - width of a RAPL time window (nanoseconds)
    """
    entry_time = syscall_data["entry_time"]
    exit_time = syscall_data["exit_time"]

    syscall_duration = exit_time - entry_time
    ratio = syscall_duration / (rapl_sample_width * len(rapl_data_window))

    # Use the duration of the system call as the time information
    syscall_rapl_data = RaplData(exit_time - entry_time, 0, 0, 0, 0, 0)
    for rapl_data in rapl_data_window:
        # TODO: The edge entries need to be interpolated because the syscall
        # most likely did not start the moment the RAPL sampl was made. For now
        # that is not a concern.

        syscall_rapl_data.add(rapl_data)

    syscall_rapl_data.ratio(ratio)

    syscall_data["cores"] = syscall_rapl_data["cores"]
    syscall_data["pkg"] = syscall_rapl_data["pkg"]
    syscall_data["ram"] = syscall_rapl_data["ram"]
    syscall_data["gpu"] = syscall_rapl_data["gpu"]
    syscall_data["psys"] = syscall_rapl_data["psys"]


def process_syscall(profile_data: dict, values: list[str], helper_data: dict):
    """
    Procesesses raw system call trace by looking up its name based on its ID
    and gives it an energy score using a scoring algorithm.

    Parameters:
    ===========
    profile_data - dictionary wehre the processed info should be put
    values - raw system call trace
    helper_data - dictionary hodling context data
    """
    cpu: int = int(values[1])
    tid: int = int(values[2])
    syscall_id: int = int(values[3])
    entry_time_ns: int = int(values[4])
    exit_time_ns: int = int(values[5])

    # The system_calls library is not sure to know a syscall's name, so
    # syscall_name can be an empty string.
    syscall_name = _get_syscall_id_name(syscall_id)

    rapl_data_window = get_rapl_window(
        profile_data["rapl"]["events"], entry_time_ns, exit_time_ns, helper_data
    )

    # FIXME: get_rapl_window() should never return None
    # But this is a much quicker fix for the time being
    if not rapl_data_window:
        return

    syscall_data = SyscallData(
        syscall_id,
        syscall_name,
        cpu,
        tid,
        entry_time_ns,
        exit_time_ns,
    )

    score_syscall(rapl_data_window, syscall_data,
                  helper_data["rapl_sample_width"])

    # Consumption profile in time (chronological order of syscalls with
    # context and energy price
    profile_data["syscalls"].append(syscall_data)


def process_rapl(profile_rapl: dict, values: list[str], helper_data: dict):
    """
    Processes raw rapl energy readings by scaling the values and calculating
    energy consumed in between the different samples.

    Parameters:
    ===========
    profile_rapl - dictionary where theprocessed info should be put
    values - raw rapl energy readings
    helper_data - dictionary holding context data
    """
    time: int = int(values[1])
    cores: int = int(int(values[2]) * RAPL_SCALE)
    pkg: int = int(int(values[3]) * RAPL_SCALE)
    ram: int = int(int(values[4]) * RAPL_SCALE)
    gpu: int = int(int(values[5]) * RAPL_SCALE)
    psys: int = int(int(values[6]) * RAPL_SCALE)

    # Total consumption
    if "prev_rapl_event" not in helper_data:
        helper_data["prev_rapl_event"] = RaplData(
            time, cores, pkg, ram, gpu, psys)
        return

    prev_rapl_event: RaplData = helper_data["prev_rapl_event"]
    next_rapl_event: RaplData = RaplData(time, cores, pkg, ram, gpu, psys)

    diff_rapl_data: RaplData = RaplData.diff(prev_rapl_event, next_rapl_event)

    profile_rapl["events"].append(diff_rapl_data)
    profile_rapl["total"].add(diff_rapl_data)

    helper_data["prev_rapl_event"] = next_rapl_event


def post_process(raw_profile_file: IO, profile_data: dict):
    """
    Processes raw profiling data from an energy profiling session.

    Parameters:
    ===========
    raw_profile_file - file containing the raw profiling data
    profile_data - dictionary where the processed info should be put
    """
    helper_data: dict = {}
    profile_rapl: dict = profile_data["rapl"]

    helper_data["rapl_sample_width"] = 1_000_000_000 / \
        profile_data["info"]["frequency"]

    print("Processing profiled data", file=sys.stderr)

    lines = 0
    for line in raw_profile_file:
        line = line.strip()
        values: list[str] = line.split(" ")
        record_type: int = int(values[0])

        lines += 1
        if record_type != 1:
            continue

        process_rapl(profile_rapl, values, helper_data)

    if len(profile_rapl["events"]) > 0:
        profile_rapl["total"]["time"] = (
            profile_rapl["events"][-1]["time"] -
            profile_rapl["events"][0]["time"]
        )

    raw_profile_file.seek(0)

    for line in itertools.islice(raw_profile_file, lines):
        line = line.strip()
        values = line.split(" ")
        record_type = int(values[0])

        if record_type != 0:
            continue

        process_syscall(profile_data, values, helper_data)


def profile(args: argparse.Namespace, file: tempfile.NamedTemporaryFile):
    """
    Starts an energy profiling session using eBPF. The result of this session
    are raw profiling data that need to be further processed.

    Parameters:
    ===========
    args - CLI arguments
    file - output file for the raw profiling data
    """
    sysrapl: ctypes.CDLL = ctypes.CDLL("/usr/local/lib64/libsysrapl.so")

    profile_opts: SysraplProfileOpts = SysraplProfileOpts(
        ctypes.c_char_p(str.encode(file.name)),
        args.delay,
        args.pid,
        args.frequency,
    )

    # The C library implements its own SIGINT handler. Suppress the default one
    # for the duration of the collection, so that the post-processing phase is
    # not skipped.
    orig_sigint_handler = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, _sigint_empty_handler)

    print("Profiling. Press Ctrl+C to stop.", file=sys.stderr, end=None)

    if sysrapl.sysrapl_profile(ctypes.byref(profile_opts)) != 0:
        print("FAILED")
        raise Exception()

    signal.signal(signal.SIGINT, orig_sigint_handler)

    return 0


def main() -> int:
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        prog="sysrapl", description="Profiler of system call energy consumption"
    )

    parser.add_argument(
        "-d",
        "--delay",
        default=0,
        action="store",
        help="Delay before starting the profiling. [0]",
    )
    parser.add_argument(
        "-f",
        "--frequency",
        default=99,
        action="store",
        type=int,
        help="Frequency at which RAPL consumption will be read. [99]",
    )
    parser.add_argument(
        "-o",
        "--output",
        default=None,
        action="store",
        help="File where to store the profile. If not provided the profile is printed to stdout",
    )
    parser.add_argument(
        "-p",
        "--pid",
        action="store",
        type=int,
        help="PID of the profiled process",
    )
    parser.add_argument(
        "--reevaluate",
        default=None,
        type=argparse.FileType("r"),
        help="Path to profile to re-evaluate its content.",
    )
    parser.add_argument(
        "--version", default=False, action="store_true", help="Print program version"
    )

    args: argparse.Namespace = parser.parse_args()

    if args.reevaluate:
        if not args.output:
            print(
                "Option --output needs to be set when --reevaluate is used.",
                file=sys.stderr,
            )

            return 1

        with args.reevaluate as f:
            profile_data: dict = json.load(f)

            rapl_events = profile_data["rapl"]["events"]
            syscall_data = profile_data["syscalls"]
            helper_data: dict[str, Any] = {}
            helper_data["rapl_sample_width"] = (
                1_000_000_000 / profile_data["info"]["frequency"]
            )
            for syscall in syscall_data:
                rapl_data_window = get_rapl_window(
                    rapl_events,
                    syscall["entry_time"],
                    syscall["exit_time"],
                    helper_data,
                )

                score_syscall(
                    rapl_data_window, syscall, helper_data["rapl_sample_width"]
                )

            with open(args.output, "w") as f:
                print(json.dumps(profile_data), file=f)

            return 0

    if not args.pid or args.pid < 1:
        print("The profiled PID has to be greater or equal to 1", file=sys.stderr)
        return 1

    # This is a very unreadable way of passing two closable files using with
    # of which one is either sys.stdout or an actual file.
    with tempfile.NamedTemporaryFile(
        mode="w+t", prefix="sysrapl-", suffix="-raw-profile"
    ) as raw_profile_file, open(
        args.output, mode="w+t"
    ) if args.output is not None else contextlib.nullcontext(
        sys.stdout
    ) as profile_file:
        profile_start = time.monotonic_ns()

        # Profile a process using the profiling engine.
        try:
            profile(args, raw_profile_file)
        except Exception:
            print(f"Failed to profile process {args.pid}", file=sys.stderr)
            print(traceback.format_exc(), file=sys.stderr)
            return 1

        profile_end = time.monotonic_ns()

        # Initialize the final profile JSON
        profile_data: dict = {
            "info": {
                "process-name": _get_process_name(args.pid),
                "pid": args.pid,
                "time": datetime.datetime.now().isoformat(timespec="seconds"),
                "duration": profile_end - profile_start,
                "frequency": args.frequency,
                "host-info": platform.uname()._asdict(),
            },
            "rapl": {
                "unit": "Joules",
                "scale": RAPL_SCALE,
                # TODO: Should be made into a parameter and the profiler should
                # filter out unsupported ones.
                "domains": ["cores", "pkg", "ram", "gpu", "psys"],
                "total": RaplData(),
                "events": [],
            },
            "syscalls": [],
        }
        profile_data["info"]["host-info"]["processor"] = _get_processor_model()

        # Process the raw profile data and populate the final profile JSON
        try:
            post_process(raw_profile_file, profile_data)
        except Exception:
            print("Failed to process profile data", file=sys.stderr)
            print(traceback.format_exc(), file=sys.stderr)
            return 1

        print(json.dumps(profile_data), file=profile_file)

    return 0


if __name__ == "__main__":
    main()
