#ifndef _MEM_MONITOR_H_
#define _MEM_MONITOR_H_

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/time.h>
#include <sys/resource.h>

// 获取当前进程 RSS (物理内存) 使用量，单位 MB
static inline long getRSS_KB()
{
    long rss = 0;
    FILE *fp = fopen("/proc/self/status", "r");
    if (!fp) return -1;
    
    char line[256];
    while (fgets(line, sizeof(line), fp))
    {
        if (strncmp(line, "VmRSS:", 6) == 0)
        {
            // 格式: "VmRSS:    12345 kB"
            const char *p = line + 6;
            while (*p == ' ' || *p == '\t') p++;
            rss = atol(p);
            break;
        }
    }
    fclose(fp);
    return rss;
}

// 获取当前进程内存使用，打印格式化日志
static inline void printMemUsage(const char *tag, const char *extra)
{
    long rss = getRSS_KB();
    struct rusage usage;
    getrusage(RUSAGE_SELF, &usage);
    printf("[MEM] %-30s | VmRSS: %6ld MB | maxRSS: %6ld MB | %s\n",
           tag, rss / 1024, usage.ru_maxrss / 1024, extra ? extra : "");
}

// 获取产品检测 pipeline 队列深度信息快照
#define MEM_LOG(tag, fmt, ...) \
    do { \
        long _rss = getRSS_KB(); \
        printf("[MEM] %-30s | VmRSS: %6ld MB | " fmt "\n", tag, _rss / 1024, ##__VA_ARGS__); \
    } while(0)

// 获取 GPU 显存使用情况（通过 nvidia-smi），返回所有 GPU 累加的已用/总量，单位 MB
// 成功返回 0，失败返回 -1（此时 usedMB/totalMB 置为 -1）
static inline int getGPUMemInfo(long &usedMB, long &totalMB)
{
    usedMB = -1;
    totalMB = -1;
    FILE *fp = popen("nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader,nounits 2>/dev/null", "r");
    if (!fp)
    {
        return -1;
    }
    char line[256];
    long used = 0;
    long total = 0;
    int n = 0;
    while (fgets(line, sizeof(line), fp))
    {
        long u = 0;
        long t = 0;
        if (sscanf(line, "%ld, %ld", &u, &t) == 2)
        {
            used += u;
            total += t;
            n++;
        }
    }
    pclose(fp);
    if (n == 0)
    {
        return -1;
    }
    usedMB = used;
    totalMB = total;
    return 0;
}

// 获取系统内存信息（/proc/meminfo），单位 KB
// 成功返回 0，失败返回 -1（此时各值置为 -1）
static inline int getSystemMemInfo(long &totalKB, long &freeKB, long &availableKB)
{
    totalKB = -1;
    freeKB = -1;
    availableKB = -1;
    FILE *fp = fopen("/proc/meminfo", "r");
    if (!fp)
    {
        return -1;
    }
    char line[256];
    while (fgets(line, sizeof(line), fp))
    {
        long val = -1;
        if (sscanf(line, "MemTotal: %ld kB", &val) == 1)
        {
            totalKB = val;
        }
        else if (sscanf(line, "MemFree: %ld kB", &val) == 1)
        {
            freeKB = val;
        }
        else if (sscanf(line, "MemAvailable: %ld kB", &val) == 1)
        {
            availableKB = val;
        }
    }
    fclose(fp);
    return 0;
}

// 系统 CPU 时间快照（/proc/stat 第一行 cpu 累计值），用于两次采样计算 CPU 使用率
struct CPUStat
{
    unsigned long long user;
    unsigned long long nice;
    unsigned long long system;
    unsigned long long idle;
    unsigned long long iowait;
    unsigned long long irq;
    unsigned long long softirq;
    unsigned long long steal;

    CPUStat()
    {
        user = nice = system = idle = iowait = irq = softirq = steal = 0;
    }
    unsigned long long total() const
    {
        return user + nice + system + idle + iowait + irq + softirq + steal;
    }
    unsigned long long idleTotal() const
    {
        return idle + iowait;
    }
};

// 读取系统 CPU 时间快照，成功返回 0，失败返回 -1
static inline int getSystemCPUStat(CPUStat &stat)
{
    FILE *fp = fopen("/proc/stat", "r");
    if (!fp)
    {
        return -1;
    }
    char line[256];
    int ret = -1;
    if (fgets(line, sizeof(line), fp))
    {
        // 格式: cpu  user nice system idle iowait irq softirq steal guest guest_nice
        if (sscanf(line, "cpu %llu %llu %llu %llu %llu %llu %llu %llu",
                   &stat.user, &stat.nice, &stat.system, &stat.idle,
                   &stat.iowait, &stat.irq, &stat.softirq, &stat.steal) >= 4)
        {
            ret = 0;
        }
    }
    fclose(fp);
    return ret;
}

// 进程累计消耗的 CPU 时间（用户态 + 内核态），单位秒
static inline double cpuTimeSec(const struct rusage &usage)
{
    return (double)usage.ru_utime.tv_sec + (double)usage.ru_utime.tv_usec / 1e6 +
           (double)usage.ru_stime.tv_sec + (double)usage.ru_stime.tv_usec / 1e6;
}

#endif // _MEM_MONITOR_H_