import * as Blockly from 'blockly';
import { pythonGenerator } from 'blockly/python';
import { registerFieldMultilineInput } from '@blockly/field-multilineinput';

registerFieldMultilineInput();


// Utility to create a unique input reference block type
function createInputRefBlockType(inputName) {
  const blockType = `input_reference_${inputName}`;
  if (!Blockly.Blocks[blockType]) {
    Blockly.Blocks[blockType] = {
      init: function () {
        this.jsonInit({
          type: blockType,
          message0: "%1",
          args0: [
            {
              type: "field_label_serializable",
              name: "VARNAME",
              text: inputName
            }
          ],
          output: null,
          colour: 255,
          outputShape: 2
        });
      },

      // Save the owner block ID for proper restoration after deserialization
      saveExtraState: function () {
        return {
          ownerBlockId: this._ownerBlockId || null
        };
      },

      // Restore the owner block ID
      loadExtraState: function (state) {
        if (state.ownerBlockId) {
          this._ownerBlockId = state.ownerBlockId;
        }
      }
    };
    pythonGenerator.forBlock[blockType] = function () {
      return [inputName, pythonGenerator.ORDER_ATOMIC];
    };
  }
  return blockType;
}

// Global input reference tracking map
const inputRefs = new Map();

// Core mutator registration for dynamic tool and input creation
Blockly.Extensions.registerMutator(
  'test_mutator',
  {
    initialize: function () {
      if (!this.initialized_) {
        this.inputCount_ = 0;
        this.inputNames_ = [];
        this.inputTypes_ = [];
        this.inputRefBlocks_ = new Map();
        this.outputCount_ = 0;
        this.outputNames_ = [];
        this.outputTypes_ = [];
        this.initialized_ = true;
        // Mark all reference blocks with their owner for later identification
        this._ownerBlockId = this.id;
      }
    },

    decompose: function (workspace) {
      const containerBlock = workspace.newBlock('container');
      containerBlock.initSvg();
      let connection = containerBlock.getInput('STACK').connection;

      this.inputCount_ = this.inputCount_ || 0;
      this.inputNames_ = this.inputNames_ || [];
      this.inputTypes_ = this.inputTypes_ || [];
      this.outputCount_ = this.outputCount_ || 0;
      this.outputNames_ = this.outputNames_ || [];
      this.outputTypes_ = this.outputTypes_ || [];

      // Restore dynamically added input items
      for (let i = 0; i < this.inputCount_; i++) {
        const itemBlock = workspace.newBlock('container_input');
        itemBlock.initSvg();
        const typeVal = this.inputTypes_[i] || 'string';
        const nameVal = this.inputNames_[i] || typeVal;
        itemBlock.setFieldValue(typeVal, 'TYPE');
        itemBlock.setFieldValue(nameVal, 'NAME');

        const input = this.getInput('X' + i);
        if (input && input.connection && input.connection.targetConnection) {
          itemBlock.valueConnection_ = input.connection.targetConnection;
        }

        connection.connect(itemBlock.previousConnection);
        connection = itemBlock.nextConnection;
      }

      // Restore dynamically added output items
      let connection2 = containerBlock.getInput('STACK2').connection;
      for (let i = 0; i < this.outputCount_; i++) {
        const itemBlock = workspace.newBlock('container_output');
        itemBlock.initSvg();
        const typeVal = this.outputTypes_[i] || 'string';
        const nameVal = this.outputNames_[i] || typeVal;
        itemBlock.setFieldValue(typeVal, 'TYPE');
        itemBlock.setFieldValue(nameVal, 'NAME');

        connection2.connect(itemBlock.previousConnection);
        connection2 = itemBlock.nextConnection;
      }

      return containerBlock;
    },

    compose: function (containerBlock) {
      Blockly.Events.disable();
      try {
        if (!this.initialized_) this.initialize();

        // Ensure the inputRefBlocks_ map is properly initialized
        if (!this.inputRefBlocks_) {
          this.inputRefBlocks_ = new Map();
        }

        // Check if we need to find existing reference blocks in the workspace
        // This happens after deserialization when the blocks exist but aren't tracked
        if (this.inputRefBlocks_.size === 0 && this.inputNames_ && this.inputNames_.length > 0) {
          for (let i = 0; i < this.inputNames_.length; i++) {
            const name = this.inputNames_[i];
            const input = this.getInput('X' + i);

            // Check if there's already a connected block in this input
            if (input && input.connection && input.connection.targetBlock()) {
              const connectedBlock = input.connection.targetBlock();
              const expectedType = `input_reference_${name}`;

              // If this is the expected reference block, track it
              if (connectedBlock.type === expectedType && connectedBlock._ownerBlockId === this.id) {
                this.inputRefBlocks_.set(name, connectedBlock);
              }
            }
          }
        }

        const oldNames = [...(this.inputNames_ || [])];
        const oldOutputNames = [...(this.outputNames_ || [])];
        const connections = [];
        const returnConnections = [];
        let itemBlock = containerBlock.getInputTargetBlock('STACK');

        // Collect all child connections from mutator stack
        while (itemBlock) {
          connections.push(itemBlock.valueConnection_);
          itemBlock = itemBlock.nextConnection && itemBlock.nextConnection.targetBlock();
        }

        // Save existing return connections before removing them
        let rIdx = 0;
        while (this.getInput('R' + rIdx)) {
          const returnInput = this.getInput('R' + rIdx);
          if (returnInput && returnInput.connection && returnInput.connection.targetConnection) {
            returnConnections.push(returnInput.connection.targetConnection);
          } else {
            returnConnections.push(null);
          }
          rIdx++;
        }

        // Collect output specifications from STACK2
        const outputSpecs = [];
        let outputBlock = containerBlock.getInputTargetBlock('STACK2');
        while (outputBlock) {
          outputSpecs.push(outputBlock);
          outputBlock = outputBlock.nextConnection && outputBlock.nextConnection.targetBlock();
        }

        const newCount = connections.length;
        const newOutputCount = outputSpecs.length;
        this.inputCount_ = newCount;
        this.outputCount_ = newOutputCount;
        this.inputNames_ = this.inputNames_ || [];
        this.inputTypes_ = this.inputTypes_ || [];
        this.outputNames_ = this.outputNames_ || [];
        this.outputTypes_ = this.outputTypes_ || [];

        // Rebuild the new list of input names and types
        let idx = 0;
        let it = containerBlock.getInputTargetBlock('STACK');
        const newNames = [];
        while (it) {
          this.inputTypes_[idx] = it.getFieldValue('TYPE') || 'string';
          this.inputNames_[idx] = it.getFieldValue('NAME') || 'arg' + idx;
          newNames.push(this.inputNames_[idx]);
          it = it.nextConnection && it.nextConnection.targetBlock();
          idx++;
        }

        // Rebuild the new list of output names and types
        let oidx = 0;
        const newOutputNames = [];
        for (const outBlock of outputSpecs) {
          this.outputTypes_[oidx] = outBlock.getFieldValue('TYPE') || 'string';
          this.outputNames_[oidx] = outBlock.getFieldValue('NAME') || 'output' + oidx;
          newOutputNames.push(this.outputNames_[oidx]);
          oidx++;
        }

        // Dispose of removed input reference blocks when inputs shrink
        if (newCount < oldNames.length) {
          for (let i = newCount; i < oldNames.length; i++) {
            const oldName = oldNames[i];
            const block = this.inputRefBlocks_.get(oldName);
            if (block && !block.disposed) block.dispose(true);
            this.inputRefBlocks_.delete(oldName);
          }
        }

        // Rename reference blocks when variable names change
        // Only update reference blocks that belong to THIS block
        for (let i = 0; i < Math.min(oldNames.length, newNames.length); i++) {
          const oldName = oldNames[i];
          const newName = newNames[i];
          if (oldName !== newName) {
            const oldBlockType = `input_reference_${oldName}`;
            const newBlockType = `input_reference_${newName}`;

            if (this.inputRefBlocks_.has(oldName)) {
              const refBlock = this.inputRefBlocks_.get(oldName);
              if (refBlock && !refBlock.disposed) {
                this.inputRefBlocks_.delete(oldName);
                this.inputRefBlocks_.set(newName, refBlock);
                refBlock.setFieldValue(newName, 'VARNAME');
                // Properly update the block type in workspace tracking
                if (refBlock.workspace && refBlock.workspace.removeTypedBlock) {
                  refBlock.workspace.removeTypedBlock(refBlock);
                  refBlock.type = newBlockType;
                  refBlock.workspace.addTypedBlock(refBlock);
                } else {
                  refBlock.type = newBlockType;
                }
              }
            }

            // Update all clones of this reference block that share the same owner
            // (i.e., all clones that were created from the same parent block)
            const refBlock = this.inputRefBlocks_.get(newName);
            const ownerBlockId = this.id;

            if (refBlock && !refBlock.disposed) {
              const allBlocks = this.workspace.getAllBlocks(false);
              for (const block of allBlocks) {
                if (block.type === oldBlockType) {
                  // Update if this block has the same owner as our reference block
                  // This includes both connected and cloned blocks
                  if (block._ownerBlockId === ownerBlockId) {
                    // Properly update the block type in workspace tracking
                    if (block.workspace && block.workspace.removeTypedBlock) {
                      block.workspace.removeTypedBlock(block);
                      block.type = newBlockType;
                      block.workspace.addTypedBlock(block);
                    } else {
                      block.type = newBlockType;
                    }
                    block.setFieldValue(newName, 'VARNAME');
                  }
                }
              }
            }

            pythonGenerator.forBlock[newBlockType] = function () {
              return [newName, pythonGenerator.ORDER_ATOMIC];
            };
          }
        }

        // Remove all dynamic and temporary inputs before reconstruction
        let i = 0;
        while (this.getInput('X' + i)) this.removeInput('X' + i++);
        let r = 0;
        while (this.getInput('R' + r)) this.removeInput('R' + r++);
        let t = 0;
        while (this.getInput('T' + t)) this.removeInput('T' + t++);
        ['INPUTS_TEXT', 'RETURNS_TEXT', 'TOOLS_TEXT'].forEach(name => {
          if (this.getInput(name)) this.removeInput(name);
        });

        if (newCount > 0) {
          const inputsText = this.appendDummyInput('INPUTS_TEXT');
          inputsText.appendField('with inputs:');
          this.moveInputBefore('INPUTS_TEXT', 'BODY');
        }

        // Add each dynamic input, reconnecting to reference blocks
        for (let j = 0; j < newCount; j++) {
          const type = this.inputTypes_[j] || 'string';
          const name = this.inputNames_[j] || type;
          let check = null;
          if (type === 'integer') check = 'Number';
          if (type === 'float') check = 'Number';
          if (type === 'string') check = 'String';

          const existingRefBlock = this.inputRefBlocks_.get(name);
          const input = this.appendValueInput('X' + j);
          if (check) input.setCheck(check);
          input.appendField(type);
          this.moveInputBefore('X' + j, 'BODY');

          const blockType = createInputRefBlockType(name);

          // Check if there's already a block connected to this input
          const currentlyConnected = input.connection ? input.connection.targetBlock() : null;

          if (currentlyConnected && currentlyConnected.type === blockType) {
            // There's already the correct reference block connected, just track it
            currentlyConnected._ownerBlockId = this.id;
            this.inputRefBlocks_.set(name, currentlyConnected);
          } else if (!existingRefBlock) {
            // Only create a new reference block if none exists and nothing is connected
            const refBlock = this.workspace.newBlock(blockType);
            refBlock.initSvg();
            refBlock.setDeletable(false);
            refBlock.render();
            // Mark the reference block with its owner
            refBlock._ownerBlockId = this.id;
            this.inputRefBlocks_.set(name, refBlock);
            if (input.connection && refBlock.outputConnection) {
              input.connection.connect(refBlock.outputConnection);
            }
          } else {
            // Reference block exists - only reconnect if not already connected to this input
            existingRefBlock._ownerBlockId = this.id;
            if (input.connection && !input.connection.targetBlock() && existingRefBlock.outputConnection) {
              input.connection.connect(existingRefBlock.outputConnection);
            }
          }

          pythonGenerator.forBlock[blockType] = function () {
            return [name, pythonGenerator.ORDER_ATOMIC];
          };
        }

        // Reconnect preserved connections to new structure
        for (let k = 0; k < newCount; k++) {
          const conn = connections[k];
          if (conn) {
            try {
              conn.connect(this.getInput('X' + k).connection);
            } catch { }
          }
        }

        // Handle return inputs based on outputs
        if (newOutputCount > 0) {
          // Remove the default RETURN input if it exists
          if (this.getInput('RETURN')) {
            this.removeInput('RETURN');
          }

          // Add the "and return" label
          const returnsText = this.appendDummyInput('RETURNS_TEXT');
          returnsText.appendField('and return');

          // Add each return value input slot
          for (let j = 0; j < newOutputCount; j++) {
            const type = this.outputTypes_[j] || 'string';
            const name = this.outputNames_[j] || ('output' + j);
            let check = null;
            if (type === 'integer') check = 'Number';
            if (type === 'float') check = 'Number';
            if (type === 'string') check = 'String';

            const returnInput = this.appendValueInput('R' + j);
            if (check) returnInput.setCheck(check);
            returnInput.appendField(type);
            returnInput.appendField('"' + name + '":');

            // Reconnect previous connection if it exists
            if (returnConnections[j]) {
              try {
                returnInput.connection.connect(returnConnections[j]);
              } catch { }
            }
          }
        }

        this.workspace.render();
      } finally {
        Blockly.Events.enable();
      }
    },

    saveExtraState: function () {
      return {
        inputCount: this.inputCount_,
        inputNames: this.inputNames_,
        inputTypes: this.inputTypes_,
        outputCount: this.outputCount_,
        outputNames: this.outputNames_,
        outputTypes: this.outputTypes_,
        toolCount: this.toolCount_ || 0
      };
    },

    loadExtraState: function (state) {
      this.inputCount_ = state.inputCount || 0;
      this.inputNames_ = state.inputNames || [];
      this.inputTypes_ = state.inputTypes || [];
      this.outputCount_ = state.outputCount || 0;
      this.outputNames_ = state.outputNames || [];
      this.outputTypes_ = state.outputTypes || [];
      this.toolCount_ = state.toolCount || 0;

      // Ensure the reference block map is initialized
      if (!this.inputRefBlocks_) {
        this.inputRefBlocks_ = new Map();
      }

      // Immediately rebuild the inputs structure so they exist when connections are loaded
      // This must happen BEFORE Blockly tries to restore connections
      if (this.inputCount_ > 0) {
        const inputsText = this.appendDummyInput('INPUTS_TEXT');
        inputsText.appendField('with inputs:');
        this.moveInputBefore('INPUTS_TEXT', 'BODY');

        for (let j = 0; j < this.inputCount_; j++) {
          const type = this.inputTypes_[j] || 'string';
          const name = this.inputNames_[j] || ('arg' + j);
          let check = null;
          if (type === 'integer') check = 'Number';
          if (type === 'float') check = 'Number';
          if (type === 'string') check = 'String';

          const input = this.appendValueInput('X' + j);
          if (check) input.setCheck(check);
          input.appendField(type);
          this.moveInputBefore('X' + j, 'BODY');

          // Create the block type definition if it doesn't exist
          const blockType = createInputRefBlockType(name);
          pythonGenerator.forBlock[blockType] = function () {
            return [name, pythonGenerator.ORDER_ATOMIC];
          };
        }
      }

      // Also rebuild return inputs if there are outputs
      if (this.outputCount_ > 0) {
        if (this.getInput('RETURN')) {
          this.removeInput('RETURN');
        }

        const returnsText = this.appendDummyInput('RETURNS_TEXT');
        returnsText.appendField('and return');

        for (let j = 0; j < this.outputCount_; j++) {
          const type = this.outputTypes_[j] || 'string';
          const name = this.outputNames_[j] || ('output' + j);
          let check = null;
          if (type === 'integer') check = 'Number';
          if (type === 'float') check = 'Number';
          if (type === 'string') check = 'String';

          const returnInput = this.appendValueInput('R' + j);
          if (check) returnInput.setCheck(check);
          returnInput.appendField(type);
          returnInput.appendField('"' + name + '":');
        }
      }
    }
  },
  null,
  ['container_input']
);

