#include "skate3_native_ghost_relay.h"

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <string>
#include <unordered_map>
#include <vector>

// The relay core is compiled into BOTH the game (rex logging available) and
// the dependency-free standalone ghost_relay binary; route its logs through
// a small build-flag shim so the standalone stays rexglue-free.
#if defined(SKATE3_RELAY_STANDALONE)
namespace skate3::native_ghost {
namespace {
inline std::string StandaloneArg(uint32_t v) { return std::to_string(v); }
inline std::string StandaloneArg(int v) { return std::to_string(v); }
inline std::string StandaloneArg(const std::string& v) { return v; }
inline std::string StandaloneArg(const char* v) {
  return v != nullptr ? v : "(null)";
}

// Minimal "{}"-placeholder formatter (the relay logs use fmt-style strings
// because the game build routes them through REXLOG).
template <typename... Args>
void StandaloneLog(const char* level, const char* fmt, Args&&... args) {
  const std::vector<std::string> vals = {StandaloneArg(args)...};
  std::string out;
  size_t k = 0;
  const char* p = fmt;
  while (*p != '\0') {
    if (p[0] == '{' && p[1] == '}' && k < vals.size()) {
      out += vals[k++];
      p += 2;
    } else {
      out += *p++;
    }
  }
  if (level[0] == 'W') {
    std::fprintf(stderr, "[ghost_relay] %s: %s\n", level, out.c_str());
  } else {
    std::printf("[ghost_relay] %s: %s\n", level, out.c_str());
  }
}
}  // namespace
}  // namespace skate3::native_ghost

#define GHOST_RELAY_INFO(...) \
  ::skate3::native_ghost::StandaloneLog("I", __VA_ARGS__)
#define GHOST_RELAY_WARN(...) \
  ::skate3::native_ghost::StandaloneLog("W", __VA_ARGS__)
#else
#include <rex/logging.h>
#define GHOST_RELAY_INFO(...) REXLOG_INFO(__VA_ARGS__)
#define GHOST_RELAY_WARN(...) REXLOG_WARN(__VA_ARGS__)
#endif

