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

extern "C" {

void *__wrap_malloc(size_t size) {
    void *address = __real_malloc(size);
    auto self = MemoryAllocationStackTracer::instance();
    if (address != nullptr && self->enabled()) {
        self->push(reinterpret_cast<uint64_t>(address), MemoryAllocationStackTracer::Frame());
    }
    //    printf("%p = malloc(%zu)\n", address, size);
    return address;
}

void __wrap_free(void *__ptr) {
    if (__ptr != nullptr) {
        MemoryAllocationStackTracer::instance()->pop(reinterpret_cast<uint64_t>(__ptr));
    }
    //    printf("free(%p)\n", __ptr);
    return __real_free(__ptr);
}
}

void *operator new(std::size_t size) {
    // printf("void *operator new(std::size_t size),%zu\n", size);

    return malloc(size);
}

void operator delete(void *p) noexcept {
    //    printf("%s", "void operator delete(void *p)\n");
    return free(p);
}

void *operator new[](std::size_t count) {
    //    printf("%s", "void *operator new[](std::size_t count)\n");
    return malloc(count);
}

void operator delete[](void *ptr) noexcept {
    //    printf("%s", "void operator delete[](void *ptr) noexcept\n");
    return free(ptr);
}

MemoryAllocationStackTracer *MemoryAllocationStackTracer::m_instance = nullptr;

void MemoryAllocationStackTracer::start() {
    if (!m_enabled) {
        m_enabled = true;
    }
}

void MemoryAllocationStackTracer::push(uint64_t address, Frame &&stacktrace) {
    Infomation info;
    info.time = std::chrono::system_clock::now();
    info.frame = std::move(stacktrace);
    m_stacktraces.insert({address, info});
    // m_stacktraces.insert({address, stacktrace});
}

void MemoryAllocationStackTracer::pop(uint64_t address) {
    if (m_stacktraces.count(address) > 0) {
        m_stacktraces.erase(address);
    }
}

MemoryAllocationStackTracer *MemoryAllocationStackTracer::instance() {
    if (m_instance == nullptr) {
        auto buffer = __real_malloc(sizeof(MemoryAllocationStackTracer));
        if (buffer != nullptr) {
            m_instance = new (buffer) MemoryAllocationStackTracer();
        }
    }
    return m_instance;
}
void MemoryAllocationStackTracer::dump() {
    using namespace std::chrono;
    m_enabled = false;
    std::cout << "size: " << m_stacktraces.size() << std::endl;
    int index = 0;
    auto now = std::chrono::system_clock::now();
    std::ofstream ofs("test.txt");
    for (auto &info : m_stacktraces) {
        std::ostringstream oss;
        auto duration = duration_cast<seconds>(now - info.second.time);
        oss << "-----index[" << index << "] duration: " << duration.count() << "s:-----" << std::endl;
        oss << info.second.frame << std::endl;
        auto content = oss.str();
        ofs << content;
        std::cout << content;
        index++;
    }
    m_enabled = true;
}