#include <stdint.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <string.h>

#include <Windows.h>

#define INITGUID
#include <evntrace.h>
#include <evntprov.h>
#include <evntcons.h>

#include "events.h"

typedef uint64_t    u64;
typedef uint32_t    u32;
typedef uint8_t     u8;
typedef int8_t      i8;

DEFINE_GUID(PMC_TRACE_GUID, 0x7e5ac468, 0x6864, 0x4b43, 0xa0, 0xf4, 0xef, 0x4d, 0xa7, 0x6e, 0x07, 0x4a);
DEFINE_GUID(PMC_TRACE_START_GUID, 0x7e5ac468, 0x6864, 0x4b43, 0xa0, 0xf4, 0xef, 0x4d, 0xa7, 0x6e, 0x07, 0x4b);

enum PMC_CAPTURE_STATUS {
    PMC_CAPTURE_STATUS_INACTIVE = 0U,
    PMC_CAPTURE_STATUS_STARTED = 1U,
    PMC_CAPTURE_STATUS_STOPPED = 2U,
};

struct PMC_marker_data {
    u32 unique_id;
    const char* name;
    u32 thread_id;
};
struct begin_capture_event {
    EVENT_TRACE_HEADER header;
    struct PMC_marker_data data;
};

struct end_capture_event {
    EVENT_TRACE_HEADER header;
    struct PMC_marker_data data;
};

struct capture {
    TRACEHANDLE trace_handle;
    TRACEHANDLE logger_handle;
    TRACEHANDLE event_provider_handle;
};

#define REAL_TIME_MODE

#if !defined(REAL_TIME_MODE)
#define LOGFILE_NAME "test_etw_log.etl"
#endif

struct trace_process_ctx {
    u32 PMCs_count;
    u32* PMCs;
    u64* values;
    u64* start_values;
    u32 thread_id;
    u8* started;
};

#define SOURCE_TIMER "Timer"
#define SOURCE_TIMER_INDEX 0U
#define SOURCE_TOTAL_ISSUES "TotalIssues"
#define SOURCE_TOTAL_ISSUES_INDEX 2U
#define SOURCE_BRANCH_INSTRUCTIONS "BranchInstructions"
#define SOURCE_BRANCH_INSTRUCTIONS_INDEX 6U

const char* PMC_SOURCES[] = {
    SOURCE_TIMER,
    "",
    SOURCE_TOTAL_ISSUES,
    "",
    "",
    "",
    SOURCE_BRANCH_INSTRUCTIONS,
    "",
    "DcacheMisses",
    "IcacheMisses",
    "CacheMisses",
    "BranchMispredictions",
    "",
    "",
    "",
    "",
    "",
    "",
    "",
    "TotalCycles",
    "IcacheIssues",
    "DcacheAccesses",
    "",
    "",
    "",
    "",
    "InstructionsRetired",
    "IcacheL2Hit",
    "DcacheL2Hit",
    "RetiredMacroOps",
    "RetiredTakenBranchInstructions",
    "RetiredTakenBranchMispredictions",
    "RetiredFarControlTransfers",
    "RetiredNearReturns",
    "RetiredNearReturnMispredictions",
    "RetiredIndirectBranchMispredictions",
    "RetiredConditionalBranchInstructions",
    "RetiredConditionalBranchMispredictions",
    "RetiredFusedInstructions",
    "VariableTargetPredictions",
    "EarlyRedirects",
    "InefectiveSoftwarePrefetches",
    "TlbFlushes",
    "SmiReceived",
    "DemandDcacheFills",
    "RetiredLockInstructions",
    "BusLock",
    "Retiredx87Instructions",
    "RetiredMmxInstructions",
    "RetiredSseInstructions",
};
const u32 PMC_SOURCES_COUNT = sizeof(PMC_SOURCES) / sizeof(PMC_SOURCES[0U]);

void log_error(DWORD dw) {
    // Retrieve the system error message for the last-error code
    LPVOID lpMsgBuf;

    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | 
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        dw,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR) &lpMsgBuf,
        0, NULL );

    // Display the error message and exit the process
    printf("error %lu: %s\n", dw, (char*)lpMsgBuf);

    LocalFree(lpMsgBuf);
}

#define TRACE_NAME "test_etw"

struct event_trace_property_data {
    EVENT_TRACE_PROPERTIES_V2 props;
#if !defined(REAL_TIME_MODE)
    char log_file_name[sizeof(LOGFILE_NAME)];
#endif
    char logger_name[sizeof(TRACE_NAME)];
} data = { 0U };