// Base block definitions
const container = {
  type: "container",
  message0: "inputs %1 %2 outputs %3 %4",
  args0: [
    { type: "input_dummy", name: "title" },
    { type: "input_statement", name: "STACK" },
    { type: "input_dummy", name: "title2" },
    { type: "input_statement", name: "STACK2" },
  ],
  colour: 210,
  inputsInline: false
};

const container_input = {
  type: "container_input",
  message0: "%1 %2",
  args0: [
    {
      type: "field_dropdown",
      name: "TYPE",
      options: [
        ["String", "string"],
        ["Integer", "integer"],
        ["Float", "float"],
        ["List", "list"],
        ["Boolean", "boolean"],
        ["Any", "any"],
      ]
    },
    { type: "field_input", name: "NAME" },
  ],
  previousStatement: null,
  nextStatement: null,
  colour: 210,
};

const container_output = {
  type: "container_output",
  message0: "%1 %2",
  args0: [
    {
      type: "field_dropdown",
      name: "TYPE",
      options: [
        ["String", "string"],
        ["Integer", "integer"],
        ["Float", "float"],
        ["List", "list"],
        ["Boolean", "boolean"],
        ["Any", "any"],
      ]
    },
    { type: "field_input", name: "NAME" },
  ],
  previousStatement: null,
  nextStatement: null,
  colour: 210,
};

