#include "MemoryAllocationStackTracer.h"
#include <iostream>
#include <stdio.h>

extern "C" {

void *__wrap_malloc(size_t size) {
    void *address = __real_malloc(size);
    // if (address != nullptr) {
    //     MemoryAllocationStackTracer::instance()->push(reinterpret_cast<uint64_t>(address),
    //                                                   boost::stacktrace::stacktrace());
    // }
    // std::cout << "malloc " << size << " bytes." << std::endl;
    return address;
}

void __wrap_free(void *__ptr) {
    // if (__ptr != nullptr) {
    //     MemoryAllocationStackTracer::instance()->pop(reinterpret_cast<uint64_t>(__ptr));
    // }
    return __real_free(__ptr);
}
}

void *operator new(std::size_t size) {
    printf("new %d %d\n ", size, sizeof(boost::stacktrace::stacktrace));
    void *address = __real_malloc(size);

    // if (size != sizeof(boost::stacktrace::stacktrace) || size != 2048) {
    //     MemoryAllocationStackTracer::instance()->push(reinterpret_cast<uint64_t>(address),
    //                                                   boost::stacktrace::stacktrace());
    // }

    return address;
}

void operator delete(void *p) {
    free(p);
}

MemoryAllocationStackTracer *MemoryAllocationStackTracer::m_instance = nullptr;

void MemoryAllocationStackTracer::initialize() {
    if (m_instance == nullptr) {
        auto buffer = __real_malloc(sizeof(MemoryAllocationStackTracer));
        if (buffer != nullptr) {
            m_instance = new (buffer) MemoryAllocationStackTracer();
        }
    }
}

void MemoryAllocationStackTracer::push(uint64_t address, const boost::stacktrace::stacktrace &stacktrace) {
    m_stacktraces.insert({address, stacktrace});
}

void MemoryAllocationStackTracer::pop(uint64_t address) {
    m_stacktraces.erase(address);
}

MemoryAllocationStackTracer *MemoryAllocationStackTracer::instance() {

    return m_instance;
}