#include "device/mcu_prox.h"

#include <windows.h>

#include <hidsdi.h>
#include <setupapi.h>

#include <cstdio>
#include <cstring>
#include <vector>

namespace sauna {
namespace {

constexpr USHORT kVid = 0x35BD;
constexpr USHORT kPid = 0x0101;

// Open every 35BD:0101 HID interface; the caller probes which one speaks the
// control protocol (the MCU exposes several collections).
std::vector<HANDLE> openMcuInterfaces() {
  std::vector<HANDLE> out;
  GUID hidGuid;
  HidD_GetHidGuid(&hidGuid);
  HDEVINFO info = SetupDiGetClassDevsW(&hidGuid, nullptr, nullptr,
                                       DIGCF_PRESENT | DIGCF_DEVICEINTERFACE);
  if (info == INVALID_HANDLE_VALUE) return out;
  SP_DEVICE_INTERFACE_DATA iface{};
  iface.cbSize = sizeof(iface);
  for (DWORD i = 0;
       SetupDiEnumDeviceInterfaces(info, nullptr, &hidGuid, i, &iface); i++) {
    DWORD need = 0;
    SetupDiGetDeviceInterfaceDetailW(info, &iface, nullptr, 0, &need, nullptr);
    std::vector<uint8_t> buf(need);
    auto* detail = reinterpret_cast<SP_DEVICE_INTERFACE_DETAIL_DATA_W*>(buf.data());
    detail->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);
    if (!SetupDiGetDeviceInterfaceDetailW(info, &iface, detail, need, nullptr,
                                          nullptr))
      continue;
    HANDLE h = CreateFileW(detail->DevicePath, GENERIC_READ | GENERIC_WRITE,
                           FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                           OPEN_EXISTING, FILE_FLAG_OVERLAPPED, nullptr);
    if (h == INVALID_HANDLE_VALUE) continue;
    HIDD_ATTRIBUTES attr{};
    attr.Size = sizeof(attr);
    if (HidD_GetAttributes(h, &attr) && attr.VendorID == kVid &&
        attr.ProductID == kPid) {
      out.push_back(h);
    } else {
      CloseHandle(h);
    }
  }
  SetupDiDestroyDeviceInfoList(info);
  return out;
}

bool setFeature(HANDLE h, const uint8_t* payload, size_t len) {
  uint8_t buf[65] = {0};  // report id 0 + payload, padded like bey-closer-t1
  memcpy(buf + 1, payload, len < 64 ? len : 64);
  return HidD_SetFeature(h, buf, sizeof(buf)) != 0;
}

// One overlapped interrupt read. Returns bytes read, 0 on timeout, or -1 on
// a hard error (device unplugged) — callers must NOT retry on -1, a dead
// handle fails instantly and retrying busy-spins.
int readReport(HANDLE h, HANDLE evt, uint8_t* buf, DWORD len, DWORD timeoutMs) {
  OVERLAPPED ov{};
  ov.hEvent = evt;
  ResetEvent(evt);
  DWORD got = 0;
  if (!ReadFile(h, buf, len, &got, &ov)) {
    if (GetLastError() != ERROR_IO_PENDING) return -1;
    if (WaitForSingleObject(evt, timeoutMs) == WAIT_TIMEOUT) {
      CancelIoEx(h, &ov);
      GetOverlappedResult(h, &ov, &got, TRUE);
      return 0;
    }
    if (!GetOverlappedResult(h, &ov, &got, FALSE)) return -1;
  }
  return (int)got;
}

}  // namespace

bool McuProx::sendFeature(const uint8_t* payload, size_t len) {
  return setFeature(dev_, payload, len);
}

bool McuProx::setProxBypass(bool bypass) {
  if (!dev_) return false;  // unplugged / reconnecting
  // Firmware debug commands (usbhid_interface.c): 'p' = prox disabled until
  // power cycle or '[' — firmware believes worn, displays turn on regardless.
  const uint8_t cmd[1] = {bypass ? (uint8_t)'p' : (uint8_t)'['};
  return sendFeature(cmd, sizeof(cmd));
}

bool McuProx::setDoze(bool sleep) {
  if (!dev_) return false;  // unplugged / reconnecting
  // Firmware doze command pair (usbhid_interface.c, fw >= 0.4.0): 'H' =
  // display-sleep, 'h' = display-wake. Firmware owns the choreography.
  const uint8_t cmd[1] = {sleep ? (uint8_t)'H' : (uint8_t)'h'};
  return sendFeature(cmd, sizeof(cmd));
}

bool McuProx::setBrightness(uint16_t value) {
  if (value > 1023) value = 1023;
  if (!dev_) return false;  // unplugged / reconnecting
  const uint8_t cmd[3] = {'I', (uint8_t)(value >> 8), (uint8_t)(value & 0xFF)};
  return sendFeature(cmd, sizeof(cmd));
}