const llm_call = {
  type: "llm_call",
  message0: "call model %1 with prompt %2",
  args0: [
    {
      type: "field_dropdown",
      name: "MODEL",
      options: [
        ["gpt-3.5-turbo", "gpt-3.5-turbo-0125"],
        ["gpt-4o", "gpt-4o-2024-08-06"],
        ["gpt-5-mini", "gpt-5-mini-2025-08-07"],
        ["gpt-5", "gpt-5-2025-08-07"],
        ["gpt-4o search", "gpt-4o-search-preview-2025-03-11"],
      ]
    },
    { type: "input_value", name: "PROMPT", check: "String" },
  ],
  inputsInline: true,
  output: "String",
  colour: 160,
  tooltip: "Call the selected OpenAI model to get a response.",
  helpUrl: "",
};

const call_api = {
  "type": "call_api",
  "message0": "call API with method %1 url %2 headers %3",
  "args0": [
    {
      type: "field_dropdown",
      name: "METHOD",
      options: [
        ["GET", "GET"],
        ["POST", "POST"],
        ["PUT", "PUT"],
        ["DELETE", "DELETE"],
      ]
    },
    {
      "type": "input_value",
      "name": "URL",
    },
    {
      "type": "input_value",
      "name": "HEADERS",
    },
  ],
  "output": ["String", "Integer", "List"],
  "colour": 165,
  "inputsInline": true
}

