WRNexusJS/packages/cli/src/migrations/mode-functions.ts

/**
 * Migration: move mode-scoped helpers out of legacy `ssr { functions { … } }`
 * / `client { functions { … } }` blocks into the page-level `functions { }`
 * block, tagged `shared`.
 *
 * IMPORTANT: this is a text-level transform, not a parse-and-rewrite, for the
 * same reason as `apis-block.ts`: the `@wrnexus/syntax` parser deliberately
 * THROWS a `ParseError` on `ssr { … }` / `client { … }` data blocks — that
 * legacy syntax is exactly what this migration exists to read, so it cannot
 * be reached by parsing first. Instead this scans the source with the same
 * balanced-brace technique used throughout `update.ts` (`findMatching`), and
 * only parses the RESULT, as a validity check.
 */

import { parse } from "@wrnexus/syntax";

/** Balanced-brace scan: returns the index of the brace matching `open`, or -1. */
function findMatching(source: string, open: number, openChar = "{", closeChar = "}"): number {
  let depth = 0;
  let quote = "";
  for (let index = open; index < source.length; index++) {
    const char = source[index]!;
    if (quote) {
      if (char === "\\") index++;
      else if (char === quote) quote = "";
      continue;
    }
    if (char === '"' || char === "'" || char === "`") quote = char;
    else if (char === openChar) depth++;
    else if (char === closeChar && --depth === 0) return index;
  }
  return -1;
}

/** Net `{` minus `}` outside string literals -- a splice that loses a brace changes it. */
function braceBalance(source: string): number {
  let balance = 0;
  let quote = "";
  for (let index = 0; index < source.length; index++) {
    const char = source[index]!;
    if (quote) {
      if (char === "\\") index++;
      else if (char === quote) quote = "";
      continue;
    }
    if (char === '"' || char === "'" || char === "`") quote = char;
    else if (char === "{") balance++;
    else if (char === "}") balance--;
  }
  return balance;
}

interface ModeBlock {
  mode: "ssr" | "client";
  /** Start of the `ssr`/`client` keyword. */
  start: number;
  /** One past the block's closing `}`. */
  end: number;
  bodyStart: number;
  bodyEnd: number;
}