namespace {

// Find the MCU control collection: send 'R' <period BE> to each 35BD:0101
// interface; the right one ACKs '$' on the interrupt pipe (skip any '#'
// telemetry already streaming). Returns INVALID_HANDLE_VALUE if none.
HANDLE probeControlInterface(uint16_t periodMs, HANDLE evt) {
  std::vector<HANDLE> handles = openMcuInterfaces();
  HANDLE chosen = INVALID_HANDLE_VALUE;
  uint8_t rateCmd[3] = {'R', (uint8_t)(periodMs >> 8), (uint8_t)(periodMs & 0xFF)};
  for (HANDLE h : handles) {
    if (chosen == INVALID_HANDLE_VALUE && setFeature(h, rateCmd, sizeof(rateCmd))) {
      uint8_t buf[65];
      for (int attempt = 0; attempt < 10; attempt++) {
        int got = readReport(h, evt, buf, sizeof(buf), 500);
        if (got <= 0) break;
        // Unnumbered reports (id 0): the HID class driver prepends the 0x00
        // report id on ReadFile — payload starts at byte 1 (hidapi strips
        // this; the bey-closer-t1 offsets are hidapi-relative).
        const uint8_t* r = buf[0] == 0 && got > 1 ? buf + 1 : buf;
        if (r[0] == '#') continue;  // periodic telemetry, keep looking
        if (r[0] == '$') { chosen = h; }
        break;
      }
    }
    if (chosen != h) CloseHandle(h);
  }
  return chosen;
}

// Query the firmware version string over HID_CODE_FOR_SW_VER ('*'): send the
// feature, then read interrupt reports skipping '#' telemetry until the '*'
// reply (firmware echoes the request code, then a null-terminated ASCII version
// string at byte 1). Empty string on timeout. Telemetry is already streaming by
// the time this runs, so a few '#' reports interleave before the reply.
std::string queryVersion(HANDLE h, HANDLE evt) {
  const uint8_t verCmd[1] = {'*'};
  if (!setFeature(h, verCmd, sizeof(verCmd))) return {};
  uint8_t buf[65];
  for (int attempt = 0; attempt < 12; attempt++) {
    int got = readReport(h, evt, buf, sizeof(buf), 500);
    if (got <= 0) break;
    const uint8_t* r = buf[0] == 0 && got > 1 ? buf + 1 : buf;
    int rlen = r == buf ? got : got - 1;
    if (rlen < 2) continue;
    if (r[0] == '#') continue;  // periodic telemetry, keep looking
    if (r[0] == '*') {
      // ASCII version string at byte 1, null-terminated; bound by the report.
      const char* s = reinterpret_cast<const char*>(r + 1);
      size_t maxLen = (size_t)(rlen - 1);
      size_t n = 0;
      while (n < maxLen && s[n] != '\0') n++;
      return std::string(s, n);
    }
    break;  // some other reply code — not the version
  }
  return {};
}

// Parse "major.minor.patch" and report whether it meets the doze contract
// (>= 0.4.0, ADR-0005). Anything unparseable / empty is treated as incapable.
bool versionHasDoze(const std::string& v) {
  int major = 0, minor = 0;
  if (sscanf(v.c_str(), "%d.%d", &major, &minor) < 2) return false;
  return major > 0 || (major == 0 && minor >= 4);
}

}  // namespace

bool McuProx::start(uint16_t periodMs) {
  if (running_.load()) return true;
  periodMs_ = periodMs;
  HANDLE evt = CreateEventW(nullptr, TRUE, FALSE, nullptr);
  HANDLE chosen = probeControlInterface(periodMs, evt);
  if (chosen == INVALID_HANDLE_VALUE) {
    CloseHandle(evt);
    fprintf(stderr,
            "McuProx: no 35BD:0101 interface ACKed the rate command "
            "(Beyond unplugged?)\n");
    return false;
  }

  // Capability probe before the read thread takes the handle: single-threaded
  // here, so the synchronous version query is safe (the doze gate, ADR-0005).
  fwVersion_ = queryVersion(chosen, evt);
  dozeCapable_.store(versionHasDoze(fwVersion_), std::memory_order_relaxed);

  dev_ = chosen;
  readEvent_ = evt;
  stopRequested_.store(false);
  running_.store(true);
  thread_ = std::thread(&McuProx::readLoop, this);
  printf("McuProx: telemetry at %u ms, threshold %u +/- %u; firmware %s "
         "(doze %s)\n",
         periodMs, threshold_, hysteresis_,
         fwVersion_.empty() ? "unknown" : fwVersion_.c_str(),
         dozeCapable_.load() ? "capable" : "unavailable -> parked fallback");
  return true;
}