const in_json = {
  "type": "in_json",
  "message0": "get %1 from JSON %2",
  "args0": [
    {
      "type": "input_value",
      "name": "NAME",
    },
    {
      "type": "input_value",
      "name": "JSON",
    },
  ],
  "output": ["String", "Integer", "List"],
  "colour": 165,
  "inputsInline": true
}

const json_field = {
  type: "json_field",
  message0: "field",
  args0: [],
  previousStatement: null,
  nextStatement: null,
  colour: 165,
};

const make_json_container = {
  type: "make_json_container",
  message0: "fields %1",
  args0: [
    { type: "input_statement", name: "STACK" },
  ],
  colour: 165,
  inputsInline: false
};

const make_json = {
  type: "make_json",
  message0: "make JSON %1",
  args0: [
    { type: "input_dummy" },
  ],
  colour: 165,
  inputsInline: false,
  output: ["String", "Integer"],
  mutator: "json_mutator",
  fieldCount_: 1,
};

const lists_contains = {
  type: "lists_contains",
  message0: "item %1 in list %2",
  args0: [
    { type: "input_value", name: "ITEM", check: null },
    { type: "input_value", name: "LIST", check: "Array" },
  ],
  output: "Boolean",
  colour: 260,
  inputsInline: true,
  tooltip: "Check if an item exists in a list",
  helpUrl: "",
};