ULONG WINAPI Wmidprequest(WMIDPREQUESTCODE RequestCode, PVOID RequestContext, ULONG *BufferSize, PVOID Buffer) {
    return 0U;
}

void PeventRecordCallback(PEVENT_RECORD event_record) {
    struct trace_process_ctx* ctx = (struct trace_process_ctx*)event_record->UserContext;
    GUID provider = event_record->EventHeader.ProviderId;

    if (memcmp(&provider, &PMC_TRACE_GUID, sizeof(provider)) == 0) {
        *ctx->started = PMC_CAPTURE_STATUS_STARTED;
        return;
    }

    if (memcmp(&provider, &PMC_TRACE_GUID, sizeof(provider)) == 0) {
        *ctx->started = PMC_CAPTURE_STATUS_STOPPED;
        return;
    }

    u32 opcode = event_record->EventHeader.EventDescriptor.Opcode;
    GUID thread_guid = THREAD_EVENTS_GUID;
    GUID perf_info_guid = PERF_INFO_EVENTS_GUID;
    if (memcmp(&provider, &thread_guid, sizeof(provider)) != 0 || opcode != EventThreadContextSwitch.Type) {
        return;
    }

    if (memcmp(&provider, &perf_info_guid, sizeof(provider)) != 0 || (opcode != EventSysCallEnter.Type && opcode != EventSysCallExit.Type)) {
        return;
    }

    PEVENT_HEADER_EXTENDED_DATA_ITEM cur_data = event_record->ExtendedData;
    for (u32 index = 0U; index < event_record->ExtendedDataCount; index++, cur_data++) {
        if (cur_data->ExtType == EVENT_HEADER_EXT_TYPE_PMC_COUNTERS) {
            ULONG64* pmc = (ULONG64*)cur_data->DataPtr;
            size_t pmc_count = cur_data->DataSize / sizeof(pmc[0U]);

            if (opcode == EventThreadContextSwitch.Type) {
                u32 new_thread_id = *(u32*)event_record->UserData;
                u32 old_thread_id = *((u32*)event_record->UserData + 1U);
                if (new_thread_id == ctx->thread_id) {
                    for (u32 pmc_index = 0U; pmc_index < pmc_count; pmc_index++) {
                        ctx->start_values[pmc_index] = pmc[pmc_index];
                    }
                } else if (old_thread_id == ctx->thread_id) {
                    for (u32 pmc_index = 0U; pmc_index < pmc_count; pmc_index++) {
                        if (ctx->start_values[pmc_index] != 0U) {
                            ctx->values[pmc_index] += pmc[pmc_index] - ctx->start_values[pmc_index];
                        }
                    }
                }
            } else if (opcode == EventSysCallExit.Type) {
                for (u32 pmc_index = 0U; pmc_index < pmc_count; pmc_index++) {
                    if (ctx->start_values[pmc_index] == 0U) {
                        ctx->start_values[pmc_index] = pmc[pmc_index];
                    }
                }
            } else {
                for (u32 pmc_index = 0U; pmc_index < pmc_count; pmc_index++) {
                    if (ctx->start_values[pmc_index] != 0U) {
                        ctx->values[pmc_index] += pmc[pmc_index] - ctx->start_values[pmc_index];
                    }
                }
            }
        }
    }
}

u32 end_trace(struct capture* capture) {
    ULONG res = 0U;

    if ((res = UnregisterTraceGuids(capture->event_provider_handle)) != ERROR_SUCCESS) {
        log_error(res);
    }

    if ((res = ControlTraceA(capture->trace_handle, TRACE_NAME, (EVENT_TRACE_PROPERTIES*)&data.props, EVENT_TRACE_CONTROL_STOP)) != ERROR_SUCCESS) {
        log_error(res);
        return 1U;
    }

    return 0U;
}