void McuProx::stop() {
  if (!running_.load() && !thread_.joinable()) return;
  stopRequested_.store(true);
  if (dev_) CancelIoEx(dev_, nullptr);
  if (thread_.joinable()) thread_.join();
  if (dev_) { CloseHandle(dev_); dev_ = nullptr; }
  if (readEvent_) { CloseHandle(readEvent_); readEvent_ = nullptr; }
  running_.store(false);
}

void McuProx::readLoop() {
  // 8-sample moving average + hysteresis, per the bey-closer-t1 port.
  uint16_t ring[8] = {0};
  int idx = 0;
  uint64_t n = 0;
  int consecTimeouts = 0;

  while (!stopRequested_.load()) {
    // SteamVR coexistence: release the 0101 handle so SteamVR's BeyondProximity
    // driver gets exclusive MCU access (concurrent opens wedge the app). Park
    // here, holding nothing, until resume() clears the flag — then fall through
    // to the reconnect path below, which reopens + re-probes once the MCU is
    // free again (BeyondProximity may take a beat to let go on compositor exit).
    if (suspended_.load()) {
      if (dev_) {
        CloseHandle(dev_);
        dev_ = nullptr;
        worn_.store(false, std::memory_order_relaxed);
        prox_.store(0, std::memory_order_relaxed);
        fprintf(stderr, "McuProx: suspended — released 0101 to SteamVR's "
                        "BeyondProximity driver\n");
      }
      Sleep(100);
      continue;
    }
    // Hot-unplug recovery: reopen + re-probe (the rate command must be
    // resent after the MCU reboots) when the handle dies.
    if (!dev_) {
      for (int i = 0; i < 20 && !stopRequested_.load(); i++) Sleep(100);
      if (stopRequested_.load()) break;
      HANDLE h = probeControlInterface(periodMs_, readEvent_);
      if (h == INVALID_HANDLE_VALUE) continue;
      dev_ = h;
      idx = 0;
      n = 0;  // refill the average window before judging worn state
      consecTimeouts = 0;
      fprintf(stderr, "McuProx: device reconnected, telemetry restarted\n");
    }

    uint8_t buf[65];
    int got = readReport(dev_, readEvent_, buf, sizeof(buf), 1000);
    if (got < 0) {
      if (stopRequested_.load()) break;
      fprintf(stderr, "McuProx: read failed (err %lu) — device gone, "
              "waiting for reconnect\n", GetLastError());
      CloseHandle(dev_);
      dev_ = nullptr;
      worn_.store(false, std::memory_order_relaxed);  // unplugged = not worn
      prox_.store(0, std::memory_order_relaxed);
      continue;
    }
    if (got == 0) {
      if (stopRequested_.load()) break;
      // Silence watchdog: telemetry is commanded at periodMs_; a silently-
      // dead handle after a cable pull shows as timeouts, not read errors.
      if (++consecTimeouts >= 3) {
        fprintf(stderr, "McuProx: telemetry silent %ds — device gone, "
                "waiting for reconnect\n", consecTimeouts);
        CloseHandle(dev_);
        dev_ = nullptr;
        worn_.store(false, std::memory_order_relaxed);
        prox_.store(0, std::memory_order_relaxed);
      }
      continue;
    }
    consecTimeouts = 0;
    const uint8_t* r = buf[0] == 0 && got > 1 ? buf + 1 : buf;
    int rlen = r == buf ? got : got - 1;
    if (r[0] != '#' || rlen < 6) continue;

    uint16_t prox = ((uint16_t)r[4] << 8) | r[5];  // big-endian
    prox_.store(prox, std::memory_order_relaxed);
    // Display status word (firmware video_get_display_status, bytes 24-25 BE):
    // bit0 DISPLAYS_ON, bit1 PROX_ON, bit2 PROX_TIMER_EXPIRED, bit3 DSC,
    // high nibble byte0 = video mode, byte1 = link rate / lane count.
    if (rlen >= 26)
      displayStatus_.store(((uint16_t)r[24] << 8) | r[25],
                           std::memory_order_relaxed);
    reports_.fetch_add(1, std::memory_order_relaxed);

    ring[idx] = prox;
    idx = (idx + 1) % 8;
    n++;
    if (n < 8) continue;  // let the window fill before judging
    uint32_t sum = 0;
    for (int i = 0; i < 8; i++) sum += ring[i];
    uint32_t avg = sum / 8;

    bool cur = worn_.load(std::memory_order_relaxed);
    if (cur) {
      if (threshold_ >= hysteresis_ && avg <= (uint32_t)(threshold_ - hysteresis_))
        worn_.store(false, std::memory_order_relaxed);
    } else {
      if (avg >= (uint32_t)threshold_ + hysteresis_)
        worn_.store(true, std::memory_order_relaxed);
    }
  }
  running_.store(false);
}

}  // namespace sauna