// Cast block for type conversion
const cast_as = {
  type: "cast_as",
  message0: "cast %1 as %2",
  args0: [
    { type: "input_value", name: "VALUE", check: null },
    {
      type: "field_dropdown",
      name: "TYPE",
      options: [
        ["int", "int"],
        ["float", "float"],
        ["str", "str"],
        ["bool", "bool"],
      ],
    },
  ],
  output: null,
  colour: 210,
  inputsInline: true,
  tooltip: "Convert a value to a different type",
  helpUrl: "",
};

// Dynamic function call block
const func_call = {
  type: "func_call",
  message0: "func %1",
  args0: [
    {
      type: "field_dropdown",
      name: "FUNC_NAME",
      options: function () {
        // This will be populated dynamically
        const options = [["<no functions>", "NONE"]];
        if (this.sourceBlock_) {
          const workspace = this.sourceBlock_.workspace;
          const funcBlocks = workspace.getAllBlocks(false).filter(b => b.type === 'func_def');
          if (funcBlocks.length > 0) {
            options.length = 0;
            funcBlocks.forEach(block => {
              const name = block.getFieldValue('NAME');
              if (name) {
                options.push([name, name]);
              }
            });
          }
        }
        return options;
      }
    }
  ],
  inputsInline: true,
  output: null,
  colour: 210,
  tooltip: "Call a function defined in the workspace",
  helpUrl: "",
  extensions: ["func_call_dynamic"]
};

const create_mcp = {
  type: "create_mcp",
  message0: "create MCP %1 %2",
  args0: [
    { type: "input_dummy" },
    { type: "input_statement", name: "BODY" },
  ],
  colour: 210,
  inputsInline: true,
  mutator: "test_mutator",
  inputCount_: 0,
  deletable: false,
  extensions: ["test_cleanup_extension"]
};

