client/src/components/BlockEditor/customFunctionBlocks.ts

import * as Blockly from 'blockly';
import { javascriptGenerator } from 'blockly/javascript';
import { encodeName } from '../../blocks/utils';

// ============================================================
// Custom function blocks ("define" hats, callers, param ovals)
//
// A user-defined function is stored as a `rb_custom_fn` hat block on the
// workspace. Its description is a list of "segments":
//   { kind: 'text', text: 'Generate Primes ' }
//   { kind: 'param', name: 'n', type: 'variable' | 'boolean' }
// The block renders the segments as an inline phrase where each `param`
// segment becomes a value input socket:
//   define myFunction  Generate Primes ▢ through ▢  ...
//
// `param` segments spawn small variable ovals (`rb_custom_param_var`) that sit
// in the hat's sockets. They can be dragged out into other blocks. When used
// outside their own function they are greyed out (disabled).
//
// Caller blocks (`rb_custom_call` / `rb_custom_call_value`) are added to the
// functions category of the flyout and render the same phrase with empty
// sockets to fill in. A caller placed inside its own definition hat is greyed
// out.
//
// Everything is serialized to workspace XML: the segment list travels in the
// block's <mutation> element (and its `data` property), the param ovals carry
// their owning function name in `data`, and the normal field/connection
// elements carry the rest. Definitions live on the workspace like any other
// block, so saving + WebSocket sync of the blocks XML already persists them.
// ============================================================

export type FnSegment =
  | { kind: 'text'; text: string }
  | { kind: 'param'; name: string; type: 'variable' | 'boolean' };

export interface FnData {
  name: string;
  segments: FnSegment[];
  asyncFn?: boolean;
  returns?: boolean;
}

const FN_COLOR = '#E91E63';
const OVAL_COLOR = '#FF8C1A';

function paramSegsOf(segments: FnSegment[]): Extract<FnSegment, { kind: 'param' }>[] {
  return (segments || []).filter(
    (s): s is Extract<FnSegment, { kind: 'param' }> => s.kind === 'param'
  );
}

