MCMapper-Backend/api/src/ws-gateway.ts

import { eq } from "drizzle-orm";
import { db } from "./db/client";
import { chunkColumns, chunkSections, servers } from "./db/schema";
import { markChunkDirty } from "./redis";
import { storeLinkCode } from "./link";
import { recordAndPublishChat } from "./chat";
import { arePlayerPositionsVisible, publishPlayerPositions, type PlayerPosition } from "./players";
import { storeBlockRegistry, storeBlockTextures, type BlockRegistryEntry, type BlockTextureEntry } from "./textures";

// Wire protocol (mod <-> api), one JSON object per WS text frame:
//
//   mod -> api  {"type":"hello","token":"<serverToken>"}
//   api -> mod  {"type":"hello_ack","ok":true,"serverId":"<uuid>"}
//               {"type":"hello_ack","ok":false,"error":"..."}  (connection closed after)
//
//   mod -> api  {"type":"columns","dimension":0,"columns":[{"x":..,"z":..,"height":..,"blockId":..,"blockMeta":..}]}
//   mod -> api  {"type":"sections","dimension":0,"chunkX":..,"chunkZ":..,"sections":[{"sectionY":4,"blocks":"<base64>"}]}
//
// `columns` doubles as both initial backfill (one message per loaded chunk, 256 columns) and
// live deltas (one message per flush tick, just the columns that changed) — both are just "here
// is the current topmost block + height for these XZ columns", the mod recomputes it from its
// own world access rather than the api trying to infer a post-break top block from a raw diff.
//
// `sections` is the Phase 2 addition for full-voxel 3D meshing, additive to `columns` (see
// db/schema.ts's chunkColumns/chunkSections comments) — one message per loaded chunk at load
// time (all non-empty 16x16x16 sections), and again at flush time for chunks touched since the
// last flush (the whole section is resent, same "current state, not a diff" philosophy as
// columns — see DeltaEvent's javadoc on the mod side). `blocks` is 4096 little-endian u16
// blockStateIds, base64-encoded, indexed by `(ly*16 + lz)*16 + lx` within the section.
// A "sections" message marks the chunk dirty the same way "columns" does — one dirty-chunk
// event now triggers the worker to re-render both the 2D tile and any 3D meshes for that chunk.
//
//   mod -> api  {"type":"link_request","code":"AB3XQ9","uuid":"...","username":"...","authMode":"online"}
//
// Phase 3: sent when a player runs `/mcmapper link` — see LinkRequest.java on the mod side and
// link.ts's storeLinkCode/redeemLinkCode for the rest of the flow (redemption happens over plain
// HTTP from the browser, not over this WS connection — see index.ts's /api/link/redeem).
//
//   mod -> api  {"type":"chat","uuid":"...","username":"...","message":"..."}
//   api -> mod  {"type":"chat","username":"...","message":"..."}
//
// In-game chat, both directions. mod -> api is this module (records + publishes for web
// viewers, see chat.ts); api -> mod is sent by the browser-facing /ws/chat route (chat-gateway.ts)
// looking up the mod's connection via getModSocket() and calling ChatBridge.injectWebChatMessage
// on the mod side.
//
//   api -> mod  {"type":"waypoint_share","name":"...","x":..,"y":..,"z":..,"dimension":..,"color":"#RRGGBB","format":"journeymap"|"xaero"}
//
// Phase 4: sent by markers.ts's shareMarkerToChat() via getModSocket() when a linked account
// shares a placed marker to chat. The api only resolves *which* format the target server wants
// (per-server `waypointFormat` config) and forwards the structured point — the mod owns building
// the actual chat text in that format (see WaypointShare.java's javadoc for the researched wire
// formats and their attribution).
//
//   mod -> api  {"type":"player_positions","dimension":0,"players":[{"uuid":"...","username":"...","x":..,"y":..,"z":..}]}
//
// Phase 7b: the mod's throttled online-player roster, always the full current list (not a diff —
// see PlayerPosition.java's javadoc). Not persisted (no meaningful history) — just fanned out via
// players.ts's publishPlayerPositions to whatever browsers are subscribed on /ws/players/:serverId
// (players-gateway.ts), gated on the per-server `playerPositionsVisible` admin toggle (independent
// of the mod-local `playerTrackingEnabled` config that decides whether this message is sent at all).
//
//   mod -> api  {"type":"block_registry","entries":[{"id":4000,"name":"botania:manapool"}]}
//   mod -> api  {"type":"block_textures","textures":[{"name":"botania:manapool","dataBase64":"..."}]}
//
// Phase 11: sent once after `hello_ack` (a world's numeric block-id assignments and mod-jar
// contents are both stable for the server's lifetime, so there's no need to resend on a timer).
// `block_registry` is the numeric-id -> registry-name mapping needed to make sense of
// `chunkColumns.blockId`/`blockMeta` for modded blocks (see MCMapperMod.java's connect-time
// registry dump); `block_textures` is the mod's best-effort classloader extraction of each
// block's texture PNG (see textures.ts's doc comment for storage). Both are stored now but not
// yet read by the render pipeline — see worker/README's Phase 11 note on the deferred per-server
// palette-resolution work this unlocks.

interface Column {
  x: number;
  z: number;
  height: number;
  blockId: number;
  blockMeta: number;
}

interface Section {
  sectionY: number;
  blocks: string;
}

interface ConnState {
  serverId: string;
}