const func_def = {
  type: "func_def",
  message0: "function %1 %2 %3",
  args0: [
    { type: "field_input", name: "NAME", text: "newFunction" },
    { type: "input_dummy" },
    { type: "input_statement", name: "BODY" },
  ],
  colour: 210,
  inputsInline: true,
  mutator: "test_mutator",
  inputCount_: 0,
  deletable: true,
  extensions: ["test_cleanup_extension"]
};

// Cleanup extension ensures that dynamic reference blocks are deleted when parent is
Blockly.Extensions.register('test_cleanup_extension', function () {
  const oldDispose = this.dispose;
  this.dispose = function (healStack, recursive) {
    if (this.inputRefBlocks_) {
      for (const [, refBlock] of this.inputRefBlocks_) {
        if (refBlock && !refBlock.disposed) refBlock.dispose(false);
      }
      this.inputRefBlocks_.clear();
    }
    if (oldDispose) oldDispose.call(this, healStack, recursive);
  };
});

// JSON mutator for dynamic field creation
Blockly.Extensions.registerMutator(
  'json_mutator',
  {

    decompose: function (workspace) {
      const containerBlock = workspace.newBlock('make_json_container');
      containerBlock.initSvg();
      let connection = containerBlock.getInput('STACK').connection;

      // Initialize defaults if not set
      if (this.fieldCount_ === undefined) {
        this.fieldCount_ = 0;
        this.fieldKeys_ = [];
      }

      // Restore dynamically added field items
      for (let i = 0; i < this.fieldCount_; i++) {
        const itemBlock = workspace.newBlock('json_field');
        itemBlock.initSvg();
        // Store the connection for compose
        const input = this.getInput('FIELD' + i);
        if (input && input.connection && input.connection.targetConnection) {
          itemBlock.valueConnection_ = input.connection.targetConnection;
        }

        connection.connect(itemBlock.previousConnection);
        connection = itemBlock.nextConnection;
      }

      return containerBlock;
    },

    compose: function (containerBlock) {
      Blockly.Events.disable();
      try {
        // Initialize defaults if not set
        if (this.fieldCount_ === undefined) {
          this.fieldCount_ = 0;
          this.fieldKeys_ = [];
        }

        const connections = [];
        let itemBlock = containerBlock.getInputTargetBlock('STACK');

        // Collect all child connections from mutator stack
        while (itemBlock) {
          connections.push(itemBlock.valueConnection_);
          itemBlock = itemBlock.nextConnection && itemBlock.nextConnection.targetBlock();
        }

        const newCount = connections.length;
        const oldCount = this.fieldCount_;
        this.fieldCount_ = newCount;

        // Preserve old keys and extend array if needed
        if (!this.fieldKeys_) {
          this.fieldKeys_ = [];
        }

        // Remove all dynamic inputs before reconstruction
        let i = 0;
        while (this.getInput('FIELD' + i)) this.removeInput('FIELD' + i++);

        // Add each dynamic field input with editable key name
        for (let j = 0; j < newCount; j++) {
          // Use existing key or create new one
          if (!this.fieldKeys_[j]) {
            this.fieldKeys_[j] = 'key' + j;
          }
          const key = this.fieldKeys_[j];
          const input = this.appendValueInput('FIELD' + j);
          const field = new Blockly.FieldTextInput(key);
          field.setValidator((newValue) => {
            // Update the stored key when user edits it
            this.fieldKeys_[j] = newValue || 'key' + j;
            return newValue;
          });
          input.appendField(field, 'KEY' + j);
          input.appendField(':');
          this.moveInputBefore('FIELD' + j, null);
        }

        // Trim fieldKeys array if fields were removed
        if (newCount < oldCount) {
          this.fieldKeys_.length = newCount;
        }

        // Reconnect preserved connections to new structure
        for (let k = 0; k < newCount; k++) {
          const conn = connections[k];
          if (conn) {
            try {
              this.getInput('FIELD' + k).connection.connect(conn);
            } catch { }
          }
        }

        this.workspace.render();
      } finally {
        Blockly.Events.enable();
      }
    },

    saveExtraState: function () {
      return {
        fieldCount: this.fieldCount_,
        fieldKeys: this.fieldKeys_,
      };
    },

    loadExtraState: function (state) {
      this.fieldCount_ = state.fieldCount || 0;
      this.fieldKeys_ = state.fieldKeys || [];

      // Immediately rebuild the inputs structure so they exist when connections are loaded
      if (this.fieldCount_ > 0) {
        for (let j = 0; j < this.fieldCount_; j++) {
          const key = this.fieldKeys_[j] || ('key' + j);
          const input = this.appendValueInput('FIELD' + j);
          const field = new Blockly.FieldTextInput(key);
          field.setValidator((newValue) => {
            // Update the stored key when user edits it
            this.fieldKeys_[j] = newValue || 'key' + j;
            return newValue;
          });
          input.appendField(field, 'KEY' + j);
          input.appendField(':');
        }
      }
    }
  },
  null,
  ['json_field']
);

