#include "device/tundra_imu.h"

#include <windows.h>

#include <initguid.h>

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

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

namespace sauna {
namespace {

constexpr USHORT kVid = 0x28DE;
constexpr USHORT kPid = 0x2300;

// The IMU stream lives on MI_00; the interface number is encoded in the
// device interface path ("&mi_00").
bool pathIsMi00(const wchar_t* path) {
  std::wstring p(path);
  for (auto& c : p) c = towlower(c);
  return p.find(L"&mi_00") != std::wstring::npos;
}

// First location path of a devnode, e.g.
// "PCIROOT(0)#PCI(0114,0)#USBROOT(0)#USB(1)#USB(4)#USB(2)".
std::wstring locationPath(DEVINST inst) {
  DEVPROPTYPE type = 0;
  wchar_t buf[512];
  ULONG size = sizeof(buf);
  if (CM_Get_DevNode_PropertyW(inst, &DEVPKEY_Device_LocationPaths, &type,
                               (PBYTE)buf, &size, 0) != CR_SUCCESS)
    return L"";
  return buf;  // first string of the multi-sz
}

// Ancestor hub prefixes: the path with 1 and 2 trailing "#USB(n)" hops
// stripped. The Beyond's Tundra and MCU sit on sibling internal hubs
// (observed: ...#USB(1)#USB(3)#USB(3) vs ...#USB(1)#USB(1)#USB(3)), so they
// share a grandparent, not a parent — pair on any ancestor intersection.
std::vector<std::wstring> ancestorPrefixes(const std::wstring& path) {
  std::vector<std::wstring> out;
  std::wstring p = path;
  for (int i = 0; i < 2; i++) {
    size_t pos = p.rfind(L"#USB(");
    if (pos == std::wstring::npos) break;
    p = p.substr(0, pos);
    out.push_back(p);
  }
  return out;
}

bool sharesAncestor(const std::wstring& a, const std::wstring& b) {
  if (a.empty() || b.empty()) return false;
  for (auto& pa : ancestorPrefixes(a))
    for (auto& pb : ancestorPrefixes(b))
      if (pa == pb) return true;
  return false;
}

// Walk a HID devnode up to its USB composite root (instance id
// "USB\VID_xxxx&PID_xxxx\serial", no "&MI_").
bool compositeRoot(DEVINST inst, const wchar_t* vidpid, DEVINST* out) {
  for (int depth = 0; depth < 6; depth++) {
    DEVINST parent;
    if (CM_Get_Parent(&parent, inst, 0) != CR_SUCCESS) return false;
    wchar_t id[MAX_DEVICE_ID_LEN];
    if (CM_Get_Device_IDW(parent, id, MAX_DEVICE_ID_LEN, 0) != CR_SUCCESS)
      return false;
    std::wstring s(id);
    for (auto& c : s) c = towupper(c);
    if (s.find(vidpid) == 0 && s.find(L"&MI_") == std::wstring::npos) {
      *out = parent;
      return true;
    }
    inst = parent;
  }
  return false;
}

// Location paths of every attached Beyond MCU (35BD:0101 composite root).
std::vector<std::wstring> mcuLocationPaths() {
  std::vector<std::wstring> hubs;
  HDEVINFO info = SetupDiGetClassDevsW(nullptr, L"USB\\VID_35BD&PID_0101",
                                       nullptr,
                                       DIGCF_PRESENT | DIGCF_ALLCLASSES);
  if (info == INVALID_HANDLE_VALUE) return hubs;
  SP_DEVINFO_DATA dd{};
  dd.cbSize = sizeof(dd);
  for (DWORD i = 0; SetupDiEnumDeviceInfo(info, i, &dd); i++) {
    wchar_t id[MAX_DEVICE_ID_LEN];
    if (CM_Get_Device_IDW(dd.DevInst, id, MAX_DEVICE_ID_LEN, 0) != CR_SUCCESS)
      continue;
    std::wstring s(id);
    for (auto& c : s) c = towupper(c);
    if (s.find(L"&MI_") != std::wstring::npos) continue;  // composite root only
    std::wstring lp = locationPath(dd.DevInst);
    if (!lp.empty()) hubs.push_back(lp);
  }
  SetupDiDestroyDeviceInfoList(info);
  return hubs;
}

// Any Watchman device exposes 28DE:2300 MI_00 "IMU" — a wired Index
// controller or Tundra tracker on the same host is indistinguishable from
// the Beyond by VID/PID alone. Disambiguation: the Beyond's Tundra shares
// the headset's internal USB hub with the MCU (35BD:0101), so pair them by
// USB location-path prefix (same trick as the firmware updater). An explicit
// LHR serial overrides. outSerial receives the chosen unit's serial so
// reconnects can be locked to the same unit.
HANDLE openImuInterface(const std::wstring& wantSerial,
                        std::wstring* outSerial = nullptr,
                        bool verbose = true) {
  GUID hidGuid;
  HidD_GetHidGuid(&hidGuid);
  HDEVINFO info = SetupDiGetClassDevsW(&hidGuid, nullptr, nullptr,
                                       DIGCF_PRESENT | DIGCF_DEVICEINTERFACE);
  if (info == INVALID_HANDLE_VALUE) return INVALID_HANDLE_VALUE;

  std::vector<std::wstring> mcuPaths = mcuLocationPaths();

  struct Candidate { HANDLE h; std::wstring serial; bool pairedWithMcu; };
  std::vector<Candidate> candidates;
  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);
    SP_DEVINFO_DATA devinfo{};
    devinfo.cbSize = sizeof(devinfo);
    if (!SetupDiGetDeviceInterfaceDetailW(info, &iface, detail, need, nullptr,
                                          &devinfo))
      continue;
    if (!pathIsMi00(detail->DevicePath)) 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) {
      wchar_t serial[128] = L"";
      HidD_GetSerialNumberString(h, serial, sizeof(serial));
      bool paired = false;
      DEVINST root;
      if (compositeRoot(devinfo.DevInst, L"USB\\VID_28DE&PID_2300\\", &root)) {
        std::wstring lp = locationPath(root);
        for (auto& m : mcuPaths)
          if (sharesAncestor(lp, m)) paired = true;
      }
      candidates.push_back({h, serial, paired});
    } else {
      CloseHandle(h);
    }
  }
  SetupDiDestroyDeviceInfoList(info);

  // Selection: explicit serial wins; otherwise Beyond presence rides on the
  // MCU — only an MCU-paired candidate is acceptable. NO bare single-device
  // fallback: with the headset unplugged a desk Index controller becomes the
  // sole candidate and would be latched onto (field bug).
  size_t pairedCount = 0;
  for (auto& c : candidates) pairedCount += c.pairedWithMcu ? 1 : 0;
  if (verbose)
    printf("TundraImu: %zu Watchman device(s), %zu MCU-paired\n",
           candidates.size(), pairedCount);
  HANDLE found = INVALID_HANDLE_VALUE;
  for (auto& c : candidates) {
    bool match;
    if (!wantSerial.empty())
      match = c.serial == wantSerial;
    else
      match = pairedCount == 1 && c.pairedWithMcu;
    if (match && found == INVALID_HANDLE_VALUE) {
      found = c.h;
      if (outSerial) *outSerial = c.serial;
    } else {
      CloseHandle(c.h);
    }
  }
  if (verbose && found == INVALID_HANDLE_VALUE && !candidates.empty()) {
    fprintf(stderr,
            "TundraImu: cannot pick among %zu Watchman device(s) "
            "(%zu MCU-paired%s) — is the Beyond attached? Or pass its LHR "
            "serial:\n",
            candidates.size(), pairedCount,
            mcuPaths.empty() ? ", no Beyond MCU on the bus" : "");
    for (auto& c : candidates)
      fprintf(stderr, "  %ls%s\n", c.serial.c_str(),
              c.pairedWithMcu ? "  [shares hub with Beyond MCU]" : "");
  }
  return found;
}