const connections = new Map<string | number, ConnState>();
// Reverse lookup for chat-gateway.ts to forward web-originated chat into the right mod
// connection — populated on a successful hello, cleaned up on close (see wsGateway.close()).
const modSocketsByServer = new Map<string, any>();

export function chunkOf(coord: number): number {
  return Math.floor(coord / 16);
}

export function getModSocket(serverId: string) {
  return modSocketsByServer.get(serverId);
}

export const wsGateway = {
  async open() {
    // Nothing to do until `hello` arrives — see message().
  },

  async message(ws: any, raw: unknown) {
    let msg: any;
    try {
      msg = typeof raw === "string" ? JSON.parse(raw) : raw;
    } catch {
      ws.send(JSON.stringify({ type: "error", error: "invalid_json" }));
      return;
    }

    const state = connections.get(ws.id);

    if (msg.type === "hello") {
      const [row] = await db.select().from(servers).where(eq(servers.token, msg.token)).limit(1);
      if (!row) {
        ws.send(JSON.stringify({ type: "hello_ack", ok: false, error: "invalid_token" }));
        ws.close();
        return;
      }
      connections.set(ws.id, { serverId: row.id });
      modSocketsByServer.set(row.id, ws);
      ws.send(JSON.stringify({ type: "hello_ack", ok: true, serverId: row.id }));
      console.log(`[ws] server '${row.name}' (${row.id}) authenticated`);
      return;
    }

    if (!state) {
      ws.send(JSON.stringify({ type: "error", error: "not_authenticated" }));
      ws.close();
      return;
    }

    if (msg.type === "columns") {
      const dimension: number = msg.dimension;
      const columns: Column[] = msg.columns ?? [];
      if (columns.length === 0) return;

      await db
        .insert(chunkColumns)
        .values(
          columns.map((c) => ({
            serverId: state.serverId,
            dimension,
            x: c.x,
            z: c.z,
            blockId: c.blockId,
            blockMeta: c.blockMeta,
            height: c.height,
          })),
        )
        .onConflictDoUpdate({
          target: [chunkColumns.serverId, chunkColumns.dimension, chunkColumns.x, chunkColumns.z],
          set: {
            blockId: sqlExcluded("block_id"),
            blockMeta: sqlExcluded("block_meta"),
            height: sqlExcluded("height"),
            updatedAt: new Date(),
          },
        });

      const dirtyChunks = new Set<string>();
      for (const c of columns) dirtyChunks.add(`${chunkOf(c.x)},${chunkOf(c.z)}`);
      for (const key of dirtyChunks) {
        const [cx, cz] = key.split(",").map(Number);
        await markChunkDirty(state.serverId, dimension, cx, cz);
      }
      return;
    }

    if (msg.type === "sections") {
      const dimension: number = msg.dimension;
      const chunkX: number = msg.chunkX;
      const chunkZ: number = msg.chunkZ;
      const sections: Section[] = msg.sections ?? [];
      if (sections.length === 0) return;

      await db
        .insert(chunkSections)
        .values(
          sections.map((s) => ({
            serverId: state.serverId,
            dimension,
            x: chunkX,
            z: chunkZ,
            sectionY: s.sectionY,
            blocks: s.blocks,
          })),
        )
        .onConflictDoUpdate({
          target: [
            chunkSections.serverId,
            chunkSections.dimension,
            chunkSections.x,
            chunkSections.z,
            chunkSections.sectionY,
          ],
          set: { blocks: sqlExcluded("blocks"), updatedAt: new Date() },
        });

      await markChunkDirty(state.serverId, dimension, chunkX, chunkZ);
      return;
    }

    if (msg.type === "link_request") {
      // The mod generates `code` itself (so it can show it to the player immediately) — this
      // just remembers what it means until redeemed. See link.ts's storeLinkCode doc comment.
      await storeLinkCode(msg.code, state.serverId, msg.uuid, msg.username, msg.authMode);
      return;
    }

    if (msg.type === "chat") {
      await recordAndPublishChat(state.serverId, {
        username: msg.username,
        message: msg.message,
        source: "game",
      });
      return;
    }

    if (msg.type === "player_positions") {
      if (!(await arePlayerPositionsVisible(state.serverId))) return;
      const dimension: number = msg.dimension;
      const players: PlayerPosition[] = msg.players ?? [];
      await publishPlayerPositions(state.serverId, dimension, players);
      return;
    }

    if (msg.type === "block_registry") {
      const entries: BlockRegistryEntry[] = msg.entries ?? [];
      await storeBlockRegistry(state.serverId, entries);
      return;
    }

    if (msg.type === "block_textures") {
      const textures: BlockTextureEntry[] = msg.textures ?? [];
      await storeBlockTextures(state.serverId, textures);
      return;
    }
  },

  close(ws: any) {
    const state = connections.get(ws.id);
    if (state && modSocketsByServer.get(state.serverId) === ws) {
      modSocketsByServer.delete(state.serverId);
    }
    connections.delete(ws.id);
  },
};

// Small helper: drizzle's onConflictDoUpdate `set` needs a raw `excluded.<col>` reference for
// "use the value that would have been inserted" — drizzle-orm doesn't expose a typed helper for
// this on postgres-js yet, so build the sql fragment directly.
import { sql } from "drizzle-orm";
function sqlExcluded(column: string) {
  return sql.raw(`excluded.${column}`);
}