#include "app/gpu_hogs.h"

#include <windows.h>

#include <pdh.h>
#include <pdhmsg.h>
#include <psapi.h>

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <vector>

namespace sauna {

std::string TopGpuConsumers(int sampleMs, int maxEntries) {
  // Every failure path reports itself — a silent "" during a live episode
  // is a wasted forensic opportunity (field finding: first episode on the
  // instrumented build printed nothing).
  char err[64];
  PDH_HQUERY query = nullptr;
  PDH_STATUS st = PdhOpenQueryW(nullptr, 0, &query);
  if (st != ERROR_SUCCESS) {
    snprintf(err, sizeof(err), "(PdhOpenQuery 0x%08lx)", (unsigned long)st);
    return err;
  }
  std::string out;
  PDH_HCOUNTER counter = nullptr;
  // Instance names look like "pid_1234_luid_..._engtype_3D"; utilization
  // needs two collections to form a rate.
  st = PdhAddEnglishCounterW(query, L"\\GPU Engine(*)\\Utilization Percentage",
                             0, &counter);
  if (st == ERROR_SUCCESS) st = PdhCollectQueryData(query);
  if (st == ERROR_SUCCESS) {
    Sleep(sampleMs);
    st = PdhCollectQueryData(query);
  }
  if (st == ERROR_SUCCESS) {
    DWORD bufSize = 0, count = 0;
    st = PdhGetFormattedCounterArrayW(counter, PDH_FMT_DOUBLE, &bufSize,
                                      &count, nullptr);
    if (st != PDH_MORE_DATA && st != ERROR_SUCCESS) {
      snprintf(err, sizeof(err), "(PdhGetFormattedCounterArray size 0x%08lx)",
               (unsigned long)st);
      PdhCloseQuery(query);
      return err;
    }
    std::vector<uint8_t> buf(bufSize);
    auto* items = (PDH_FMT_COUNTERVALUE_ITEM_W*)buf.data();
    st = bufSize ? PdhGetFormattedCounterArrayW(counter, PDH_FMT_DOUBLE,
                                                &bufSize, &count, items)
                 : (PDH_STATUS)ERROR_SUCCESS;
    if (st == ERROR_SUCCESS) {
      std::map<DWORD, double> byPid;  // all engines summed per process
      for (DWORD i = 0; i < count; i++) {
        const wchar_t* name = items[i].szName;
        if (wcsncmp(name, L"pid_", 4) != 0) continue;
        const DWORD pid = (DWORD)wcstoul(name + 4, nullptr, 10);
        if (pid) byPid[pid] += items[i].FmtValue.doubleValue;
      }
      std::vector<std::pair<DWORD, double>> top(byPid.begin(), byPid.end());
      std::sort(top.begin(), top.end(), [](const auto& a, const auto& b) {
        return a.second > b.second;
      });
      int shown = 0;
      for (const auto& [pid, pct] : top) {
        // Always include the top entry — even a low number is signal
        // (means the load is NOT 3D/compute-engine work).
        if (shown >= maxEntries || (shown > 0 && pct < 2.0)) break;
        char exe[MAX_PATH] = "?";
        HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);
        if (h) {
          DWORD len = MAX_PATH;
          char full[MAX_PATH];
          if (QueryFullProcessImageNameA(h, 0, full, &len)) {
            const char* base = strrchr(full, '\\');
            snprintf(exe, sizeof(exe), "%s", base ? base + 1 : full);
          }
          CloseHandle(h);
        }
        char entry[MAX_PATH + 32];
        snprintf(entry, sizeof(entry), "%s%s(%lu) %.0f%%", shown ? "  " : "",
                 exe, (unsigned long)pid, pct);
        out += entry;
        shown++;
      }
      if (out.empty()) out = "(no per-process GPU activity reported)";
    } else {
      snprintf(err, sizeof(err), "(PdhGetFormattedCounterArray 0x%08lx)",
               (unsigned long)st);
      out = err;
    }
  } else {
    snprintf(err, sizeof(err), "(PDH add/collect 0x%08lx)", (unsigned long)st);
    out = err;
  }
  PdhCloseQuery(query);
  return out;
}

std::string CpuStallForensics(int sampleMs, int maxEntries) {
  char err[64];
  PDH_HQUERY query = nullptr;
  PDH_STATUS st = PdhOpenQueryW(nullptr, 0, &query);
  if (st != ERROR_SUCCESS) {
    snprintf(err, sizeof(err), "(PdhOpenQuery 0x%08lx)", (unsigned long)st);
    return err;
  }
  PDH_HCOUNTER cDpc = nullptr, cInt = nullptr, cPriv = nullptr,
               cProc = nullptr;
  PdhAddEnglishCounterW(query, L"\\Processor(_Total)\\% DPC Time", 0, &cDpc);
  PdhAddEnglishCounterW(query, L"\\Processor(_Total)\\% Interrupt Time", 0,
                        &cInt);
  PdhAddEnglishCounterW(query, L"\\Processor(_Total)\\% Privileged Time", 0,
                        &cPriv);
  PdhAddEnglishCounterW(query, L"\\Process(*)\\% Processor Time", 0, &cProc);
  st = PdhCollectQueryData(query);
  if (st == ERROR_SUCCESS) {
    Sleep(sampleMs);
    st = PdhCollectQueryData(query);
  }
  if (st != ERROR_SUCCESS) {
    snprintf(err, sizeof(err), "(PDH collect 0x%08lx)", (unsigned long)st);
    PdhCloseQuery(query);
    return err;
  }
  auto scalar = [](PDH_HCOUNTER c) {
    PDH_FMT_COUNTERVALUE v{};
    if (!c || PdhGetFormattedCounterValue(c, PDH_FMT_DOUBLE, nullptr, &v) !=
                  ERROR_SUCCESS)
      return -1.0;
    return v.doubleValue;
  };
  char head[96];
  snprintf(head, sizeof(head), "dpc=%.1f%% int=%.1f%% priv=%.0f%%",
           scalar(cDpc), scalar(cInt), scalar(cPriv));
  std::string out = head;
  // Top per-process CPU (instance = exe base name; per-core scale).
  DWORD bufSize = 0, count = 0;
  st = PdhGetFormattedCounterArrayW(cProc, PDH_FMT_DOUBLE, &bufSize, &count,
                                    nullptr);
  if ((st == PDH_MORE_DATA || st == ERROR_SUCCESS) && bufSize) {
    std::vector<uint8_t> buf(bufSize);
    auto* items = (PDH_FMT_COUNTERVALUE_ITEM_W*)buf.data();
    if (PdhGetFormattedCounterArrayW(cProc, PDH_FMT_DOUBLE, &bufSize, &count,
                                     items) == ERROR_SUCCESS) {
      std::vector<std::pair<std::wstring, double>> top;
      for (DWORD i = 0; i < count; i++) {
        if (!wcscmp(items[i].szName, L"_Total") ||
            !wcscmp(items[i].szName, L"Idle"))
          continue;
        top.emplace_back(items[i].szName, items[i].FmtValue.doubleValue);
      }
      std::sort(top.begin(), top.end(), [](const auto& a, const auto& b) {
        return a.second > b.second;
      });
      out += "  top:";
      int shown = 0;
      for (const auto& [name, pct] : top) {
        if (shown >= maxEntries || (shown > 0 && pct < 5.0)) break;
        char entry[96];
        snprintf(entry, sizeof(entry), " %ls %.0f", name.c_str(), pct);
        out += entry;
        shown++;
      }
    }
  }
  PdhCloseQuery(query);
  return out;
}

}  // namespace sauna