u32 begin_trace(struct capture* capture, u32* pmc_sources, u32 sources_size) {
    ULONG res = 0U;

    data.props.Wnode.BufferSize = sizeof(struct event_trace_property_data);
    data.props.Wnode.Flags      = WNODE_FLAG_TRACED_GUID | WNODE_FLAG_VERSIONED_PROPERTIES;

    data.props.VersionNumber    = 2UL;
    data.props.BufferSize       = 16384U;
    data.props.MinimumBuffers   = 4U;
    data.props.MaximumBuffers   = 4U;
    data.props.MaximumFileSize  = 0U;
    data.props.VersionNumber    = 2U;
    data.props.EnableFlags      = EVENT_TRACE_FLAG_NO_SYSCONFIG |
                                  EVENT_TRACE_FLAG_CSWITCH |
                                  EVENT_TRACE_FLAG_SYSTEMCALL;
    data.props.LogFileMode      = EVENT_TRACE_SYSTEM_LOGGER_MODE;
    data.props.LoggerNameOffset = offsetof(struct event_trace_property_data, logger_name),

#if defined(REAL_TIME_MODE)
    data.props.LogFileMode      |= EVENT_TRACE_REAL_TIME_MODE;
#else
    data.props.log_file_name    = LOGFILE_NAME,
    data.props.LogFileNameOffset= offsetof(struct event_trace_property_data, log_file_name),
#endif

    strcpy(TRACE_NAME, data.logger_name);

    if ((res = StartTraceA(&capture->trace_handle, TRACE_NAME, (EVENT_TRACE_PROPERTIES*)&data.props)) != ERROR_SUCCESS) {
        log_error(res);
        return 1U;
    }

    TRACE_GUID_REGISTRATION trace_registration = {
        .Guid = &PMC_TRACE_START_GUID,
        .RegHandle = NULL,
    };
    if ((res = RegisterTraceGuidsA((WMIDPREQUEST)Wmidprequest, NULL, &PMC_TRACE_GUID, 1U, &trace_registration, NULL, NULL, &capture->event_provider_handle)) != ERROR_SUCCESS) {
        log_error(res);
        return 1U;
    }

    // if ((res = EnableTraceEx2(capture->trace_handle, &PMC_TRACE_GUID, EVENT_CONTROL_CODE_ENABLE_PROVIDER, TRACE_LEVEL_VERBOSE, 0U, 0U, 0U, NULL) != ERROR_SUCCESS)) {
    //     log_error(res);
    //     return 1U;
    // }

    if ((res = TraceSetInformation(capture->trace_handle, TracePmcCounterListInfo, pmc_sources, sources_size)) != ERROR_SUCCESS) { // PMCs to collect
        log_error(res);
        return 1U;
    }

    CLASSIC_EVENT_ID events[] = {
        EventThreadContextSwitch,
        EventSysCallExit,
        EventSysCallEnter,
    };
    if ((res = TraceSetInformation(capture->trace_handle, TracePmcEventListInfo, &events, sizeof(events))) != ERROR_SUCCESS) { // Events on which we collect PMCs
        log_error(res);
        return 1U;
    }

    return 0U;
}

u32 begin_PMCs_capture(struct capture* capture) {
    struct begin_capture_event event = {
        .header = {
            .Size = sizeof(event),
            .Guid = PMC_TRACE_START_GUID,
            .Flags = WNODE_FLAG_TRACED_GUID,
        },
        .data = {
            .thread_id = GetThreadId(GetCurrentThread()),
        },
    };
    ULONG res = TraceEvent(capture->trace_handle, &event.header);
    if (res != ERROR_SUCCESS) {
        log_error(res);
        return 1U;
    }
    return 0U;
}

u32 end_PMCs_capture(struct capture* capture) {
    struct end_capture_event event = {
        .header = {
            .Size = sizeof(event),
            .Guid = PMC_TRACE_GUID,
            .Flags = WNODE_FLAG_TRACED_GUID,
        },
        .data = {
            .thread_id = GetThreadId(GetCurrentThread()),
        },
    };
    ULONG res = TraceEvent(capture->trace_handle, &event.header);
    if (res != ERROR_SUCCESS) {
        log_error(res);
        return 1U;
    }
    return 0U;
}

u32 PMC_sources_from_names(const char** names, u32 names_count, u32* sources, u32* sources_count) {
    for (u32 name_index = 0U; name_index < names_count; name_index++) {
        for (u32 list_index = 0U; list_index < PMC_SOURCES_COUNT; list_index++) {
            if (strcmp(names[name_index], PMC_SOURCES[list_index]) == 0U) {
                sources[*sources_count] = list_index;
                (*sources_count)++;
            }
        }
    }

    return (names_count != *sources_count);
}

u32 list_sources() {
    char* profile_source_infos[sizeof(PROFILE_SOURCE_INFO) * 256U];
    ULONG bwrite = 0U;
    ULONG res = 0U;
    if ((res = TraceQueryInformation(0U, TraceProfileSourceListInfo, (void*)profile_source_infos, sizeof(profile_source_infos), &bwrite)) != ERROR_SUCCESS) {
        log_error(res);
        return 1U;
    }

    PROFILE_SOURCE_INFO* cur_profile_source_info = (PROFILE_SOURCE_INFO*)&profile_source_infos[0U];
    while (1) {
        printf("%lu: %S\n", cur_profile_source_info->Source, cur_profile_source_info->Description);

        if (cur_profile_source_info->NextEntryOffset == 0U) {
            break;
        }

        cur_profile_source_info = (PROFILE_SOURCE_INFO*)((char*)cur_profile_source_info + cur_profile_source_info->NextEntryOffset);
    }

    return 0U;
}