namespace skate3::native_ghost {

namespace {
constexpr double kClientTimeoutSec = 8.0;
}  // namespace

GhostRelay::~GhostRelay() {
  Stop();
}

bool GhostRelay::Start() {
  if (running_.load(std::memory_order_relaxed)) {
    return true;
  }
  sock_ = GhostSocketCreate(/*allow_broadcast=*/true);
  if (sock_ == GHOST_INVALID_SOCKET) {
    return false;
  }
  if (!GhostSockBind(sock_, port_)) {
    GhostSockClose(sock_);
    sock_ = GHOST_INVALID_SOCKET;
    return false;
  }
  GhostSockSetRecvTimeout(sock_, 200);
  running_.store(true, std::memory_order_relaxed);
  stop_.store(false, std::memory_order_relaxed);
  thread_ = std::thread(&GhostRelay::Run, this);
  return true;
}

void GhostRelay::Stop() {
  if (!running_.load(std::memory_order_relaxed)) {
    return;
  }
  stop_.store(true, std::memory_order_relaxed);
  if (thread_.joinable()) {
    thread_.join();
  }
  if (sock_ != GHOST_INVALID_SOCKET) {
    GhostSockClose(sock_);
    sock_ = GHOST_INVALID_SOCKET;
  }
  running_.store(false, std::memory_order_relaxed);
  std::lock_guard<std::mutex> lock(mu_);
  clients_.clear();
}

size_t GhostRelay::ClientCount() const {
  std::lock_guard<std::mutex> lock(mu_);
  return clients_.size();
}

void GhostRelay::SetLockedRoom(const char* room) {
  std::lock_guard<std::mutex> lock(mu_);
  RoomSet(locked_room_, locked_room_len_, room);
  locked_ = locked_room_len_ != 0;
}

void GhostRelay::Run() {
  while (!stop_.load(std::memory_order_relaxed)) {
    try {
      GhostAddr from;
      PacketBuf p;
      const int n = GhostSockRecvFrom(sock_, &from, &p, sizeof(p));
      if (n <= 0) {
        Sweep(0.0);
        continue;
      }
      if (n < static_cast<int>(sizeof(PacketHeader))) {
        continue;
      }
      if (p.h.magic != kGhostMagic) {
        continue;
      }
      const uint16_t expect = static_cast<uint16_t>(n - sizeof(PacketHeader));
      if (p.h.payload_size != expect ||
          p.h.payload_size > kGhostMaxPayload) {
        continue;
      }
      const double now =
          double(std::chrono::steady_clock::now().time_since_epoch().count()) /
          1e9;
      HandlePacket(from, p, now);
      Sweep(now);
    } catch (const std::exception& e) {
      GHOST_RELAY_WARN("relay thread exception: {}", e.what());
    } catch (...) {
      GHOST_RELAY_WARN("relay thread unknown exception");
    }
  }
}

void GhostRelay::HandlePacket(const GhostAddr& from, const PacketBuf& p,
                              double now) {
  std::lock_guard<std::mutex> lock(mu_);

  // Resolve the sender's client record. Client records are only created by
  // JOIN/PING (a client must announce itself before posing); any other
  // packet from an unknown address is dropped, so spoofed POSE/LEAVE floods
  // cannot fill the client table. LIST is the exception: a lobby browser can
  // ask for the room list before joining anything.
  Client* me = nullptr;
  for (Client& c : clients_) {
    if (GhostAddrEqual(c.addr, from)) {
      me = &c;
      break;
    }
  }
  if (me == nullptr) {
    if (p.h.type == kTypeList) {
      BuildRoomList(from);
      return;
    }
    if (p.h.type != kTypeJoin && p.h.type != kTypePing) {
      return;
    }
    if (clients_.size() >= max_clients_) {
      GHOST_RELAY_WARN("client cap reached ({}); dropping {}", max_clients_,
                  GhostAddrStr(from));
      return;
    }
    clients_.push_back(Client{});
    me = &clients_.back();
    me->addr = from;
    me->session_id = next_session_++;
    RoomClear(me->room, me->room_len);
  }

  const auto send = [this, &from](const PacketBuf& out) {
    GhostSockSendTo(sock_, from, &out, PacketSize(out));
  };
  PacketBuf out = {};

  switch (p.h.type) {
    case kTypeJoin:
    case kTypePing: {
      if (p.h.payload_size < sizeof(JoinPayload)) {
        return;
      }
      JoinPayload jp = {};
      std::memcpy(&jp, p.payload, sizeof(jp));
      // Protocol handshake: a mismatched version is REFUSED (HELLO with
      // session 0) and its poses are dropped, so two builds can never
      // silently corrupt each other's pose stream.
      if (jp.version != kGhostVersion) {
        me->refused = true;
        if (!me->name[0]) {
          std::memcpy(me->name, jp.name, std::min<uint8_t>(jp.name_len, kGhostMaxNameLen));
          me->name_len = std::min<uint8_t>(jp.name_len, kGhostMaxNameLen);
        }
        HelloPayload hp = {};
        hp.session_id = 0;  // refused
        hp.version = kGhostVersion;
        BuildPacket(out, kTypeHello, &hp, static_cast<uint16_t>(sizeof(hp)));
        send(out);
        GHOST_RELAY_WARN("refused {} (client v{}, relay v{})", GhostAddrStr(from),
                    jp.version, kGhostVersion);
        return;
      }
      if (locked_) {
        std::memcpy(me->room, locked_room_, locked_room_len_);
        me->room_len = locked_room_len_;
      } else if (RoomValid(jp.room, jp.room_len)) {
        std::memcpy(me->room, jp.room, jp.room_len);
        me->room_len = jp.room_len;
      }
      // Keep the client's display name (JOIN carries it; PINGs from probes
      // may omit it, so only overwrite with a non-empty name).
      if (jp.name_len != 0 && jp.name_len <= kGhostMaxNameLen) {
        std::memcpy(me->name, jp.name, jp.name_len);
        me->name_len = jp.name_len;
      }
      // v5: the client's virtual IP identity (used to route kTypeGame and to
      // answer kTypeSessionQuery). PING probes may omit it.
      if (jp.virtual_ip != 0) {
        me->virtual_ip = jp.virtual_ip;
      }
      // Empty room: assign one (uppercase alnum, 4 chars, seeded by the
      // session id and the wall of the receiving thread).
      if (me->room_len == 0) {
        uint32_t state = (uint32_t(now * 1e3) * 2654435761u) ^
                         (me->session_id * 40503u) ^ 0x9E3779B9u;
        char room[kGhostMaxRoomLen] = {};
        for (uint8_t i = 0; i < 4; ++i) {
          state = state * 1664525u + 1013904223u;
          room[i] = (state & 1u) != 0
                        ? static_cast<char>('A' + (state >> 1) % 26u)
                        : static_cast<char>('0' + (state >> 2) % 10u);
        }
        std::memcpy(me->room, room, 4);
        me->room_len = 4;
      }
      me->last_seen = now;
      if (p.h.type == kTypeJoin) {
        // HELLO: our session id, room and the roster of the OTHER members.
        HelloPayload hp = {};
        hp.session_id = me->session_id;
        hp.version = kGhostVersion;
        hp.room_len = me->room_len;
        std::memcpy(hp.room, me->room, me->room_len);
        for (const Client& c : clients_) {
          if (&c != me && c.room_len == me->room_len &&
              RoomEquals(c.room, c.room_len, me->room, me->room_len) &&
              hp.roster_count < kGhostMaxRoster) {
            RosterEntry& entry = hp.roster[hp.roster_count++];
            entry.session_id = c.session_id;
            entry.name_len = c.name_len;
            std::memcpy(entry.name, c.name, c.name_len);
          }
        }
        BuildPacket(out, kTypeHello, &hp, static_cast<uint16_t>(sizeof(hp)));
        send(out);
        if (verbose_.load(std::memory_order_relaxed)) {
          GHOST_RELAY_INFO("{} joined room '{}' (session {})", GhostAddrStr(from),
                      std::string(hp.room, hp.room_len), me->session_id);
        }
      } else {
        // PING probe: HELLO with the host's room/roster so probing clients
        // can decide whether to join this LAN host.
        HelloPayload hp = {};
        hp.session_id = me->session_id;
        hp.version = kGhostVersion;
        hp.room_len = me->room_len;
        std::memcpy(hp.room, me->room, me->room_len);
        for (const Client& c : clients_) {
          if (&c != me && c.room_len == me->room_len &&
              RoomEquals(c.room, c.room_len, me->room, me->room_len) &&
              hp.roster_count < kGhostMaxRoster) {
            RosterEntry& entry = hp.roster[hp.roster_count++];
            entry.session_id = c.session_id;
            entry.name_len = c.name_len;
            std::memcpy(entry.name, c.name, c.name_len);
          }
        }
        BuildPacket(out, kTypeHello, &hp, static_cast<uint16_t>(sizeof(hp)));
        send(out);
      }
      break;
    }

    case kTypePose:
    case kTypePoseChunk: {
      // A refused (version-mismatched) client's poses are never forwarded.
      if (me->refused) {
        return;
      }
      // Clients join before posing; a pose with no room is auto-joined to
      // the lobby so the relay stays resilient to a missed HELLO.
      if (me->room_len == 0) {
        RoomSet(me->room, me->room_len, kGhostLobbyRoom);
      }
      me->last_seen = now;
      // Forward the payload VERBATIM to the other room members, re-stamping
      // only the sender session id at payload offset 0 (little-endian). The
      // relay never parses the pose body (a full skeleton pose is a
      // multi-chunk kTypePoseChunk stream; each chunk is forwarded the same
      // way), so the protocol can grow without a relay update.
      PacketBuf out = p;
      out.h.magic = kGhostMagic;
      out.h.type = p.h.type;
      std::memcpy(out.payload, &me->session_id, sizeof(me->session_id));
      const uint16_t size = PacketSize(out);
      for (const Client& c : clients_) {
        if (&c == me) {
          continue;
        }
        if (c.room_len != me->room_len ||
            !RoomEquals(c.room, c.room_len, me->room, me->room_len)) {
          continue;
        }
        GhostSockSendTo(sock_, c.addr, &out, size);
      }
      break;
    }

    case kTypeGame: {
      // netcode-revival: the game's own RakNet UDP, opaque to us. Forward
      // to the destination member (dst_virtual_ip, or every other member
      // when 0), re-stamping the sender's virtual IP so the recipient can
      // reply. Unknown addresses are dropped (a member with a bogus dst is
      // gone or never joined).
      if (me->refused) {
        return;
      }
      if (p.h.payload_size != sizeof(GamePayload) || me->virtual_ip == 0) {
        return;
      }
      GamePayload gp = {};
      std::memcpy(&gp, p.payload, sizeof(gp));
      if (gp.data_len > kGhostGameData) {
        return;
      }
      me->last_seen = now;
      PacketBuf out = p;
      out.h.magic = kGhostMagic;
      out.h.type = kTypeGame;
      GamePayload* fwd = reinterpret_cast<GamePayload*>(out.payload);
      fwd->src_virtual_ip = me->virtual_ip;
      const uint16_t size = PacketSize(out);
      uint32_t recipients = 0;
      for (Client& c : clients_) {
        if (&c == me || c.refused) {
          continue;
        }
        if (c.room_len != me->room_len ||
            !RoomEquals(c.room, c.room_len, me->room, me->room_len)) {
          continue;
        }
        if (gp.dst_virtual_ip != 0 && c.virtual_ip != gp.dst_virtual_ip) {
          continue;
        }
        GhostSockSendTo(sock_, c.addr, &out, size);
        ++recipients;
      }
      // Verbose: a windowed summary of game datagrams relayed, so the game's
      // own netcode traffic is visible in the log.
      if (verbose_.load(std::memory_order_relaxed) && recipients != 0) {
        ++game_forwards_;
        if (now - last_game_log_ >= 2.0 || game_forwards_ == 1) {
          GHOST_RELAY_INFO("game traffic: {} datagram(s) relayed to {} "
                           "member(s) in room '{}'",
                           game_forwards_, recipients,
                           std::string(me->room, me->room_len));
          last_game_log_ = now;
          game_forwards_ = 0;
        }
      }
      break;
    }

    case kTypeSessionQuery: {
      // netcode-revival: a client asks "who's in my room?" so XNetQosLookup
      // can discover sessions. Report every OTHER member with a known
      // virtual IP (their session id + address + name).
      if (me->refused) {
        return;
      }
      me->last_seen = now;
      BuildSessionReply(from, *me);
      break;
    }

    case kTypeLeave: {
      const auto it = std::remove_if(clients_.begin(), clients_.end(),
                                     [&from](const Client& c) {
                                       return GhostAddrEqual(c.addr, from);
                                     });
      if (it != clients_.end() && verbose_.load(std::memory_order_relaxed)) {
        GHOST_RELAY_INFO("{} left", GhostAddrStr(from));
      }
      clients_.erase(it, clients_.end());
      break;
    }

    case kTypeList: {
      BuildRoomList(from);
      break;
    }

    default:
      break;  // PONG etc. are client-side only; ignore at the relay.
  }
}

void GhostRelay::BuildRoomList(const GhostAddr& to) {
  // Aggregate the client table by room (the mu_ lock is already held by the
  // caller), most-populated first, capped at kGhostMaxRooms with up to
  // kGhostMaxRoomMembers names per room.
  struct RoomAgg {
    uint32_t count = 0;
    uint32_t shown = 0;
    uint32_t names[kGhostMaxRoomMembers] = {};  // indices into clients_
  };
  std::unordered_map<std::string, RoomAgg> agg;
  for (size_t i = 0; i < clients_.size(); ++i) {
    const Client& c = clients_[i];
    if (c.room_len == 0 || c.refused) {
      continue;
    }
    const std::string key(c.room, c.room_len);
    RoomAgg& a = agg[key];
    ++a.count;
    if (a.shown < kGhostMaxRoomMembers && c.name_len != 0) {
      a.names[a.shown++] = uint32_t(i);
    }
  }

  RoomsPayload rp = {};
  rp.version = kGhostVersion;
  uint8_t listed = 0;
  for (const auto& kv : agg) {
    if (listed >= kGhostMaxRooms) {
      break;
    }
    const std::string& key = kv.first;
    const RoomAgg& a = kv.second;
    RoomEntry& e = rp.rooms[listed];
    e.room_len = uint8_t(std::min<size_t>(key.size(), kGhostMaxRoomLen));
    std::memcpy(e.room, key.data(), e.room_len);
    e.member_count = uint8_t(std::min<uint32_t>(a.count, 255));
    e.shown_members = uint8_t(a.shown);
    for (uint32_t n = 0; n < a.shown; ++n) {
      const Client& c = clients_[a.names[n]];
      e.member_name_len[n] = c.name_len;
      std::memcpy(e.member_names[n], c.name, c.name_len);
    }
    ++listed;
  }
  rp.room_count = listed;

  PacketBuf out;
  BuildPacket(out, kTypeRooms, &rp, static_cast<uint16_t>(sizeof(rp)));
  GhostSockSendTo(sock_, to, &out, PacketSize(out));
}

void GhostRelay::BuildSessionReply(const GhostAddr& to, const Client& me) {
  // mu_ is held by the caller. Report the OTHER members of the requester's
  // room that have a known virtual IP (their relay session id + 10.80.0.x
  // address + display name) - exactly the discovery data XNetQosLookup needs
  // to hand the game.
  SessionReplyPayload sr = {};
  for (const Client& c : clients_) {
    if (&c == &me || c.refused || c.virtual_ip == 0) {
      continue;
    }
    if (c.room_len != me.room_len || !RoomEquals(c.room, c.room_len, me.room, me.room_len)) {
      continue;
    }
    if (sr.count >= kGhostMaxRoster) {
      break;
    }
    SessionEntry& e = sr.list[sr.count++];
    e.session_id = c.session_id;
    e.virtual_ip = c.virtual_ip;
    e.name_len = c.name_len;
    std::memcpy(e.name, c.name, c.name_len);
  }

  PacketBuf out;
  BuildPacket(out, kTypeSessionReply, &sr, static_cast<uint16_t>(sizeof(sr)));
  GhostSockSendTo(sock_, to, &out, PacketSize(out));
}

void GhostRelay::Sweep(double now) {
  std::lock_guard<std::mutex> lock(mu_);
  if (now <= 0.0) {
    return;
  }
  const bool verbose = verbose_.load(std::memory_order_relaxed);
  clients_.erase(std::remove_if(
                    clients_.begin(), clients_.end(),
                    [now, verbose](const Client& c) {
                      const bool stale = now - c.last_seen > kClientTimeoutSec;
                      if (stale && verbose) {
                        GHOST_RELAY_INFO("{} timed out ({}s silence; {} joined "
                                         "room '{}')",
                                         GhostAddrStr(c.addr),
                                         int(now - c.last_seen), c.session_id,
                                         std::string(c.room, c.room_len));
                      }
                      return stale;
                    }),
                 clients_.end());
}

}  // namespace skate3::native_ghost