// Extension for dynamic function call blocks
Blockly.Extensions.register('func_call_dynamic', function () {
  const block = this;

  // Store current function being called
  block.currentFunction_ = null;
  block.paramCount_ = 0;

  // Function to update the block based on selected function
  block.updateShape_ = function () {
    const funcName = this.getFieldValue('FUNC_NAME');

    // Temporarily disable events to prevent recursive updates
    const eventsEnabled = Blockly.Events.isEnabled();
    if (eventsEnabled) {
      Blockly.Events.disable();
    }

    try {
      // Remove all existing parameter inputs
      let i = 0;
      while (this.getInput('ARG' + i)) {
        this.removeInput('ARG' + i);
        i++;
      }

      if (funcName && funcName !== 'NONE') {
        // Find the function definition block
        const workspace = this.workspace;
        const funcBlock = workspace.getAllBlocks(false).find(b =>
          b.type === 'func_def' && b.getFieldValue('NAME') === funcName
        );

        if (funcBlock) {
          this.currentFunction_ = funcName;

          // Get the function's parameters
          const inputCount = funcBlock.inputCount_ || 0;
          const inputNames = funcBlock.inputNames_ || [];
          const inputTypes = funcBlock.inputTypes_ || [];

          this.paramCount_ = inputCount;

          // Add parameter inputs matching the function definition
          for (let j = 0; j < inputCount; j++) {
            const paramName = inputNames[j] || ('arg' + j);
            const paramType = inputTypes[j] || 'string';

            let check = null;
            if (paramType === 'integer') check = 'Number';
            if (paramType === 'string') check = 'String';

            const input = this.appendValueInput('ARG' + j);
            if (check) input.setCheck(check);
            input.appendField(`${paramType} "${paramName}"`);
          }

          // Set output type based on function's return type
          if (funcBlock.outputCount_ && funcBlock.outputCount_ > 0) {
            const outputType = funcBlock.outputTypes_[0] || 'string';
            if (outputType === 'integer') {
              this.setOutput(true, 'Number');
            } else if (outputType === 'string') {
              this.setOutput(true, 'String');
            } else {
              this.setOutput(true, null);
            }
          } else {
            this.setOutput(true, null);
          }
        } else {
          this.currentFunction_ = null;
          this.paramCount_ = 0;
        }
      } else {
        this.currentFunction_ = null;
        this.paramCount_ = 0;
      }
    } finally {
      // Re-enable events if they were enabled before
      if (eventsEnabled) {
        Blockly.Events.enable();
      }
    }
  };

  // Listen for dropdown changes
  block.getField('FUNC_NAME').setValidator(function (newValue) {
    const block = this.getSourceBlock();
    // Ensure the update happens after the dropdown value is set
    setTimeout(() => {
      block.updateShape_();
    }, 0);
    return newValue;
  });

  // Listen for workspace changes to update when functions are modified
  const workspaceListener = function (event) {
    // Skip if block has been disposed
    if (block.disposed) {
      return;
    }

    if (event.type === Blockly.Events.BLOCK_CHANGE ||
      event.type === Blockly.Events.BLOCK_DELETE ||
      event.type === Blockly.Events.BLOCK_CREATE ||
      event.type === Blockly.Events.BLOCK_MOVE) {

      // Check if a func_def block was changed
      const changedBlock = block.workspace.getBlockById(event.blockId);

      if (event.type === Blockly.Events.BLOCK_DELETE) {
        // Delay check to ensure workspace has been updated
        setTimeout(() => {
          if (!block.disposed && block.currentFunction_ &&
            !block.workspace.getAllBlocks(false).some(b =>
              b.type === 'func_def' && b.getFieldValue('NAME') === block.currentFunction_)) {
            block.dispose(true);
          }
        }, 0);
      } else if (changedBlock && changedBlock.type === 'func_def') {
        // If the function being called was modified, update shape
        const funcName = changedBlock.getFieldValue('NAME');
        if (funcName === block.currentFunction_ ||
          (event.oldValue && event.oldValue === block.currentFunction_)) {
          // Handle renaming
          if (event.type === Blockly.Events.BLOCK_CHANGE &&
            event.name === 'NAME' && event.oldValue === block.currentFunction_) {
            block.currentFunction_ = funcName;
            block.setFieldValue(funcName, 'FUNC_NAME');
          }
          setTimeout(() => {
            if (!block.disposed) {
              block.updateShape_();
            }
          }, 0);
        }
      }

      // Update dropdown options
      const dropdown = block.getField('FUNC_NAME');
      if (dropdown) {
        const workspace = block.workspace;
        const funcBlocks = workspace.getAllBlocks(false).filter(b => b.type === 'func_def');
        const options = funcBlocks.length > 0
          ? funcBlocks.map(b => [b.getFieldValue('NAME'), b.getFieldValue('NAME')])
          : [["<no functions>", "NONE"]];

        dropdown.menuGenerator_ = options;

        // If current function no longer exists, reset
        if (block.currentFunction_ && !options.some(opt => opt[1] === block.currentFunction_)) {
          block.setFieldValue('NONE', 'FUNC_NAME');
          setTimeout(() => {
            if (!block.disposed) {
              block.updateShape_();
            }
          }, 0);
        }
      }
    }
  };

  block.workspace.addChangeListener(workspaceListener);

  // Clean up the listener when block is disposed
  const oldDispose = block.dispose;
  block.dispose = function (healStack) {
    if (block.workspace) {
      block.workspace.removeChangeListener(workspaceListener);
    }
    if (oldDispose) {
      oldDispose.call(this, healStack);
    }
  };
});

