#include "MemoryAllocationStackTracer.h"
#include "DateTime.h"
#include <boost/stacktrace/detail/addr_base.hpp>
#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(1, static_cast<std::size_t>(-1)));
    }
    //    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) {
    void *address = __real_malloc(size);
    auto self = MemoryAllocationStackTracer::instance();
    if (address != nullptr && self->enabled()) {
        self->push(reinterpret_cast<uint64_t>(address),
                   MemoryAllocationStackTracer::Frame(1, static_cast<std::size_t>(-1)));
    }
    return address;
}

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

void *operator new[](std::size_t count) {
    void *address = __real_malloc(count);
    auto self = MemoryAllocationStackTracer::instance();
    if (address != nullptr && self->enabled()) {
        self->push(reinterpret_cast<uint64_t>(address),
                   MemoryAllocationStackTracer::Frame(1, static_cast<std::size_t>(-1)));
    }
    return address;
}

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

void operator delete(void *ptr, std::size_t sz, std::align_val_t al) noexcept {
    return free(ptr);
}
void operator delete[](void *ptr, std::size_t sz, std::align_val_t al) noexcept {
    return free(ptr);
}

MemoryAllocationStackTracer *MemoryAllocationStackTracer::m_instance = nullptr;

void MemoryAllocationStackTracer::push(uint64_t address, Frame &&stacktrace) {
    if (stacktrace.empty()) return;
    auto beginFrame = stacktrace.cbegin()->address();
    if (m_counts.count(beginFrame) <= 0) {
        m_counts.insert({beginFrame, 1});
    } else {
        m_counts[beginFrame]++;
    }

    Infomation info;
    info.time = std::chrono::system_clock::now();
    info.frame = std::move(stacktrace);
    m_stacktraces.insert({address, info});
}

void MemoryAllocationStackTracer::pop(uint64_t address) {
    if (m_stacktraces.count(address) > 0) {
        auto &info = m_stacktraces.at(address);
        auto beginFrame = info.frame.cbegin()->address();
        if (m_counts.count(beginFrame) > 0) {
            if (m_counts.at(beginFrame) > 0) {
                m_counts[beginFrame]--;
            }
        }
        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;
}

std::string MemoryAllocationStackTracer::dump() {
    using namespace std::chrono;
    m_enabled = false;
    std::ostringstream oss;
    oss << "size: " << m_stacktraces.size() << std::endl;
    int index = 0;
    auto now = std::chrono::system_clock::now();
    for (auto &info : m_stacktraces) {
        auto duration = duration_cast<seconds>(now - info.second.time);
        oss << "-----index[" << index << "] duration: " << duration.count() << "s:-----" << std::endl;
        oss << info.second.frame << std::endl;
        index++;
    }
    // aarch64-linux-gnu-addr2line -e build/Tools/LeakTracer/LeakTracer 0x55caea8e4f62

    for (auto &count : m_counts) {
        auto addr_base = boost::stacktrace::detail::get_own_proc_addr_base(count.first);
        oss << "address: " << std::hex << (count.first - addr_base) << ", count: " << count.second << std::endl;
    }
    m_enabled = true;
    return oss.str();
}

void MemoryAllocationStackTracer::run() {
    using namespace std::chrono;
    using namespace std::chrono_literals;
    while (!m_exit) {
        auto start = system_clock::now();

        while (!m_exit) {
            std::this_thread::sleep_for(50ms);
            auto now = system_clock::now();
            auto elapsed = duration_cast<milliseconds>(now - start);
            if (elapsed <= m_intervals) continue;

            std::ostringstream oss;
            oss << m_path << "/" << m_applicationName << "_" << DateTime(start).toString("%Y%m%d%H%M%S") << "-"
                << DateTime::currentDateTime().toString("%Y%m%d%H%M%S") << ".dump";

            std::ofstream ofs(oss.str());
            ofs << dump();
            break;
        }
    }
}

void MemoryAllocationStackTracer::stop() {
    if (!m_exit) {
        m_exit = true;
        if (m_thread.joinable()) m_thread.join();
    }
    if (m_enabled) {
        m_enabled = false;
    }
    m_stacktraces.clear();
    m_counts.clear();
}

MemoryAllocationStackTracer::~MemoryAllocationStackTracer() {
    stop();
}