// LHR serials are ASCII; lossless for the values that ever appear here.
std::string narrowSerial(const std::wstring& w) {
  std::string s;
  for (wchar_t c : w) s.push_back(static_cast<char>(c));
  return s;
}

}  // namespace

void* OpenWatchmanHid(const char* serial, std::string* outSerial) {
  std::wstring wserial;
  for (const char* p = serial; *p; p++) wserial.push_back((wchar_t)*p);
  std::wstring chosen;
  HANDLE h = openImuInterface(wserial, &chosen);
  if (h == INVALID_HANDLE_VALUE) return nullptr;
  if (outSerial) *outSerial = narrowSerial(chosen);
  return h;
}

bool TundraImu::start(const char* serial) {
  if (running_.load()) return true;
  std::wstring wserial;
  for (const char* p = serial; *p; p++) wserial.push_back((wchar_t)*p);
  std::wstring chosen;
  HANDLE h = openImuInterface(wserial, &chosen);
  if (h == INVALID_HANDLE_VALUE) {
    fprintf(stderr,
            "TundraImu: no matching 28DE:2300 MI_00 device "
            "(Beyond unplugged? wrong --imu-serial?)\n");
    return false;
  }
  // Reconnects re-run the same selection with the same user filter — the
  // MCU-paired-only rule keeps a desk Index controller from being latched
  // onto (field bug), while still allowing a different Beyond to be swapped
  // in mid-session.
  serial_ = wserial;
  printf("TundraImu: using %ls\n", chosen.c_str());
  {
    std::lock_guard<std::mutex> lk(mu_);
    connectedSerial_ = narrowSerial(chosen);
  }
  dev_ = h;
  readEvent_ = CreateEventW(nullptr, TRUE, FALSE, nullptr);
  stopRequested_.store(false);
  running_.store(true);
  thread_ = std::thread(&TundraImu::readLoop, this);
  return true;
}

