/**
 * netcode-revival: relay transport for the game's own netcode.
 * See xnet_transport.h for the design.
 */

#include <rex/system/xnet_transport.h>

#include <algorithm>
#include <chrono>
#include <cstring>
#include <thread>

#include <rex/platform.h>
#include <rex/logging.h>
#include <rex/system/kernel_state.h>
#include <rex/system/xsocket.h>
#include <rex/system/xam/user_profile.h>

#if REX_PLATFORM_WIN32
#include <WinSock2.h>
#include <WS2tcpip.h>
#else
#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
#endif

namespace rex {
namespace system {

// Pin the packed wire sizes to the values the relay expects (the game-side
// proto packs them identically). Catches accidental un-packing or member
// edits before they corrupt the wire.
static_assert(sizeof(XNetRelayHeader) == 7, "relay header must be 7 bytes");
static_assert(sizeof(XNetRelayJoin) == 31, "relay join must be 31 bytes");
static_assert(sizeof(XNetRelayGame) == 1410, "relay game payload must be 1410 bytes");
static_assert(sizeof(XNetRelaySessionEntry) == 25, "session entry must be 25 bytes");
static_assert(sizeof(XNetRelaySessionReply) == 801, "session reply must be 801 bytes");

namespace {

// Virtual address range: 10.80.0.0/16 (host order prefix 0x0A500000).
constexpr uint32_t kVirtualPrefix = 0x0A500000u;

int64_t SocketCreate() {
#if REX_PLATFORM_WIN32
  return int64_t(::socket(AF_INET, SOCK_DGRAM, 0));
#else
  return int64_t(::socket(AF_INET, SOCK_DGRAM, 0));
#endif
}

void SocketClose(int64_t s) {
#if REX_PLATFORM_WIN32
  ::closesocket(SOCKET(s));
#else
  ::close(int(s));
#endif
}

// The receive loop must wake regularly so the heartbeat fires even with no
// inbound traffic (otherwise the relay sweeps a silent transport client
// after ~8s) and so StopLocked()'s join completes promptly instead of
// deadlocking SetRelay on a room change.
void SocketSetRecvTimeout(int64_t s, int timeout_ms) {
#if REX_PLATFORM_WIN32
  DWORD tv = DWORD(timeout_ms);
  ::setsockopt(SOCKET(s), SOL_SOCKET, SO_RCVTIMEO,
               reinterpret_cast<const char*>(&tv), sizeof(tv));
#else
  timeval tv = {};
  tv.tv_sec = timeout_ms / 1000;
  tv.tv_usec = (timeout_ms % 1000) * 1000;
  ::setsockopt(int(s), SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
#endif
}

bool SocketSendTo(int64_t s, const sockaddr_in* to, const void* buf, int size) {
#if REX_PLATFORM_WIN32
  const int n = ::sendto(SOCKET(s), reinterpret_cast<const char*>(buf), size, 0,
                         reinterpret_cast<const sockaddr*>(to), sizeof(*to));
#else
  const int n = ::sendto(int(s), buf, size, 0,
                         reinterpret_cast<const sockaddr*>(to), sizeof(*to));
#endif
  return n == size;
}

// Resolve "host" to a sockaddr_in (UDP), defaulting the port. Returns true
// on success.
bool Resolve(const std::string& host, uint16_t port, sockaddr_in* out) {
  addrinfo hints = {};
  hints.ai_family = AF_INET;
  hints.ai_socktype = SOCK_DGRAM;
  hints.ai_protocol = IPPROTO_UDP;
  addrinfo* res = nullptr;
  const std::string port_str = std::to_string(port);
  if (getaddrinfo(host.c_str(), port_str.c_str(), &hints, &res) != 0 ||
      res == nullptr) {
    if (res != nullptr) {
      freeaddrinfo(res);
    }
    return false;
  }
  *out = *reinterpret_cast<sockaddr_in*>(res->ai_addr);
  out->sin_port = htons(port);
  freeaddrinfo(res);
  return true;
}

}  // namespace

XNetTransport& XNetTransport::Instance() {
  static XNetTransport s_instance;
  return s_instance;
}

XNetTransport::~XNetTransport() {
  StopThread();  // join outside the lock (see StopThread)
  std::lock_guard<std::mutex> lock(mu_);
  if (sock_ != -1) {
    SocketClose(sock_);
    sock_ = -1;
  }
  active_ = false;
  joined_ = false;
  sessions_.clear();
  sockets_.clear();
}

uint32_t XNetTransport::LocalVirtualIp() {
  // Mirrors the XNet identity: machine id from the signed-in profile XUID,
  // mapped onto 10.80.0.x. Stable across runs on one box.
  uint64_t seed = 0;
  if (xam::UserProfile* profile = REX_KERNEL_STATE()->user_profile()) {
    seed = profile->xuid();
  }
  if (seed == 0) {
    seed = 0x52454B4E45544CULL;
  }
  const uint8_t machine_id = uint8_t((seed ^ (seed >> 32)) & 0xFEu) + 1;
  return kVirtualPrefix | machine_id;
}

bool XNetTransport::IsVirtualIp(uint32_t host_order_ip) {
  return (host_order_ip & 0xFFFF0000u) == kVirtualPrefix;
}

bool XNetTransport::Active() const {
  std::lock_guard<std::mutex> lock(mu_);
  return active_;
}

void XNetTransport::ForceLiveFlag() {
  // The Skate 3 frontend gates "Xbox Live features" on an internal static
  // byte (the game's own "live user" flag) and refuses WITHOUT calling any
  // rexglue API when it is 0. The game sets it from its sign-in manager, but
  // with no real Live the flag never becomes non-zero. Force it here.
  // Address from the game: sub_8263E788 does lis r11,-32003 -> r11 =
  // 0x83030000 (recompiler's -2097348608) then lbz at -30639 =>
  // 0x82FC8851. (Fragile: tied to this title's build; revival only.)
  uint8_t* flag = REX_KERNEL_MEMORY()->TranslateVirtual<uint8_t*>(0x82FC8851);
  // Rate-limited status so we can SEE whether the address maps and its value.
  static double last_log = 0.0;
  const double now =
      double(std::chrono::steady_clock::now().time_since_epoch().count()) / 1e9;
  const bool rate = now - last_log >= 2.0;
  if (rate) {
    last_log = now;
    REXLOG_INFO("xnet trace: live flag 0x82FC8851 {} value={}",
                flag != nullptr ? "mapped" : "NOT MAPPED",
                flag != nullptr ? int(*flag) : -1);
  }
  // The flag is the LIVE USER INDEX (not a boolean). The profile lives in
  // slot 0, so force 0 (any slot answers after the shim fix, but 0 is the
  // safest for index-into-live-array readers).
  if (flag != nullptr && *flag != 0) {
    if (!rate) {
      REXLOG_INFO("xnet trace: forcing game live flag 0x82FC8851 = 0");
    }
    *flag = 0;
  }
  // The game also caches "live user exists" flags from its sign-in manager
  // (set only when it processes a real signin notification, which never
  // happens with no Live): 0x8308396C (=14700, 'live user exists') and
  // 0x83084690 (=18064, checked by the frontend online-access controllers
  // sub_826147C8 / sub_82614F38). Force both so Party Play ungreys.
  const auto force_byte = [rate](uint32_t addr, const char* what) {
    uint8_t* p = REX_KERNEL_MEMORY()->TranslateVirtual<uint8_t*>(addr);
    if (p != nullptr && *p == 0) {
      if (!rate) {
        REXLOG_INFO("xnet trace: forcing game live cache {} = 1", what);
      }
      *p = 1;
    }
  };
  force_byte(0x8308396C, "0x8308396C (live user exists)");
  force_byte(0x83084690, "0x83084690 (online access)");
}

XNetTransport::GateStatus XNetTransport::GetGateStatus() const {
  GateStatus gs;
  uint8_t* flag = REX_KERNEL_MEMORY()->TranslateVirtual<uint8_t*>(0x82FC8851);
  gs.live_flag_mapped = flag != nullptr;
  gs.live_flag_value = flag != nullptr ? *flag : 0;
  const uint32_t state = REX_KERNEL_STATE()->user_profile()->signin_state();
  const bool signed_in = REX_KERNEL_STATE()->user_profile()->is_signed_in();
  for (int slot = 0; slot < 4; ++slot) {
    gs.signin_state[slot] = state;
    gs.signin_info_ok[slot] = signed_in;
  }
  gs.link_active = true;   // XNetGetEthernetLinkStatus now reports active
  gs.title_online = true;  // XNetGetTitleXnAddr now reports 0x60
  return gs;
}

void XNetTransport::SetRelay(const char* host, uint16_t port, const char* room) {
  {
    std::lock_guard<std::mutex> lock(mu_);
    const std::string h = host ? host : "";
    const std::string r = room ? room : "";
    if (h == host_ && port == port_ && r == room_ && active_) {
      return;  // unchanged and already connected
    }
    host_ = h;
    port_ = port;
    room_ = r;
  }
  StopThread();  // join the old thread OUTSIDE the lock
  Reconnect();
}

void XNetTransport::StopThread() {
  std::thread old;
  {
    std::lock_guard<std::mutex> lock(mu_);
    running_ = false;
    old = std::move(thread_);
    thread_ = {};
  }
  if (old.joinable()) {
    old.join();
  }
}

void XNetTransport::Reconnect() {
  std::lock_guard<std::mutex> lock(mu_);
  if (host_.empty() || port_ == 0) {
    return;
  }
  if (sock_ != -1) {
    SocketClose(sock_);
    sock_ = -1;
  }
  active_ = false;
  joined_ = false;
  sessions_.clear();
  relay_resolved_ = false;
  sockaddr_in relay_addr = {};
  if (!Resolve(host_, port_, &relay_addr)) {
    REXLOG_WARN("xnet transport: cannot resolve relay '{}:{}'", host_, port_);
    return;
  }
  std::memcpy(relay_addr_, &relay_addr, sizeof(relay_addr));
  relay_resolved_ = true;
  sock_ = SocketCreate();
  if (sock_ == -1) {
    REXLOG_WARN("xnet transport: cannot open relay socket");
    return;
  }
  // Wake the receive loop every 200ms (heartbeat + prompt shutdown).
  SocketSetRecvTimeout(sock_, 200);
  REXLOG_INFO("xnet transport: connecting to relay '{}:{}' room '{}' (vip {}.{}.{}.{})", host_,
              port_, room_.c_str(), (LocalVirtualIp() >> 24) & 0xFF, (LocalVirtualIp() >> 16) & 0xFF,
              (LocalVirtualIp() >> 8) & 0xFF, LocalVirtualIp() & 0xFF);
  // Join the room so the relay routes our game traffic among its members.
  XNetRelayJoin jp = {};
  jp.version = kVersion;
  jp.virtual_ip = LocalVirtualIp();
  const uint32_t room_len =
      std::min<uint32_t>(room_.size(), sizeof(jp.room));
  std::memcpy(jp.room, room_.data(), room_len);
  jp.room_len = uint8_t(room_len);
  const char name[] = "skate3";
  const uint32_t name_len = std::min<uint32_t>(sizeof(name) - 1, sizeof(jp.name));
  std::memcpy(jp.name, name, name_len);
  jp.name_len = uint8_t(name_len);
  uint8_t pkt[7 + sizeof(jp)] = {};
  XNetRelayHeader* h = reinterpret_cast<XNetRelayHeader*>(pkt);
  h->magic = kMagic;
  h->type = kJoin;
  h->payload_size = uint16_t(sizeof(jp));
  std::memcpy(pkt + 7, &jp, sizeof(jp));
  SocketSendTo(sock_, &relay_addr, pkt, sizeof(pkt));

  // Keep the JOIN around for the heartbeat (the relay drops silent clients
  // after ~8s).
  std::memcpy(join_pkt_, pkt, sizeof(pkt));
  join_pkt_size_ = uint16_t(sizeof(pkt));
  last_join_ = 0.0;

  running_ = true;
  thread_ = std::thread(&XNetTransport::Run, this);
}

int XNetTransport::SendGame(const void* buf, uint32_t len, uint32_t dst_virtual_ip_host_order) {
  if (len == 0 || len > sizeof(XNetRelayGame::data)) {
    return -1;
  }
  if (!IsVirtualIp(dst_virtual_ip_host_order) && dst_virtual_ip_host_order != 0) {
    return -1;
  }
  std::lock_guard<std::mutex> lock(mu_);
  if (!active_ || sock_ == -1 || !relay_resolved_) {
    return -1;
  }
  const sockaddr_in* relay = reinterpret_cast<const sockaddr_in*>(relay_addr_);
  XNetRelayGame gp = {};
  gp.dst_virtual_ip = dst_virtual_ip_host_order;
  gp.data_len = uint16_t(len);
  std::memcpy(gp.data, buf, len);
  uint8_t pkt[7 + sizeof(gp)] = {};
  XNetRelayHeader* h = reinterpret_cast<XNetRelayHeader*>(pkt);
  h->magic = kMagic;
  h->type = kGame;
  h->payload_size = uint16_t(sizeof(gp));
  std::memcpy(pkt + 7, &gp, sizeof(gp));
  if (!SocketSendTo(sock_, relay, pkt, sizeof(pkt))) {
    return -1;
  }
  return int(len);
}

void XNetTransport::QuerySessions() {
  std::lock_guard<std::mutex> lock(mu_);
  if (!active_ || sock_ == -1 || !relay_resolved_) {
    return;
  }
  const sockaddr_in* relay = reinterpret_cast<const sockaddr_in*>(relay_addr_);
  uint8_t pkt[7] = {};
  XNetRelayHeader* h = reinterpret_cast<XNetRelayHeader*>(pkt);
  h->magic = kMagic;
  h->type = kSessionQuery;
  h->payload_size = 0;
  SocketSendTo(sock_, relay, pkt, sizeof(pkt));
}

void XNetTransport::WaitForReply(int timeout_ms) {
  const uint64_t target = [this] {
    std::lock_guard<std::mutex> lock(mu_);
    return reply_seq_;
  }();
  const auto deadline = std::chrono::steady_clock::now() +
                        std::chrono::milliseconds(timeout_ms);
  while (std::chrono::steady_clock::now() < deadline) {
    std::lock_guard<std::mutex> lock(mu_);
    if (reply_seq_ > target) {
      return;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(5));
  }
}

std::vector<XNetSession> XNetTransport::Sessions() const {
  std::lock_guard<std::mutex> lock(mu_);
  return sessions_;
}

void XNetTransport::RegisterSocket(XSocket* s) {
  if (s == nullptr) {
    return;
  }
  std::lock_guard<std::mutex> lock(mu_);
  sockets_.insert(s);
}

void XNetTransport::UnregisterSocket(XSocket* s) {
  std::lock_guard<std::mutex> lock(mu_);
  sockets_.erase(s);
}

void XNetTransport::Run() {
  while (true) {
    {
      std::lock_guard<std::mutex> lock(mu_);
      if (!running_) {
        return;
      }
      // Heartbeat: re-JOIN every ~4s so the relay never times us out.
      const double now =
          double(std::chrono::steady_clock::now().time_since_epoch().count()) / 1e9;
      if (join_pkt_size_ != 0 && now - last_join_ >= 4.0) {
        SocketSendTo(sock_, reinterpret_cast<const sockaddr_in*>(relay_addr_), join_pkt_,
                     join_pkt_size_);
        last_join_ = now;
      }
    }
    uint8_t buf[7 + sizeof(XNetRelayGame)];
    sockaddr_in from = {};
    socklen_t from_len = sizeof(from);
#if REX_PLATFORM_WIN32
    int n = ::recvfrom(SOCKET(sock_), reinterpret_cast<char*>(buf), sizeof(buf), 0,
                       reinterpret_cast<sockaddr*>(&from), &from_len);
#else
    int n = ::recvfrom(int(sock_), buf, sizeof(buf), 0,
                       reinterpret_cast<sockaddr*>(&from), &from_len);
#endif
    if (n < 0) {
      std::this_thread::sleep_for(std::chrono::milliseconds(5));
      continue;
    }
    if (n < 7) {
      continue;
    }
    HandlePacket(buf, n);
  }
}

void XNetTransport::HandlePacket(const uint8_t* data, int size) {
  const XNetRelayHeader* h = reinterpret_cast<const XNetRelayHeader*>(data);
  if (h->magic != kMagic || h->payload_size > 3200 ||
      size < int(sizeof(*h) + h->payload_size)) {
    return;
  }
  const uint8_t* payload = data + sizeof(*h);

  switch (h->type) {
    case kHello: {
      std::lock_guard<std::mutex> lock(mu_);
      // session_id != 0 means accepted (0 = refused).
      active_ = payload[0] != 0;
      joined_ = true;
      REXLOG_INFO("xnet transport: relay {} (vip {}.{}.{}.{})", active_ ? "joined" : "refused",
                  (LocalVirtualIp() >> 24) & 0xFF, (LocalVirtualIp() >> 16) & 0xFF,
                  (LocalVirtualIp() >> 8) & 0xFF, LocalVirtualIp() & 0xFF);
      break;
    }
    case kGame: {
      if (h->payload_size < sizeof(XNetRelayGame)) {
        return;
      }
      XNetRelayGame gp = {};
      std::memcpy(&gp, payload, sizeof(gp));
      if (gp.data_len > sizeof(gp.data)) {
        return;
      }
      std::vector<XSocket*> targets;
      {
        std::lock_guard<std::mutex> lock(mu_);
        for (XSocket* s : sockets_) {
          targets.push_back(s);
        }
      }
      for (XSocket* s : targets) {
        // src_virtual_ip is host order; QueuePacket's src is host order too.
        s->QueuePacket(gp.src_virtual_ip, 0, gp.data, gp.data_len);
      }
      break;
    }
    case kSessionReply: {
      if (h->payload_size < sizeof(XNetRelaySessionReply)) {
        return;
      }
      XNetRelaySessionReply sr = {};
      std::memcpy(&sr, payload, sizeof(sr));
      std::vector<XNetSession> fresh;
      const uint8_t count = std::min<uint8_t>(sr.count, 32);
      for (uint8_t i = 0; i < count; ++i) {
        XNetSession s;
        s.session_id = sr.list[i].session_id;
        s.virtual_ip = sr.list[i].virtual_ip;
        s.name_len = std::min<uint8_t>(sr.list[i].name_len, sizeof(s.name));
        std::memcpy(s.name, sr.list[i].name, s.name_len);
        fresh.push_back(s);
      }
      std::lock_guard<std::mutex> lock(mu_);
      sessions_ = std::move(fresh);
      ++reply_seq_;
      REXLOG_INFO("xnet transport: session reply with {} session(s) in room '{}'", count,
                  room_.c_str());
      break;
    }
    default:
      break;  // kTypeRooms / Pose / Ping / Pong: not this transport's business.
  }
}

}  // namespace system
}  // namespace rex