/*
 * Copyright 2009, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include <File.h>

#include <syscalls.h>
#include <system_profiler_defs.h>

#include <DebugEventStream.h>

#include "debug_utils.h"


#define SCHEDULING_RECORDING_AREA_SIZE  (4 * 1024 * 1024)

#define DEBUG_EVENT_MASK \
        (B_SYSTEM_PROFILER_TEAM_EVENTS | B_SYSTEM_PROFILER_THREAD_EVENTS        \
                | B_SYSTEM_PROFILER_SCHEDULING_EVENTS                                                   \
                | B_SYSTEM_PROFILER_IO_SCHEDULING_EVENTS)


extern const char* __progname;
const char* kCommandName = __progname;


static const char* kUsage =
        "Usage: %s [ <options> ] <output file> [ <command line> ]\n"
        "Records thread scheduling information to a file for later analysis.\n"
        "If a command line <command line> is given, recording starts right before\n"
        "executing the command and stops when the respective team quits.\n"
        "\n"
        "Options:\n"
        "  -l           - When a command line is given: Start recording before\n"
        "                 executable has been loaded.\n"
        "  -r           - Don't profile, but evaluate recorded kernel profile data.\n"
        "  -h, --help   - Print this usage info.\n"
;


static void
print_usage_and_exit(bool error)
{
    fprintf(error ? stderr : stdout, kUsage, kCommandName);
    exit(error ? 1 : 0);
}


class Recorder {
public:
        Recorder()
                :
                fMainTeam(-1),
                fSkipLoading(true),
                fCaughtDeadlySignal(false)
        {
        }

        ~Recorder()
        {
                fOutput.Flush();
        }


        status_t Init(const char* outputFile)
        {
                // open file
                status_t error = fOutputFile.SetTo(outputFile,
                        B_READ_WRITE | B_CREATE_FILE | B_ERASE_FILE);
                if (error != B_OK) {
                        fprintf(stderr, "Error: Failed to open \"%s\": %s\n", outputFile,
                                strerror(error));
                        return error;
                }

                // create output stream
                error = fOutput.SetTo(&fOutputFile, 0, DEBUG_EVENT_MASK);
                if (error != B_OK) {
                        fprintf(stderr, "Error: Failed to initialize the output "
                                "stream: %s\n", strerror(error));
                        return error;
                }

                return B_OK;
        }

        void SetSkipLoading(bool skipLoading)
        {
                fSkipLoading = skipLoading;
        }

        void Run(const char* const* programArgs, int programArgCount)
        {
                // Load the executable, if we have to.
                thread_id threadID = -1;
                if (programArgCount >= 1) {
                        threadID = load_program(programArgs, programArgCount,
                                !fSkipLoading);
                        if (threadID < 0) {
                                fprintf(stderr, "%s: Failed to start `%s': %s\n", kCommandName,
                                        programArgs[0], strerror(threadID));
                                exit(1);
                        }
                        fMainTeam = threadID;
                }

                // install signal handlers so we can exit gracefully
                struct sigaction action;
                action.sa_handler = (__sighandler_t)_SignalHandler;
                action.sa_flags = 0;
                sigemptyset(&action.sa_mask);
                action.sa_userdata = this;
                if (sigaction(SIGHUP, &action, NULL) < 0
                        || sigaction(SIGINT, &action, NULL) < 0
                        || sigaction(SIGQUIT, &action, NULL) < 0) {
                        fprintf(stderr, "%s: Failed to install signal handlers: %s\n",
                                kCommandName, strerror(errno));
                        exit(1);
                }

                // create an area for the sample buffer
                system_profiler_buffer_header* bufferHeader;
                area_id area = create_area("profiling buffer", (void**)&bufferHeader,
                        B_ANY_ADDRESS, SCHEDULING_RECORDING_AREA_SIZE, B_NO_LOCK,
                        B_READ_AREA | B_WRITE_AREA);
                if (area < 0) {
                        fprintf(stderr, "%s: Failed to create sample area: %s\n",
                                kCommandName, strerror(area));
                        exit(1);
                }

                uint8* bufferBase = (uint8*)(bufferHeader + 1);
                size_t totalBufferSize = SCHEDULING_RECORDING_AREA_SIZE
                        - (bufferBase - (uint8*)bufferHeader);

                // start profiling
                system_profiler_parameters profilerParameters;
                profilerParameters.buffer_area = area;
                profilerParameters.flags = DEBUG_EVENT_MASK;
                profilerParameters.locking_lookup_size = 64 * 1024;

                status_t error = _kern_system_profiler_start(&profilerParameters);
                if (error != B_OK) {
                        fprintf(stderr, "%s: Failed to start profiling: %s\n", kCommandName,
                                strerror(error));
                        exit(1);
                }

                // resume the loaded team, if we have one
                if (threadID >= 0)
                        resume_thread(threadID);

                // main event loop
                while (true) {
                        // get the current buffer
                        size_t bufferStart = bufferHeader->start;
                        size_t bufferSize = bufferHeader->size;
                        uint8* buffer = bufferBase + bufferStart;
//printf("processing buffer of size %lu bytes\n", bufferSize);

                        bool quit;
                        if (bufferStart + bufferSize <= totalBufferSize) {
                                quit = _ProcessEventBuffer(buffer, bufferSize);
                        } else {
                                size_t remainingSize = bufferStart + bufferSize
                                        - totalBufferSize;
                                quit = _ProcessEventBuffer(buffer, bufferSize - remainingSize)
                                        || _ProcessEventBuffer(bufferBase, remainingSize);
                        }

                        if (quit || fCaughtDeadlySignal)
                                break;

                        // get next buffer
                        uint64 droppedEvents = 0;
                        error = _kern_system_profiler_next_buffer(bufferSize,
                                &droppedEvents);

                        if (error != B_OK) {
                                if (error == B_INTERRUPTED)
                                        continue;

                                fprintf(stderr, "%s: Failed to get next sample buffer: %s\n",
                                        kCommandName, strerror(error));
                                break;
                        }

                        if (droppedEvents > 0)
                                fprintf(stderr, "%llu events dropped\n", droppedEvents);
                }

                // stop profiling
                _kern_system_profiler_stop();
        }

        void DumpRecorded()
        {
                // retrieve recorded samples and parameters
                system_profiler_parameters profilerParameters;
                status_t error = _kern_system_profiler_recorded(&profilerParameters);
                if (error != B_OK) {
                        fprintf(stderr, "%s: Failed to get recorded profiling buffer: %s\n",
                                kCommandName, strerror(error));
                        exit(1);
                }

                // create an area for the sample buffer
                area_info info;
                error = get_area_info(profilerParameters.buffer_area, &info);
                if (error != B_OK) {
                        fprintf(stderr, "%s: Recorded profiling buffer invalid: %s\n",
                                kCommandName, strerror(error));
                        exit(1);
                }

                system_profiler_buffer_header* bufferHeader
                        = (system_profiler_buffer_header*)info.address;

                uint8* bufferBase = (uint8*)(bufferHeader + 1);
                size_t totalBufferSize = info.size - (bufferBase - (uint8*)bufferHeader);

                // get the current buffer
                size_t bufferStart = bufferHeader->start;
                size_t bufferSize = bufferHeader->size;
                uint8* buffer = bufferBase + bufferStart;

                if (bufferStart + bufferSize <= totalBufferSize) {
                        _ProcessEventBuffer(buffer, bufferSize);
                } else {
                        size_t remainingSize = bufferStart + bufferSize - totalBufferSize;
                        if (!_ProcessEventBuffer(buffer,
                                        bufferSize - remainingSize)) {
                                _ProcessEventBuffer(bufferBase, remainingSize);
                        }
                }
        }

private:
        bool _ProcessEventBuffer(uint8* buffer, size_t bufferSize)
        {
//printf("_ProcessEventBuffer(%p, %lu)\n", buffer, bufferSize);
                const uint8* bufferStart = buffer;
                const uint8* bufferEnd = buffer + bufferSize;
                size_t usableBufferSize = bufferSize;
                bool quit = false;

                while (buffer < bufferEnd) {
                        system_profiler_event_header* header
                                = (system_profiler_event_header*)buffer;

                        buffer += sizeof(system_profiler_event_header);

                        if (header->event == B_SYSTEM_PROFILER_BUFFER_END) {
                                // Marks the end of the ring buffer -- we need to ignore the
                                // remaining bytes.
                                usableBufferSize = (uint8*)header - bufferStart;
                                break;
                        }

                        if (header->event == B_SYSTEM_PROFILER_TEAM_REMOVED) {
                                system_profiler_team_removed* event
                                        = (system_profiler_team_removed*)buffer;

                                // quit, if the main team we're interested in is gone
                                if (fMainTeam >= 0 && event->team == fMainTeam) {
                                        usableBufferSize = buffer + header->size - bufferStart;
                                        quit = true;
                                        break;
                                }
                        }

                        buffer += header->size;
                }

                // write buffer to file
                if (usableBufferSize > 0) {
                        status_t error = fOutput.Write(bufferStart, usableBufferSize);
                        if (error != B_OK) {
                                fprintf(stderr, "%s: Failed to write buffer: %s\n",
                                        kCommandName, strerror(error));
                                quit = true;
                        }
                }

                return quit;
        }


        static void _SignalHandler(int signal, void* data)
        {
                Recorder* self = (Recorder*)data;
                self->fCaughtDeadlySignal = true;
        }

private:
        BFile                                   fOutputFile;
        BDebugEventOutputStream fOutput;
        team_id                                 fMainTeam;
        bool                                    fSkipLoading;
        bool                                    fCaughtDeadlySignal;
};


int
main(int argc, const char* const* argv)
{
        Recorder recorder;

        bool dumpRecorded = false;
        while (true) {
                static struct option sLongOptions[] = {
                        { "help", no_argument, 0, 'h' },
                        { 0, 0, 0, 0 }
                };

                opterr = 0; // don't print errors
                int c = getopt_long(argc, (char**)argv, "+hlr", sLongOptions, NULL);
                if (c == -1)
                        break;

                switch (c) {
                        case 'h':
                                print_usage_and_exit(false);
                                break;
                        case 'l':
                                recorder.SetSkipLoading(false);
                                break;

                        case 'r':
                                dumpRecorded = true;
                                break;

                        default:
                                print_usage_and_exit(true);
                                break;
                }
        }

        // Remaining arguments should be the output file and the optional command
        // line.
        if (optind >= argc)
                print_usage_and_exit(true);

        const char* outputFile = argv[optind++];
        const char* const* programArgs = argv + optind;
        int programArgCount = argc - optind;

        // prepare for battle
        if (recorder.Init(outputFile) != B_OK)
                exit(1);

        // start the action
        if (!dumpRecorded)
                recorder.Run(programArgs, programArgCount);
        else
                recorder.DumpRecorded();

        return 0;
}