u32 PMCs[] = {
    SOURCE_TOTAL_ISSUES_INDEX,
    SOURCE_BRANCH_INSTRUCTIONS_INDEX,
};
u64 PMCs_values[sizeof(PMCs) / sizeof(PMCs[0U])] = {
    0U,
};
u64 PMCs_start_values[sizeof(PMCs) / sizeof(PMCs[0U])] = {
    0U,
};

u8 started = 0U;
u32 open_trace(struct capture* capture) {
    ULONG res = 0U;
    struct trace_process_ctx ctx = {
        .PMCs_count = sizeof(PMCs) / sizeof(PMCs[0U]),
        .PMCs = PMCs,
        .values = PMCs_values,
        .start_values = PMCs_start_values,
        .thread_id = GetThreadId(GetCurrentThread()),
        .started = &started
    };

    EVENT_TRACE_LOGFILEA trace_log = { 0U };

    trace_log.EventRecordCallback = PeventRecordCallback;
    trace_log.Context = &ctx;
#if defined(REAL_TIME_MODE)
    trace_log.LoggerName = TRACE_NAME;
    trace_log.ProcessTraceMode = PROCESS_TRACE_MODE_REAL_TIME | PROCESS_TRACE_MODE_EVENT_RECORD | PROCESS_TRACE_MODE_RAW_TIMESTAMP;
#else
    trace_log.LogFileName = "test_etw_log.etl";
    trace_log.ProcessTraceMode = PROCESS_TRACE_MODE_EVENT_RECORD;
#endif

    if ((capture->logger_handle = OpenTraceA(&trace_log)) == INVALID_PROCESSTRACE_HANDLE) {
        log_error(GetLastError());
        return 1U;
    }
    return 0U;
}

#ifdef REAL_TIME_MODE
DWORD WINAPI process_event_routine(LPVOID params) {
    ULONG res = 0U;
    struct capture* capture = params;

    if (open_trace(capture)) {
        ExitThread(1U);
        return 1U;
    }

    if ((res = ProcessTrace(&capture->logger_handle, 1U, NULL, NULL)) != ERROR_SUCCESS) {
        log_error(res);
        ExitThread(1U);
        return 1U;
    }

    if ((res = CloseTrace(capture->logger_handle)) != ERROR_SUCCESS) {
        log_error(res);
        ExitThread(1U);
        return 1U;
    }

    ExitThread(0U);
    return 0U;
}
#endif

int main() {
    struct capture capture;
    static_assert(sizeof(PMCs) / sizeof(PMCs[0U]) <= 4U, "Too much PMCs");
    if (begin_trace(&capture, PMCs, sizeof(PMCs))) {
        end_trace(&capture);
        return 1;
    }

    ULONG res = 0U;
#ifdef REAL_TIME_MODE
    DWORD process_event_thread_id = 0U;
    HANDLE handle_process_event_thread = CreateThread(NULL, 64U * 1024U, process_event_routine, (LPVOID)&capture, 0U, &process_event_thread_id);
#endif

    int j = 0;
    while (j < 4) {
        Sleep(1000U);
        j++;
    }

    if (begin_PMCs_capture(&capture)) {
        end_trace(&capture);
        return 1;
    }

    int i = 0;
    float f = 0.5f;
    while (i < 4) {
        f = sinf(f);
        i++;
    }

    if (end_PMCs_capture(&capture)) {
        end_trace(&capture);
        return 1;
    }

#if defined(REAL_TIME_MODE)
    WaitForSingleObject(handle_process_event_thread, INFINITE);
    CloseHandle(handle_process_event_thread);
#endif

    if (end_trace(&capture)) {
        return 1;
    }

#if !defined(REAL_TIME_MODE)
    if (open_trace(&capture)) {
        return 1;
    }

    if ((res = ProcessTrace(&capture.logger_handle, 1U, NULL, NULL)) != ERROR_SUCCESS) {
        log_error(res);
    }

    if ((res = CloseTrace(capture.logger_handle)) != ERROR_SUCCESS) {
        log_error(res);
        return 1;
    }
#endif

    return 0;
}