// Function to generate a unique tool name
function generateUniqueToolName(workspace, excludeBlock) {
  const existingNames = new Set();
  const allBlocks = workspace.getAllBlocks(false);

  // Collect all existing tool names, excluding the block being created
  for (const block of allBlocks) {
    if (block.type === 'func_def' && block !== excludeBlock && block.getFieldValue('NAME')) {
      existingNames.add(block.getFieldValue('NAME'));
    }
  }

  // Generate a unique name
  let baseName = 'newTool';
  let name = baseName;
  let counter = 1;

  while (existingNames.has(name)) {
    counter++;
    name = `${baseName}${counter}`;
  }

  return name;
}

// Register create_mcp block separately to include custom init logic  
Blockly.Blocks['create_mcp'] = {
  init: function () {
    this.jsonInit(create_mcp);
    // Apply extensions
    Blockly.Extensions.apply('test_cleanup_extension', this, false);
    // Initialize mutator state
    if (this.initialize) {
      this.initialize();
    }
  }
};

// Register func_def block separately to include custom init logic
Blockly.Blocks['func_def'] = {
  init: function () {
    this.jsonInit(func_def);
    // Apply extensions
    Blockly.Extensions.apply('test_cleanup_extension', this, false);
    // Initialize mutator state
    if (this.initialize) {
      this.initialize();
    }
  }
};

// Register func_call block separately to include custom logic
Blockly.Blocks['func_call'] = {
  init: function () {
    this.jsonInit(func_call);
  },

  // Save the current function name and parameter count for serialization
  saveExtraState: function () {
    return {
      currentFunction: this.currentFunction_ || null,
      paramCount: this.paramCount_ || 0
    };
  },

  loadExtraState(state) {
    this.currentFunction_ = state.currentFunction;
    this.paramCount_ = state.paramCount;

    // Ensure ARG0..ARGN exist so the deserializer can reconnect children
    for (let i = 0; i < this.paramCount_; i++) {
      if (!this.getInput('ARG' + i)) {
        this.appendValueInput('ARG' + i).appendField("");
      }
    }
  }
};

// Register json_field block (internal mutator block, not user-visible)
Blockly.Blocks['json_field'] = {
  init: function () {
    this.jsonInit(json_field);
  }
};

// Register make_json_container block (internal mutator block, not user-visible)
Blockly.Blocks['make_json_container'] = {
  init: function () {
    this.jsonInit(make_json_container);
  }
};

// Register make_json block separately to include custom init logic
Blockly.Blocks['make_json'] = {
  init: function () {
    this.jsonInit(make_json);
    // Initialize with no fields by default
    this.fieldCount_ = 0;
    this.fieldKeys_ = [];
  }
};

export const blocks = Blockly.common.createBlockDefinitionsFromJsonArray([
  container,
  container_input,
  container_output,
  llm_call,
  call_api,
  in_json,
  lists_contains,
  cast_as,
]);