void TundraImu::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);
}

bool TundraImu::latest(ImuSample* out, double* ageSec) const {
  std::lock_guard<std::mutex> lk(mu_);
  if (!haveSample_) return false;
  *out = latest_;
  if (ageSec)
    *ageSec = std::chrono::duration<double>(std::chrono::steady_clock::now() -
                                            latestAt_).count();
  return true;
}

ImuStats TundraImu::stats() const {
  std::lock_guard<std::mutex> lk(mu_);
  return stats_;
}

std::string TundraImu::connectedSerial() const {
  std::lock_guard<std::mutex> lk(mu_);
  return connectedSerial_;
}

void TundraImu::readLoop() {
  uint8_t buf[64];
  bool haveLast = false;
  uint32_t lastTimecode = 0;
  uint8_t lastSeq = 0;
  int consecTimeouts = 0;
  std::chrono::steady_clock::time_point prevPub{};
  bool havePub = false;

  while (!stopRequested_.load()) {
    // Hot-unplug recovery: dead handle -> poll for the device coming back
    // (field finding: cable pull mid-run; present path self-recovers on
    // replug, readers must too).
    if (!dev_) {
      for (int i = 0; i < 20 && !stopRequested_.load(); i++) Sleep(100);
      if (stopRequested_.load()) break;
      std::wstring chosen;
      HANDLE h = openImuInterface(serial_, &chosen, /*verbose=*/false);
      if (h == INVALID_HANDLE_VALUE) continue;
      dev_ = h;
      haveLast = false;  // stale dedupe state across the gap
      consecTimeouts = 0;
      {
        std::lock_guard<std::mutex> lk(mu_);
        connectedSerial_ = narrowSerial(chosen);
      }
      fprintf(stderr, "TundraImu: reconnected to %ls\n", chosen.c_str());
    }

    OVERLAPPED ov{};
    ov.hEvent = readEvent_;
    ResetEvent(readEvent_);
    DWORD got = 0;
    if (!ReadFile(dev_, buf, sizeof(buf), &got, &ov)) {
      DWORD err = GetLastError();
      if (err == ERROR_IO_PENDING) {
        DWORD w = WaitForSingleObject(readEvent_, 1000);
        if (w == WAIT_TIMEOUT) {
          CancelIoEx(dev_, &ov);
          GetOverlappedResult(dev_, &ov, &got, TRUE);  // reap the cancel
          {
            std::lock_guard<std::mutex> lk(mu_);
            stats_.read_timeouts++;
          }
          // Silence watchdog: the Tundra is never legitimately quiet (it
          // streams ~994 Hz unconditionally). A silently-dead handle after
          // a cable pull shows up as timeouts, not read errors (field bug).
          if (++consecTimeouts >= 3) {
            fprintf(stderr, "TundraImu: stream silent %ds — device gone, "
                    "waiting for a Beyond\n", consecTimeouts);
            CloseHandle(dev_);
            dev_ = nullptr;
          }
          continue;
        }
        if (!GetOverlappedResult(dev_, &ov, &got, FALSE)) {
          if (stopRequested_.load()) break;
          fprintf(stderr, "TundraImu: read failed (err %lu) — device gone, "
                  "waiting for reconnect\n", GetLastError());
          CloseHandle(dev_);
          dev_ = nullptr;
          continue;
        }
      } else {
        if (stopRequested_.load()) break;
        fprintf(stderr, "TundraImu: ReadFile failed (err %lu) — device gone, "
                "waiting for reconnect\n", err);
        CloseHandle(dev_);
        dev_ = nullptr;
        continue;
      }
    }

    consecTimeouts = 0;
    if (got < 52 || buf[0] != 0x20) {
      std::lock_guard<std::mutex> lk(mu_);
      stats_.other_reports++;
      continue;
    }

    // Parse all 3 window samples, dedupe overlap by (timecode, seq).
    ImuSample parsed[3];
    for (int i = 0; i < 3; i++) {
      const uint8_t* p = buf + 1 + 17 * i;
      memcpy(parsed[i].accel, p, 6);
      memcpy(parsed[i].gyro, p + 6, 6);
      memcpy(&parsed[i].timecode, p + 12, 4);
      parsed[i].seq = p[16];
    }

    auto now = std::chrono::steady_clock::now();
    ImuSample fresh[3];
    int nFresh = 0;
    {
      std::lock_guard<std::mutex> lk(mu_);
      stats_.reports++;
      for (int i = 0; i < 3; i++) {
        const ImuSample& s = parsed[i];
        if (haveLast) {
          uint8_t dseq = (uint8_t)(s.seq - lastSeq);
          if (dseq == 0 && s.timecode == lastTimecode) continue;  // window overlap
          if (dseq >= 128) continue;  // stale (behind last) — overlap variant
          if (dseq != 1) stats_.seq_gaps++;
        }
        haveLast = true;
        lastTimecode = s.timecode;
        lastSeq = s.seq;
        latest_ = s;
        latestAt_ = now;
        haveSample_ = true;
        stats_.samples++;
        fresh[nFresh++] = s;
      }
    }
    if (nFresh > 0) {
      // Delivery-gap watermark (see maxDeliveryGapMs). Spans reconnects on
      // purpose — a reconnect gap is a real serving gap too.
      if (havePub) {
        const uint32_t gap = (uint32_t)std::chrono::duration_cast<
            std::chrono::milliseconds>(now - prevPub).count();
        uint32_t cur = maxGapMs_.load(std::memory_order_relaxed);
        while (gap > cur &&
               !maxGapMs_.compare_exchange_weak(cur, gap,
                                                std::memory_order_relaxed)) {
        }
      }
      prevPub = now;
      havePub = true;
    }
    if (sink_)
      for (int i = 0; i < nFresh; i++) sink_(fresh[i]);
  }
  running_.store(false);
}

}  // namespace sauna