import { Buffer } from "node:buffer";
import type { WSContext } from "hono/ws";
import { error as logError } from "../logger";

const textDecoder = new TextDecoder();

export const MAX_WS_MESSAGE_SIZE = 10 * 1024 * 1024;

export type DecodedWsMessage =
  | { ok: true; data: string; size: number }
  | { ok: false; reason: string; size?: number };

export function decodeWsPayload(data: unknown): DecodedWsMessage {
  if (typeof data === "string") {
    return { ok: true, data, size: Buffer.byteLength(data, "utf8") };
  }
  if (data instanceof ArrayBuffer) {
    if (data.byteLength > MAX_WS_MESSAGE_SIZE) {
      return { ok: false, reason: "message too large", size: data.byteLength };
    }
    return { ok: true, data: textDecoder.decode(data), size: data.byteLength };
  }
  if (data instanceof Uint8Array) {
    if (data.byteLength > MAX_WS_MESSAGE_SIZE) {
      return { ok: false, reason: "message too large", size: data.byteLength };
    }
    return { ok: true, data: textDecoder.decode(data), size: data.byteLength };
  }
  if (typeof SharedArrayBuffer !== "undefined" && data instanceof SharedArrayBuffer) {
    const bytes = new Uint8Array(data);
    if (bytes.byteLength > MAX_WS_MESSAGE_SIZE) {
      return { ok: false, reason: "message too large", size: bytes.byteLength };
    }
    return { ok: true, data: textDecoder.decode(bytes), size: bytes.byteLength };
  }
  return { ok: false, reason: typeof data };
}

export function handleSizedWsPayload(
  ws: WSContext,
  logPrefix: string,
  label: string,
  payload: unknown,
  handleMessage: (data: string) => void,
): boolean {
  const decoded = decodeWsPayload(payload);
  if (!decoded.ok) {
    if (decoded.reason === "message too large" && decoded.size !== undefined) {
      logError(`${logPrefix} Message too large on ${label}: size=${decoded.size} limit=${MAX_WS_MESSAGE_SIZE}`);
      ws.close(1009, "message too large");
      return false;
    }
    logError(`${logPrefix} Unsupported message payload on ${label}: ${decoded.reason}`);
    ws.close(1003, "unsupported message payload");
    return false;
  }
  if (decoded.size > MAX_WS_MESSAGE_SIZE) {
    logError(`${logPrefix} Message too large on ${label}: size=${decoded.size} limit=${MAX_WS_MESSAGE_SIZE}`);
    ws.close(1009, "message too large");
    return false;
  }
  handleMessage(decoded.data);
  return true;
}