#include "core/memory_overlay.h"
#include "imgui.h"

#ifdef TRACY_ENABLE
#include "tracy/Tracy.hpp"
#endif

#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <psapi.h>
#pragma comment(lib, "psapi.lib")
#else
#include <cstdio>
#endif

#include <cstring>

namespace {

struct MemStats {
    long rss_kb    = -1; // Resident Set Size
    long peak_kb   = -1; // Peak RSS
    long vsize_kb  = -1; // Virtual memory size
};

static MemStats s_cached;
static double   s_last_sample = -1.0;

#ifdef _WIN32
static MemStats sample_memory() {
    MemStats s;
    PROCESS_MEMORY_COUNTERS pmc{};
    if (GetProcessMemoryInfo(GetCurrentProcess(), &pmc, sizeof(pmc))) {
        s.rss_kb  = static_cast<long>(pmc.WorkingSetSize       / 1024);
        s.peak_kb = static_cast<long>(pmc.PeakWorkingSetSize   / 1024);
        s.vsize_kb = -1; // not easily available on Windows without VirtualQuery loop
    }
    return s;
}
#else
static MemStats sample_memory() {
    MemStats s;
    FILE* f = std::fopen("/proc/self/status", "r");
    if (!f) return s;

    char line[128];
    while (std::fgets(line, sizeof(line), f)) {
        long val = 0;
        if (std::sscanf(line, "VmRSS: %ld kB",  &val) == 1) s.rss_kb   = val;
        if (std::sscanf(line, "VmPeak: %ld kB", &val) == 1) s.peak_kb  = val;
        if (std::sscanf(line, "VmSize: %ld kB", &val) == 1) s.vsize_kb = val;
    }
    std::fclose(f);
    return s;
}
#endif

} // namespace

void memory_overlay() {
#ifdef TRACY_ENABLE
    ZoneScoped;
#endif

    // Sample at most once per second
    const double now = ImGui::GetTime();
    if (s_last_sample < 0.0 || (now - s_last_sample) >= 1.0) {
        s_cached      = sample_memory();
        s_last_sample = now;
    }

    ImGuiWindowFlags flags = ImGuiWindowFlags_NoDecoration
                           | ImGuiWindowFlags_AlwaysAutoResize
                           | ImGuiWindowFlags_NoSavedSettings
                           | ImGuiWindowFlags_NoFocusOnAppearing
                           | ImGuiWindowFlags_NoNav
                           | ImGuiWindowFlags_NoMove;

    const float pad = 10.0f;
    const ImGuiViewport* vp = ImGui::GetMainViewport();
    ImVec2 pos(vp->WorkPos.x + vp->WorkSize.x - pad,
               vp->WorkPos.y + vp->WorkSize.y - pad);
    ImGui::SetNextWindowPos(pos, ImGuiCond_Always, ImVec2(1.0f, 1.0f));
    ImGui::SetNextWindowBgAlpha(0.65f);

    if (ImGui::Begin("##memory_overlay", nullptr, flags)) {
        if (s_cached.rss_kb >= 0) {
            ImGui::Text("RSS:    %5ld MB", s_cached.rss_kb   / 1024);
            ImGui::Text("Peak:   %5ld MB", s_cached.peak_kb  / 1024);
#ifndef _WIN32
            ImGui::Text("VSize:  %5ld MB", s_cached.vsize_kb / 1024);
#endif
        } else {
            ImGui::TextDisabled("Memory: N/A");
        }
    }
    ImGui::End();
}