/** Find page-level `ssr { … }` / `client { … }` data blocks (not `state`/hydrate forms). */
function findModeBlocks(source: string): ModeBlock[] {
  const blocks: ModeBlock[] = [];
  const header = /\b(ssr|client)\s*\{/g;
  let match: RegExpExecArray | null;
  while ((match = header.exec(source))) {
    const open = match.index + match[0].length - 1;
    const close = findMatching(source, open);
    if (close < 0) {
      throw new Error(`unterminated "${match[1]} {" block at offset ${match.index}`);
    }
    blocks.push({
      mode: match[1] as "ssr" | "client",
      start: match.index,
      end: close + 1,
      bodyStart: open + 1,
      bodyEnd: close,
    });
    header.lastIndex = close + 1;
  }
  return blocks;
}

interface FunctionsSubBlock {
  /** Offsets relative to the mode block's body. */
  localStart: number;
  localEnd: number;
  bodyStart: number;
  bodyEnd: number;
}

/** Find `functions { … }` sub-blocks inside a mode-block body. */
function findFunctionsSubBlocks(body: string): FunctionsSubBlock[] {
  const blocks: FunctionsSubBlock[] = [];
  const header = /\bfunctions\s*\{/g;
  let match: RegExpExecArray | null;
  while ((match = header.exec(body))) {
    const open = match.index + match[0].length - 1;
    const close = findMatching(body, open);
    if (close < 0) {
      throw new Error(`unterminated "functions {" block at offset ${match.index}`);
    }
    blocks.push({
      localStart: match.index,
      localEnd: close + 1,
      bodyStart: open + 1,
      bodyEnd: close,
    });
    header.lastIndex = close + 1;
  }
  return blocks;
}

interface FunctionEntry {
  name: string;
  /** `function name(...) { ... }` text, with any leading runtime keyword stripped. */
  text: string;
}

/** Find `[client|server|shared] function <name>(...) { ... }` entries inside a functions-block body. */
function findFunctionEntries(body: string): FunctionEntry[] {
  const entries: FunctionEntry[] = [];
  const header = /\b(?:(?:client|server|shared)\s+)?function\s+([A-Za-z_$][\w$]*)\s*\(/g;
  let match: RegExpExecArray | null;
  while ((match = header.exec(body))) {
    const parenOpen = match.index + match[0].length - 1;
    const parenClose = findMatching(body, parenOpen, "(", ")");
    if (parenClose < 0) {
      throw new Error(`unterminated function parameter list at offset ${match.index}`);
    }
    const braceOpen = body.indexOf("{", parenClose + 1);
    if (braceOpen < 0) {
      throw new Error(`unterminated function body at offset ${match.index}`);
    }
    const braceClose = findMatching(body, braceOpen);
    if (braceClose < 0) {
      throw new Error(`unterminated function body at offset ${match.index}`);
    }
    const localStart = match.index;
    const localEnd = braceClose + 1;
    const text = body.slice(localStart, localEnd).replace(/^(client|server|shared)\s+/, "");
    entries.push({ name: match[1]!, text });
    header.lastIndex = localEnd;
  }
  return entries;
}

interface Edit {
  start: number;
  end: number;
  replacement: string;
}

function applyEdits(source: string, edits: Edit[]): string {
  const sorted = [...edits].sort((a, b) => a.start - b.start);
  let output = "";
  let cursor = 0;
  for (const edit of sorted) {
    output += source.slice(cursor, edit.start) + edit.replacement;
    cursor = edit.end;
  }
  output += source.slice(cursor);
  return output;
}

export function migrateModeFunctions(
  source: string,
): { source: string; changed: boolean } | { skip: string } {
  const blocks = findModeBlocks(source);
  if (blocks.length === 0) return { source, changed: false };

  const allEntries: FunctionEntry[] = [];
  const nameMode = new Map<string, "ssr" | "client">();
  const perBlockFuncBlocks = new Map<ModeBlock, FunctionsSubBlock[]>();

  for (const block of blocks) {
    const body = source.slice(block.bodyStart, block.bodyEnd);
    const funcBlocks = findFunctionsSubBlocks(body);
    if (funcBlocks.length === 0) continue;

    for (const funcBlock of funcBlocks) {
      const funcBody = body.slice(funcBlock.bodyStart, funcBlock.bodyEnd);
      for (const entry of findFunctionEntries(funcBody)) {
        const existingMode = nameMode.get(entry.name);
        if (existingMode && existingMode !== block.mode) {
          return {
            skip: `function '${entry.name}' is declared in both ssr and client — merge them by hand before this migration can run`,
          };
        }
        nameMode.set(entry.name, block.mode);
        allEntries.push(entry);
      }
    }
    perBlockFuncBlocks.set(block, funcBlocks);
  }

  if (allEntries.length === 0) return { source, changed: false };

  // Locate an existing page-level `functions { }` block, i.e. one that is not
  // inside any ssr/client mode block.
  const insideAMode = (index: number) =>
    blocks.some((block) => index >= block.start && index < block.end);

  let existingFunctions: { bodyStart: number; bodyEnd: number } | null = null;
  {
    const header = /\bfunctions\s*\{/g;
    let match: RegExpExecArray | null;
    while ((match = header.exec(source))) {
      if (insideAMode(match.index)) continue;
      const open = match.index + match[0].length - 1;
      const close = findMatching(source, open);
      if (close < 0) throw new Error(`unterminated "functions {" block at offset ${match.index}`);
      existingFunctions = { bodyStart: open + 1, bodyEnd: close };
      break;
    }
  }

  if (existingFunctions) {
    const existingBody = source.slice(existingFunctions.bodyStart, existingFunctions.bodyEnd);
    for (const entry of allEntries) {
      if (new RegExp(`\\bfunction\\s+${entry.name}\\b`).test(existingBody)) {
        return {
          skip: `'${entry.name}' already exists in the page-level functions block`,
        };
      }
    }
  }

  const designatedBlock = blocks.find((block) => perBlockFuncBlocks.has(block));
  const edits: Edit[] = [];
  // True when a mode block keeps non-function content (e.g. `api` entries)
  // after its functions are extracted. The grammar rejects any `ssr {}` /
  // `client {}` wrapper outright, regardless of what's inside, so such a
  // leftover wrapper cannot parse until move-api-blocks removes it in the
  // same `wrnexus update` run. Validating with `parse` here would therefore
  // always fail on a case this migration is specifically meant to unblock.
  let leavesModeWrapper = false;

  for (const block of blocks) {
    const funcBlocks = perBlockFuncBlocks.get(block);
    if (!funcBlocks) continue; // block had no functions sub-block; leave it untouched entirely

    const body = source.slice(block.bodyStart, block.bodyEnd);
    let remainder = body;
    for (const funcBlock of [...funcBlocks].reverse()) {
      remainder = remainder.slice(0, funcBlock.localStart) + remainder.slice(funcBlock.localEnd);
    }
    const isEmpty = remainder.trim() === "";
    if (!isEmpty) leavesModeWrapper = true;

    if (block === designatedBlock && !existingFunctions) {
      const funcsBody = allEntries.map((entry) => `    shared ${entry.text}`).join("\n\n");
      const funcsBlockText = `functions {\n${funcsBody}\n  }`;
      const replacement = isEmpty
        ? funcsBlockText
        : `${funcsBlockText}\n\n  ${block.mode} {${remainder}}`;
      edits.push({ start: block.start, end: block.end, replacement });
    } else {
      const replacement = isEmpty ? "" : `${block.mode} {${remainder}}`;
      edits.push({ start: block.start, end: block.end, replacement });
    }
  }

  if (existingFunctions) {
    const existingBody = source.slice(existingFunctions.bodyStart, existingFunctions.bodyEnd);
    const additions = allEntries.map((entry) => `    shared ${entry.text}`).join("\n\n");
    const mergedBody = `${existingBody.replace(/\s*$/, "")}\n\n${additions}\n  `;
    edits.push({
      start: existingFunctions.bodyStart,
      end: existingFunctions.bodyEnd,
      replacement: mergedBody,
    });
  }

  let after = applyEdits(source, edits);
  after = after.replace(/[ \t]+\n/g, "\n").replace(/\n{3,}/g, "\n\n");

  if (after === source) return { source, changed: false };

  if (leavesModeWrapper) {
    // The only path that writes without a parse check. A surviving `ssr {`/`client {`
    // wrapper still holding `api` entries is unparseable by design -- `move-api-blocks`
    // finishes the job later in the same run -- so `parse` cannot validate it here.
    // Brace balance is the one invariant still checkable, and it is what a bad splice
    // offset would break. Downstream nothing else would catch it: the parser rejects a
    // corrupted wrapper and an untouched one identically.
    if (braceBalance(after) !== braceBalance(source)) {
      throw new Error(
        "mode-functions migration unbalanced the source braces; refusing to write the file",
      );
    }
  } else {
    try {
      parse(after);
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      throw new Error(`mode-functions migration produced source that failed to parse: ${message}`, {
        cause: error,
      });
    }
  }

  return { source: after, changed: true };
}