function escXml(s: string): string {
  return String(s)
    .replace(/&/g, '&amp;')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;')
    .replace(/'/g, '&apos;');
}

// Read the FnData object stored on a block's `data` property.
export function parseFnData(block: Blockly.Block): FnData | null {
  try {
    const d = JSON.parse(String(block.data || '{}'));
    if (d && typeof d === 'object' && typeof d.name === 'string') {
      return d as FnData;
    }
  } catch {}
  return null;
}

export function makeFnData(
  name: string,
  segments: FnSegment[],
  asyncFn: boolean,
  returns: boolean
): FnData {
  return { name, segments, asyncFn, returns };
}

export function fnParamNames(data: FnData): string[] {
  return paramSegsOf(data.segments).map((s) => encodeName(s.name));
}

function countParams(data: FnData): number {
  return paramSegsOf(data.segments).length;
}

function isFnType(type: string): boolean {
  return type === 'rb_custom_fn';
}
function isCallType(type: string): boolean {
  return type === 'rb_custom_call' || type === 'rb_custom_call_value';
}

function createParamOval(
  workspace: Blockly.Workspace,
  funcName: string,
  name: string,
  type: 'variable' | 'boolean' = 'variable'
): Blockly.Block {
  const block = workspace.newBlock('rb_custom_param_var') as Blockly.Block;
  block.setFieldValue(name, 'name');
  block.data = String(funcName);
  if (type === 'boolean') {
    try {
      block.outputConnection!.setCheck('Boolean');
    } catch {}
  }
  const svg = block as any;
  if (typeof svg.initSvg === 'function') svg.initSvg();
  if (typeof svg.render === 'function') svg.render();
  return block;
}

// Rebuild a custom block's layout from its FnData. Used when a block is
// created (XML mutation), when it is edited, and as a safety net after load.
export function applyFnLayout(block: Blockly.Block, data: FnData, opts?: { asyncFn?: boolean }): void {
  const type = block.type;
  const isFn = isFnType(type);
  const isCall = isCallType(type);
  if (!isFn && !isCall) return;

  // Preserve the body statement input contents (fn hats only).
  const bodyInput = block.getInput('body');
  let bodyBlock: Blockly.Block | null = null;
  if (bodyInput && bodyInput.connection && bodyInput.connection.targetBlock()) {
    bodyBlock = bodyInput.connection.targetBlock() as Blockly.Block;
    try {
      bodyInput.connection.disconnect();
    } catch {}
  }

  // Preserve the current async/returns toggle values (fn hats only).
  const prevAsync = block.getFieldValue('asyncFn') === 'TRUE';
  const prevReturns = block.getFieldValue('returns') === 'TRUE';
  // Snapshot param ovals currently attached to this block (edit case).
  const attachedOvals: { segName: string; oval: Blockly.Block }[] = [];
  for (const input of [...block.inputList]) {
    if (
      input.name &&
      input.name.startsWith('PARAM') &&
      input.connection &&
      input.connection.targetBlock()
    ) {
      const t = input.connection.targetBlock() as Blockly.Block;
      if (t.type === 'rb_custom_param_var') {
        attachedOvals.push({ segName: String(t.getFieldValue('name') || ''), oval: t });
      }
    }
  }

  // Remove every input (values/connections get disconnected; ovals are kept).
  for (const input of [...block.inputList]) {
    try {
      block.removeInput(input.name);
    } catch {}
  }

  // Rebuild the phrase row (all inline so text + sockets share one row).
  let pendingText = '';
  let paramIdx = 0;
  const paramSegs = paramSegsOf(data.segments);

  if (isFn) {
    const row = block.appendDummyInput();
    row.appendField('define');
    row.appendField(new Blockly.FieldLabel(data.name || 'myFunction'), 'name');
  } else if (isCall) {
    const head = block.appendDummyInput();
    head.appendField(new Blockly.FieldLabel(data.name || 'myFunction'), 'fnName');
  }

  for (const seg of data.segments || []) {
    if (seg.kind === 'text') {
      pendingText += seg.text;
    } else {
      const vi = block.appendValueInput('PARAM' + paramIdx);
      if (pendingText) vi.appendField(pendingText);
      pendingText = '';
      if (seg.type === 'boolean') vi.setCheck('Boolean');
      paramIdx++;
    }
  }

  if (isFn) {
    // Trailing text + async/returns toggles live on one row after the phrase.
    const asyncVal = opts?.asyncFn !== undefined ? opts.asyncFn : data.asyncFn !== undefined ? data.asyncFn : prevAsync;
    const returnsVal = data.returns !== undefined ? data.returns : prevReturns;
    const props = block.appendDummyInput();
    if (pendingText) props.appendField(pendingText);
    props.appendField(new Blockly.FieldCheckbox(asyncVal), 'asyncFn');
    props.appendField('async');
    props.appendField(new Blockly.FieldCheckbox(returnsVal), 'returns');
    props.appendField('return value');
    pendingText = '';
    block.appendStatementInput('body');
    if (bodyBlock && block.getInput('body')?.connection) {
      try {
        block.getInput('body')!.connection!.connect(bodyBlock.previousConnection!);
      } catch {}
    }
  } else if (isCall && pendingText) {
    const trail = block.appendDummyInput();
    trail.appendField(pendingText);
  }

  block.setInputsInline(true);

  // Attach param ovals. Only for rendered blocks on a real workspace
  // (during load the ovals arrive from the XML <value> elements instead).
  const rendered = !!(block as any).rendered;
  const ws = block.workspace as any;
  const inFlyout = ws && typeof ws.isFlyout === 'function' ? ws.isFlyout() : false;
  if (rendered && !inFlyout && ws) {
    const used = new Set<string>();
    paramSegs.forEach((seg, i) => {
      const vi = block.getInput('PARAM' + i);
      if (!vi || !vi.connection) return;
      let oval = attachedOvals.find((o) => o.segName === seg.name)?.oval || null;
      if (!oval) {
        oval = createParamOval(block.workspace, data.name, seg.name, seg.type);
      } else {
        oval.data = String(data.name);
        try {
          if (seg.type === 'boolean') oval.outputConnection!.setCheck('Boolean');
          else oval.outputConnection!.setCheck(null);
          (oval as any).render && (oval as any).render();
        } catch {}
      }
      used.add(seg.name);
      try {
        vi.connection!.connect(oval.outputConnection!);
      } catch {}
    });

    for (const { oval } of attachedOvals) {
      const nm = String(oval.getFieldValue('name') || '');
      if (!used.has(nm)) {
        try {
          oval.dispose(false);
        } catch {}
      }
    }
  }
}

function ensureLayout(block: Blockly.Block, data: FnData): void {
  const expected = countParams(data);
  const actual = block.inputList.filter((inp) => inp.name && inp.name.startsWith('PARAM')).length;
  if (actual !== expected) {
    applyFnLayout(block, data);
  }
}

// Update a caller block's function name (data + rendered label) without
// rebuilding its layout, so connected value blocks stay attached.
export function updateFnCallerName(block: Blockly.Block, newName: string): void {
  const data = parseFnData(block);
  if (data && data.name !== newName) {
    data.name = newName;
    block.data = JSON.stringify(data);
    try {
      block.setFieldValue(newName, 'fnName');
    } catch {}
  }
}

// Rename a function everywhere it is referenced (ovals + caller blocks).
export function renameCustomFunction(ws: Blockly.Workspace, oldName: string, newName: string): void {
  if (!oldName || oldName === newName) return;
  for (const b of ws.getBlocksByType('rb_custom_param_var', false)) {
    if (String(b.data) === oldName) b.data = newName;
  }
  for (const t of ['rb_custom_call', 'rb_custom_call_value']) {
    for (const b of ws.getBlocksByType(t, false)) {
      const data = parseFnData(b);
      if (data && data.name === oldName) updateFnCallerName(b, newName);
    }
  }
}

// True when `block` lives inside the definition hat it belongs to.
function isInsideOwnFn(block: Blockly.Block): boolean {
  let parent = block.getSurroundParent();
  while (parent) {
    if (parent.type === 'rb_custom_fn') {
      const fnName = parent.getFieldValue('name');
      if (block.type === 'rb_custom_param_var') {
        return block.data === fnName;
      }
      if (isCallType(block.type)) {
        const data = parseFnData(block);
        return !!data && data.name === fnName;
      }
    }
    parent = parent.getSurroundParent();
  }
  return false;
}

// Re-apply deletable flags and grey-out rules across the whole workspace.
export function syncCustomFunctionBlocks(ws: Blockly.Workspace): void {
  for (const b of ws.getBlocksByType('rb_custom_fn', false)) {
    b.setDeletable(false);
    const data = parseFnData(b);
    if (data) {
      const inlineName = String(b.getFieldValue('name') || '');
      if (inlineName && inlineName !== data.name) {
        renameCustomFunction(ws, data.name, inlineName);
        data.name = inlineName;
        b.data = JSON.stringify(data);
      }
      ensureLayout(b, data);
      // Keep a param oval in every socket so dragging one out always leaves
      // another behind (like Scratch custom-block parameters).
      paramSegsOf(data.segments).forEach((seg, i) => {
        const vi = b.getInput('PARAM' + i);
        if (vi && vi.connection && !vi.connection.targetBlock()) {
          try {
            const oval = createParamOval(b.workspace, data.name, seg.name, seg.type);
            vi.connection.connect(oval.outputConnection!);
          } catch {}
        }
      });
    }
  }
  for (const type of ['rb_custom_call', 'rb_custom_call_value']) {
    for (const b of ws.getBlocksByType(type, false)) {
      const data = parseFnData(b);
      if (data) ensureLayout(b, data);
    }
  }
  for (const b of ws.getBlocksByType('rb_custom_param_var', false)) {
    const desiredEnabled = isInsideOwnFn(b);
    if (b.isEnabled() !== desiredEnabled) b.setEnabled(desiredEnabled);
  }
  for (const type of ['rb_custom_call', 'rb_custom_call_value']) {
    for (const b of ws.getBlocksByType(type, false)) {
      const desiredDisabled = isInsideOwnFn(b);
      if (b.isEnabled() === desiredDisabled) b.setEnabled(!desiredDisabled);
    }
  }
}

// XML string that creates a function definition hat + its param ovals.
export function buildCustomFunctionBlockXml(
  name: string,
  segments: FnSegment[],
  asyncFn: boolean,
  returns: boolean
): string {
  const segJson = JSON.stringify(segments);
  const paramSegs = segments.filter((s) => s.kind === 'param');
  const values = (paramSegs as Extract<FnSegment, { kind: 'param' }>[])
    .map(
      (seg, i) =>
        `<value name="PARAM${i}"><block type="rb_custom_param_var">` +
        `<field name="name">${escXml(seg.name)}</field>` +
        `<data>${escXml(name)}</data></block></value>`
    )
    .join('');
  return (
    `<xml><block type="rb_custom_fn">` +
    `<mutation name="${escXml(name)}" segments="${escXml(segJson)}" async="${asyncFn ? 'true' : 'false'}" returns="${returns ? 'true' : 'false'}"></mutation>` +
    `<field name="name">${escXml(name)}</field>` +
    `<field name="asyncFn">${asyncFn ? 'TRUE' : 'FALSE'}</field>` +
    `<field name="returns">${returns ? 'TRUE' : 'FALSE'}</field>` +
    `<statement name="body"></statement>` +
    values +
    `</block></xml>`
  );
}

// Toolbox entries for the caller blocks of every defined function.
export function customFunctionToolboxBlocks(ws: Blockly.Workspace): any[] {
  const result: any[] = [];
  for (const b of ws.getBlocksByType('rb_custom_fn', false)) {
    const data = parseFnData(b);
    if (!data || !data.name) continue;
    const mutation =
      `<mutation name="${escXml(data.name)}" ` +
      `segments="${escXml(JSON.stringify(data.segments || []))}" ` +
      `async="${data.asyncFn ? 'true' : 'false'}" ` +
      `returns="${data.returns ? 'true' : 'false'}"></mutation>`;
    result.push({ kind: 'block', blockxml: `<block type="rb_custom_call">${mutation}</block>` });
    if (data.returns) {
      result.push({ kind: 'block', blockxml: `<block type="rb_custom_call_value">${mutation}</block>` });
    }
  }
  return result;
}

function registerCodeGen(): void {
  // Statement blocks (fn hats, statement callers) must return a plain string:
  // Blockly's blockToCode throws `TypeError("Expecting string from statement
  // block: ...")` if a generator returns a [code, order] tuple for a block
  // without an output connection.
  //
  // Generators are invoked as func.call(block, block, generator), so generator
  // methods (statementToCode / valueToCode) come from the SECOND parameter;
  // using `this` would resolve to the block and crash code generation.
  javascriptGenerator.forBlock['rb_custom_fn'] = function (block: any, gen: any) {
    const data = parseFnData(block as Blockly.Block);
    const name = encodeName(block.getFieldValue('name') || data?.name || 'myFunction');
    const asyncFn = block.getFieldValue('asyncFn') === 'TRUE';
    const params = (data ? fnParamNames(data) : []).join(', ');
    const body = gen.statementToCode(block, 'body');
    return `${asyncFn ? 'async ' : ''}function ${name}(${params}) {\n${body}\n}`;
  };
  javascriptGenerator.forBlock['rb_custom_call'] = function (block: any, gen: any) {
    return callCode(block as Blockly.Block, gen, false);
  };
  javascriptGenerator.forBlock['rb_custom_call_value'] = function (block: any, gen: any) {
    return [callCode(block as Blockly.Block, gen, true), javascriptGenerator.ORDER_ATOMIC];
  };
  javascriptGenerator.forBlock['rb_custom_param_var'] = function (block: any) {
    return [encodeName(String(block.getFieldValue('name') || '')), javascriptGenerator.ORDER_ATOMIC];
  };
}

function callCode(block: Blockly.Block, gen: any, isValue: boolean): string {
  const data = parseFnData(block);
  const name = encodeName(data?.name || 'myFunction');
  const params = paramSegsOf(data?.segments || []);
  const args = params
    .map((_, i) => {
      try {
        // Empty param sockets become `undefined` so the call is always valid
        // JS (an empty arg would render `f(, )`, a syntax error).
        return gen.valueToCode(block, 'PARAM' + i, javascriptGenerator.ORDER_ATOMIC) || 'undefined';
      } catch {
        return 'undefined';
      }
    })
    .join(', ');
  return isValue ? `${name}(${args})` : `${name}(${args});`;
}

// Register the custom block types + code generators. Idempotent.
export function initCustomFunctionBlocks(): void {
  if ((Blockly.Blocks as any)['rb_custom_fn']) {
    return;
  }

  Blockly.Blocks['rb_custom_fn'] = {
    init: function (this: Blockly.Block) {
      this.setColour(FN_COLOR);
      this.setTooltip('A function you define with your own inputs.');
      this.setDeletable(false);
      this.appendDummyInput();
      this.setInputsInline(true);
    },
    domToMutation: function (this: Blockly.Block, mutation: Element) {
      try {
        const data: FnData = {
          name: mutation.getAttribute('name') || '',
          segments: JSON.parse(mutation.getAttribute('segments') || '[]'),
          asyncFn: mutation.getAttribute('async') === 'true',
          returns: mutation.getAttribute('returns') === 'true',
        };
        this.data = JSON.stringify(data);
        applyFnLayout(this, data);
      } catch {}
    },
    mutationToDom: function (this: Blockly.Block) {
      const m = Blockly.utils.xml.createElement('mutation');
      try {
        const d = parseFnData(this) || { name: '', segments: [] };
        m.setAttribute('name', String(d.name || ''));
        m.setAttribute('segments', JSON.stringify(d.segments || []));
        m.setAttribute('async', d.asyncFn ? 'true' : 'false');
        m.setAttribute('returns', d.returns ? 'true' : 'false');
      } catch {}
      return m;
    },
    customContextMenu: function (this: Blockly.Block, options: any[]) {
      options.push({
        text: 'Edit function...',
        enabled: true,
        callback: function (this: any) {
          const hook = (window as any).__rbEditCustomFn;
          if (typeof hook === 'function') hook(this);
        },
      });
    },
  };

  Blockly.Blocks['rb_custom_call'] = {
    init: function (this: Blockly.Block) {
      this.setColour(FN_COLOR);
      this.setTooltip('Run a function you defined.');
      this.setPreviousStatement(true);
      this.setNextStatement(true);
      this.appendDummyInput();
      this.setInputsInline(true);
    },
    domToMutation: function (this: Blockly.Block, mutation: Element) {
      try {
        const data: FnData = {
          name: mutation.getAttribute('name') || '',
          segments: JSON.parse(mutation.getAttribute('segments') || '[]'),
          returns: mutation.getAttribute('returns') === 'true',
        };
        this.data = JSON.stringify(data);
        applyFnLayout(this, data);
      } catch {}
    },
    mutationToDom: function (this: Blockly.Block) {
      const m = Blockly.utils.xml.createElement('mutation');
      try {
        const d = parseFnData(this) || { name: '', segments: [] };
        m.setAttribute('name', String(d.name || ''));
        m.setAttribute('segments', JSON.stringify(d.segments || []));
      } catch {}
      return m;
    },
  };

  Blockly.Blocks['rb_custom_call_value'] = {
    init: function (this: Blockly.Block) {
      this.setColour(FN_COLOR);
      this.setTooltip('Run a function and use its returned value.');
      this.setOutput(true);
      this.appendDummyInput();
      this.setInputsInline(true);
    },
    domToMutation: function (this: Blockly.Block, mutation: Element) {
      try {
        const data: FnData = {
          name: mutation.getAttribute('name') || '',
          segments: JSON.parse(mutation.getAttribute('segments') || '[]'),
          returns: mutation.getAttribute('returns') === 'true',
        };
        this.data = JSON.stringify(data);
        applyFnLayout(this, data);
      } catch {}
    },
    mutationToDom: function (this: Blockly.Block) {
      const m = Blockly.utils.xml.createElement('mutation');
      try {
        const d = parseFnData(this) || { name: '', segments: [] };
        m.setAttribute('name', String(d.name || ''));
        m.setAttribute('segments', JSON.stringify(d.segments || []));
      } catch {}
      return m;
    },
  };

  Blockly.Blocks['rb_custom_param_var'] = {
    init: function (this: Blockly.Block) {
      this.setColour(OVAL_COLOR);
      this.setOutput(true);
      this.setTooltip('A value belonging to a function. Use it inside that function.');
      this.appendDummyInput().appendField(new Blockly.FieldLabel('param'), 'name');
      this.setInputsInline(true);
      this.setEditable(false);
    },
  };

  registerCodeGen();
}