import { getNodeDataFromID, createNodeData, getOrCreateNode, extractNodeTypeInfo } from './strands/ir_dag.js';
import { recordInBasicBlock } from './strands/ir_cfg.js';
import { userError, internalError } from './strands/strands_FES.js';
import { BaseType, NodeType, OpCode, typeEquals, BasePriority, DataType, OpCodeToSymbol } from './strands/ir_types.js';
import { strandsBuiltinFunctions } from './strands/strands_builtins.js';

class StrandsNode {
  constructor(id, dimension, strandsContext) {
    this.id = id;
    this.strandsContext = strandsContext;
    this.dimension = dimension;
    this.structProperties = null;
    this.isStrandsNode = true;

    // Store original identifier for varying variables
    const dag = this.strandsContext.dag;
    const nodeData = getNodeDataFromID(dag, this.id);
    if (nodeData && nodeData.identifier) {
      this._originalIdentifier = nodeData.identifier;
    }
    if (nodeData) {
      this._originalBaseType = nodeData.baseType;
      this._originalDimension = nodeData.dimension;
    }
  }
  withStructProperties(properties) {
    this.structProperties = properties;
    return this;
  }
  copy() {
    return createStrandsNode(this.id, this.dimension, this.strandsContext);
  }
  typeInfo() {
    return {
      baseType: this._originalBaseType || BaseType.FLOAT,
      dimension: this.dimension
    };
  }
  bridge(value) {
    const { dag, cfg } = this.strandsContext;
    const orig = getNodeDataFromID(dag, this.id);
    const baseType = orig?.baseType ?? BaseType.FLOAT;

    let newValueID;
    if (value?.isStrandsNode) {
      newValueID = value.id;
    } else {
      const newVal = primitiveConstructorNode(
        this.strandsContext,
        { baseType, dimension: this.dimension },
        value
      );
      newValueID = newVal.id;
    }

    // For varying variables, we need both assignment generation AND a way to reference by identifier
    if (this._originalIdentifier) {
      // Create a variable node for the target (the varying variable)
      const { id: targetVarID } = variableNode(
        this.strandsContext,
        { baseType: this._originalBaseType, dimension: this._originalDimension },
        this._originalIdentifier
      );

      // Create assignment node for GLSL generation
      const assignmentNode = createNodeData({
        nodeType: NodeType.ASSIGNMENT,
        dependsOn: [targetVarID, newValueID],
        phiBlocks: []
      });
      const assignmentID = getOrCreateNode(dag, assignmentNode);
      recordInBasicBlock(cfg, cfg.currentBlock, assignmentID);

      // Track for global assignments processing
      this.strandsContext.globalAssignments.push(assignmentID);

      // Simply update this node to be a variable node with the identifier
      // This ensures it always generates the variable name in expressions
      const variableNodeData = createNodeData({
        nodeType: NodeType.VARIABLE,
        baseType: this._originalBaseType,
        dimension: this._originalDimension,
        identifier: this._originalIdentifier
      });
      const variableID = getOrCreateNode(dag, variableNodeData);

      this.id = variableID; // Point to the variable node for expression generation
    } else {
      this.id = newValueID; // For non-varying variables, just update to new value
    }

    return this;
  }
  bridgeSwizzle(swizzlePattern, value) {
    const { dag, cfg } = this.strandsContext;
    const orig = getNodeDataFromID(dag, this.id);
    const baseType = orig?.baseType ?? BaseType.FLOAT;

    let newValueID;
    if (value?.isStrandsNode) {
      newValueID = value.id;
    } else {
      const newVal = primitiveConstructorNode(
        this.strandsContext,
        { baseType, dimension: this.dimension },
        value
      );
      newValueID = newVal.id;
    }

    // For varying variables, create swizzle assignment
    if (this._originalIdentifier) {
      // Create a variable node for the target with swizzle
      const { id: targetVarID } = variableNode(
        this.strandsContext,
        { baseType: this._originalBaseType, dimension: this._originalDimension },
        this._originalIdentifier
      );

      // Create a swizzle node for the target (myVarying.xyz)
      const swizzleNode = createNodeData({
        nodeType: NodeType.OPERATION,
        opCode: OpCode.Unary.SWIZZLE,
        baseType: this._originalBaseType,
        dimension: swizzlePattern.length, // xyz = 3, xy = 2, etc.
        swizzle: swizzlePattern,
        dependsOn: [targetVarID]
      });
      const swizzleID = getOrCreateNode(dag, swizzleNode);

      // Create assignment node: myVarying.xyz = value
      const assignmentNode = createNodeData({
        nodeType: NodeType.ASSIGNMENT,
        dependsOn: [swizzleID, newValueID],
        phiBlocks: []
      });
      const assignmentID = getOrCreateNode(dag, assignmentNode);
      recordInBasicBlock(cfg, cfg.currentBlock, assignmentID);

      // Track for global assignments processing in the current hook context
      this.strandsContext.globalAssignments.push(assignmentID);

      // Simply update this node to be a variable node with the identifier
      // This ensures it always generates the variable name in expressions
      const variableNodeData = createNodeData({
        nodeType: NodeType.VARIABLE,
        baseType: this._originalBaseType,
        dimension: this._originalDimension,
        identifier: this._originalIdentifier
      });
      const variableID = getOrCreateNode(dag, variableNodeData);

      this.id = variableID; // Point to the variable node, not the assignment node
    } else {
      this.id = newValueID; // For non-varying variables, just update to new value
    }

    return this;
  }
  getValue() {
    if (this._originalIdentifier) {
      const { id, dimension } = variableNode(
        this.strandsContext,
        { baseType: this._originalBaseType, dimension: this._originalDimension },
        this._originalIdentifier
      );
      return createStrandsNode(id, dimension, this.strandsContext);
    }

    return this;
  }
}
function createStrandsNode(id, dimension, strandsContext, onRebind) {
  return new Proxy(
    new StrandsNode(id, dimension, strandsContext),
    swizzleTrap(id, dimension, strandsContext, onRebind)
  );
}

//////////////////////////////////////////////
// Builders for node graphs
//////////////////////////////////////////////
function scalarLiteralNode(strandsContext, typeInfo, value) {
  const { cfg, dag } = strandsContext;
  let { dimension, baseType } = typeInfo;
  if (dimension !== 1) {
    internalError('Created a scalar literal node with dimension > 1.');
  }
  const nodeData = createNodeData({
    nodeType: NodeType.LITERAL,
    dimension,
    baseType,
    value
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension };
}

function variableNode(strandsContext, typeInfo, identifier) {
  const { cfg, dag } = strandsContext;
  const { dimension, baseType } = typeInfo;
  const nodeData = createNodeData({
    nodeType: NodeType.VARIABLE,
    dimension,
    baseType,
    identifier
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension };
}

function unaryOpNode(strandsContext, nodeOrValue, opCode) {
  const { dag, cfg } = strandsContext;
  let dependsOn;
  let node;
  if (nodeOrValue?.isStrandsNode) {
    node = nodeOrValue;
  } else {
    const { id, dimension } = primitiveConstructorNode(strandsContext, { baseType: BaseType.FLOAT, dimension: null }, nodeOrValue);
    node = createStrandsNode(id, dimension, strandsContext);
  }
  dependsOn = [node.id];
  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    opCode,
    dependsOn,
    baseType: dag.baseTypes[node.id],
    dimension: node.dimension
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension: node.dimension };
}

function binaryOpNode(strandsContext, leftStrandsNode, rightArg, opCode) {
  const { dag, cfg } = strandsContext;
  // Construct a node for right if its just an array or number etc.
  let rightStrandsNode;
  if (rightArg[0] instanceof StrandsNode && rightArg.length === 1) {
    rightStrandsNode = rightArg[0];
  } else {
    const { id, dimension } = primitiveConstructorNode(strandsContext, { baseType: BaseType.FLOAT, dimension: null }, rightArg);
    rightStrandsNode = createStrandsNode(id, dimension, strandsContext);
  }
  let finalLeftNodeID = leftStrandsNode.id;
  let finalRightNodeID = rightStrandsNode.id;

  // Check if we have to cast either node
  const leftType = extractNodeTypeInfo(dag, leftStrandsNode.id);
  const rightType = extractNodeTypeInfo(dag, rightStrandsNode.id);
  const cast = { node: null, toType: leftType };
  const bothDeferred = leftType.baseType === rightType.baseType && leftType.baseType === BaseType.DEFER;
  if (bothDeferred) {
    cast.toType.baseType = BaseType.FLOAT;
    if (leftType.dimension === rightType.dimension) {
      cast.toType.dimension = leftType.dimension;
    }
    else if (leftType.dimension === 1 && rightType.dimension > 1) {
      cast.toType.dimension = rightType.dimension;
    }
    else if (rightType.dimension === 1 && leftType.dimension > 1) {
      cast.toType.dimension = leftType.dimension;
    }
    else {
      userError("type error", `You have tried to perform a binary operation:\n`+
        `${leftType.baseType+leftType.dimension} ${OpCodeToSymbol[opCode]} ${rightType.baseType+rightType.dimension}\n` +
        `It's only possible to operate on two nodes with the same dimension, or a scalar value and a vector.`
      );
    }
    const l = primitiveConstructorNode(strandsContext, cast.toType, leftStrandsNode);
    const r = primitiveConstructorNode(strandsContext, cast.toType, rightStrandsNode);
    finalLeftNodeID = l.id;
    finalRightNodeID = r.id;
  }
  else if (leftType.baseType !== rightType.baseType ||
    leftType.dimension !== rightType.dimension) {

    if (leftType.dimension === 1 && rightType.dimension > 1) {
      cast.node = leftStrandsNode;
      cast.toType = rightType;
    }
    else if (rightType.dimension === 1 && leftType.dimension > 1) {
      cast.node = rightStrandsNode;
      cast.toType = leftType;
    }
    else if (leftType.priority > rightType.priority) {
      // e.g. op(float vector, int vector): cast priority is float > int > bool
      cast.node = rightStrandsNode;
      cast.toType = leftType;
    }
    else if (rightType.priority > leftType.priority) {
      cast.node = leftStrandsNode;
      cast.toType = rightType;
    }
    else {
      userError('type error', `A vector of length ${leftType.dimension} operated with a vector of length ${rightType.dimension} is not allowed.`);
    }

    const casted = primitiveConstructorNode(strandsContext, cast.toType, cast.node);

    if (cast.node === leftStrandsNode) {
      leftStrandsNode = createStrandsNode(casted.id, casted.dimension, strandsContext);
      finalLeftNodeID = leftStrandsNode.id;
    } else {
      rightStrandsNode = createStrandsNode(casted.id, casted.dimension, strandsContext);
      finalRightNodeID = rightStrandsNode.id;
    }
  }

  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    opCode,
    dependsOn: [finalLeftNodeID, finalRightNodeID],
    baseType: cast.toType.baseType,
    dimension: cast.toType.dimension,
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension: nodeData.dimension };
}

function memberAccessNode(strandsContext, parentNode, componentNode, memberTypeInfo) {
  const { dag, cfg } = strandsContext;
  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    opCode: OpCode.Binary.MEMBER_ACCESS,
    dimension: memberTypeInfo.dimension,
    baseType: memberTypeInfo.baseType,
    dependsOn: [parentNode.id, componentNode.id],
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension: memberTypeInfo.dimension };
}

function structInstanceNode(strandsContext, structTypeInfo, identifier, dependsOn) {
  const { cfg, dag } = strandsContext;
  if (dependsOn.length === 0) {
    for (const prop of structTypeInfo.properties) {
      const typeInfo = prop.dataType;
      const nodeData = createNodeData({
        nodeType: NodeType.VARIABLE,
        baseType: typeInfo.baseType,
        dimension: typeInfo.dimension,
        identifier: `${identifier}.${prop.name}`,
      });
      const componentID = getOrCreateNode(dag, nodeData);
      recordInBasicBlock(cfg, cfg.currentBlock, componentID);
      dependsOn.push(componentID);
    }
  }

  const nodeData = createNodeData({
    nodeType: NodeType.VARIABLE,
    dimension: structTypeInfo.properties.length,
    baseType: structTypeInfo.typeName,
    identifier,
    dependsOn
  });
  const structID = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, structID);

  return { id: structID, dimension: 0, components: dependsOn };
}

function mapPrimitiveDepsToIDs(strandsContext, typeInfo, dependsOn) {
  const inputs = Array.isArray(dependsOn) ? dependsOn : [dependsOn];
  const mappedDependencies = [];
  let { dimension, baseType } = typeInfo;

  const dag = strandsContext.dag;
  let calculatedDimensions = 0;
  let originalNodeID = null;
  for (const dep of inputs.flat(Infinity)) {
    if (dep && dep.isStrandsNode) {
      const node = getNodeDataFromID(dag, dep.id);
      originalNodeID = dep.id;
      baseType = node.baseType;

      if (node.opCode === OpCode.Nary.CONSTRUCTOR) {
        for (const inner of node.dependsOn) {
          mappedDependencies.push(inner);
        }
      } else {
        mappedDependencies.push(dep.id);
      }

      calculatedDimensions += node.dimension;
      continue;
    }
    else if (typeof dep === 'number') {
      const { id, dimension } = scalarLiteralNode(strandsContext, { dimension: 1, baseType }, dep);
      mappedDependencies.push(id);
      calculatedDimensions += dimension;
      continue;
    }
    else {
      userError('type error', `You've tried to construct a scalar or vector type with a non-numeric value: ${dep}`);
    }
  }
  if (dimension === null) {
    dimension = calculatedDimensions;
  } else if (dimension > calculatedDimensions && calculatedDimensions === 1) {
    calculatedDimensions = dimension;
  } else if(calculatedDimensions !== 1 && calculatedDimensions !== dimension) {
    userError('type error', `You've tried to construct a ${baseType + dimension} with ${calculatedDimensions} components`);
  }
  const inferredTypeInfo = {
    dimension,
    baseType,
    priority: BasePriority[baseType],
  };
  return { originalNodeID, mappedDependencies, inferredTypeInfo };
}

function constructTypeFromIDs(strandsContext, typeInfo, strandsNodesArray) {
  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    opCode: OpCode.Nary.CONSTRUCTOR,
    dimension: typeInfo.dimension,
    baseType: typeInfo.baseType,
    dependsOn: strandsNodesArray
  });
  const id = getOrCreateNode(strandsContext.dag, nodeData);
  return id;
}

function primitiveConstructorNode(strandsContext, typeInfo, dependsOn) {
  const cfg = strandsContext.cfg;
  dependsOn = (Array.isArray(dependsOn) ? dependsOn : [dependsOn])
    .flat(Infinity)
    .map(a => {
      if (
        a.isStrandsNode &&
        a.typeInfo().baseType === BaseType.INT &&
        // TODO: handle ivec inputs instead of just int scalars
        a.typeInfo().dimension === 1
      ) {
        return castToFloat(strandsContext, a);
      } else {
        return a;
      }
    });
  const { mappedDependencies, inferredTypeInfo } = mapPrimitiveDepsToIDs(strandsContext, typeInfo, dependsOn);

  const finalType = {
    baseType: typeInfo.baseType,
    dimension: inferredTypeInfo.dimension
  };

  const id = constructTypeFromIDs(strandsContext, finalType, mappedDependencies);
  if (typeInfo.baseType !== BaseType.DEFER) {
    recordInBasicBlock(cfg, cfg.currentBlock, id);
  }

  return { id, dimension: finalType.dimension, components: mappedDependencies };
}

function castToFloat(strandsContext, dep) {
  const { id, dimension } = functionCallNode(
    strandsContext,
    strandsContext.backend.getTypeName('float', dep.typeInfo().dimension),
    [dep],
    {
      overloads: [{
        params: [dep.typeInfo()],
        returnType: {
          ...dep.typeInfo(),
          baseType: BaseType.FLOAT,
        },
      }],
    }
  );
  return createStrandsNode(id, dimension, strandsContext);
}

function structConstructorNode(strandsContext, structTypeInfo, rawUserArgs) {
  const { cfg, dag } = strandsContext;
  const { identifer, properties } = structTypeInfo;

  if (!(rawUserArgs.length === properties.length)) {
    userError('type error',
      `You've tried to construct a ${structTypeInfo.typeName} struct with ${rawUserArgs.length} properties, but it expects ${properties.length} properties.\n` +
      `The properties it expects are:\n` +
      `${properties.map(prop => prop.name + ' ' + prop.DataType.baseType + prop.DataType.dimension)}`
    );
  }

  const dependsOn = [];
  for (let i = 0; i < properties.length; i++) {
    const expectedProperty = properties[i];
    const { originalNodeID, mappedDependencies } = mapPrimitiveDepsToIDs(strandsContext, expectedProperty.dataType, rawUserArgs[i]);
    if (originalNodeID) {
      dependsOn.push(originalNodeID);
    }
    else {
      dependsOn.push(
        constructTypeFromIDs(strandsContext, expectedProperty.dataType, mappedDependencies)
      );
    }
  }

  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    opCode: OpCode.Nary.CONSTRUCTOR,
    dimension: properties.length,
    baseType: structTypeInfo.typeName ,
    dependsOn
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension: properties.length, components: structTypeInfo.components };
}

function functionCallNode(
  strandsContext,
  functionName,
  rawUserArgs,
  { overloads: rawOverloads } = {},
) {
  const { cfg, dag } = strandsContext;
  const overloads = rawOverloads || strandsBuiltinFunctions[functionName];

  const preprocessedArgs = rawUserArgs.map((rawUserArg) => mapPrimitiveDepsToIDs(strandsContext, DataType.defer, rawUserArg));
  const matchingArgsCounts = overloads.filter(overload => overload.params.length === preprocessedArgs.length);
  if (matchingArgsCounts.length === 0) {
    const argsLengthSet = new Set();
    const argsLengthArr = [];
    overloads.forEach((overload) => argsLengthSet.add(overload.params.length));
    argsLengthSet.forEach((len) => argsLengthArr.push(`${len}`));
    const argsLengthStr = argsLengthArr.join(', or ');
    userError("parameter validation error",`Function '${functionName}' has ${overloads.length} variants which expect ${argsLengthStr} arguments, but ${preprocessedArgs.length} arguments were provided.`);
  }

  const isGeneric = (T) => T.dimension === null;
  let bestOverload = null;
  let bestScore = 0;
  let inferredReturnType = null;
  let inferredDimension = null;

  for (const overload of matchingArgsCounts) {
    let isValid = true;
    let similarity = 0;

    for (let i = 0; i < preprocessedArgs.length; i++) {
      const preArg = preprocessedArgs[i];
      const argType = preArg.inferredTypeInfo;
      const expectedType = overload.params[i];
      let dimension = expectedType.dimension;

      if (isGeneric(expectedType)) {
        if (inferredDimension === null || inferredDimension === 1) {
          inferredDimension = argType.dimension;
        }

        if (inferredDimension !== argType.dimension &&
          !(argType.dimension === 1 && inferredDimension >= 1)
          ) {
          isValid = false;
        }
        dimension = inferredDimension;
      }
      else {
        if (argType.dimension > dimension) {
          isValid = false;
        }
      }

      if (argType.baseType === expectedType.baseType) {
        similarity += 2;
      }
      else if(expectedType.priority > argType.priority) {
        similarity += 1;
      }

    }

    if (isValid && (!bestOverload || similarity > bestScore)) {
      bestOverload = overload;
      bestScore = similarity;
      inferredReturnType =  {...overload.returnType };
      if (isGeneric(inferredReturnType)) {
        inferredReturnType.dimension = inferredDimension;
      }
    }
  }

  if (bestOverload === null) {
    userError('parameter validation', `No matching overload for ${functionName} was found!`);
  }

  let dependsOn = [];
  for (let i = 0; i < bestOverload.params.length; i++) {
    const arg = preprocessedArgs[i];
    const paramType = { ...bestOverload.params[i] };
    if (isGeneric(paramType)) {
      paramType.dimension = inferredDimension;
    }
    if (arg.originalNodeID && typeEquals(arg.inferredTypeInfo, paramType)) {
      dependsOn.push(arg.originalNodeID);
    }
    else {
      const castedArgID = constructTypeFromIDs(strandsContext, paramType, arg.mappedDependencies);
      recordInBasicBlock(cfg, cfg.currentBlock, castedArgID);
      dependsOn.push(castedArgID);
    }
  }

  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    opCode: OpCode.Nary.FUNCTION_CALL,
    identifier: functionName,
    dependsOn,
    baseType: inferredReturnType.baseType,
    dimension: inferredReturnType.dimension
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension: inferredReturnType.dimension  };
}

function statementNode(strandsContext, statementType) {
  const { dag, cfg } = strandsContext;
  const nodeData = createNodeData({
    nodeType: NodeType.STATEMENT,
    statementType
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return id;
}

function swizzleNode(strandsContext, parentNode, swizzle) {
  const { dag, cfg } = strandsContext;
  const baseType = dag.baseTypes[parentNode.id];
  const nodeData = createNodeData({
    nodeType: NodeType.OPERATION,
    baseType,
    dimension: swizzle.length,
    opCode: OpCode.Unary.SWIZZLE,
    dependsOn: [parentNode.id],
    swizzle,
  });
  const id = getOrCreateNode(dag, nodeData);
  recordInBasicBlock(cfg, cfg.currentBlock, id);
  return { id, dimension: swizzle.length };
}

function swizzleTrap(id, dimension, strandsContext, onRebind) {
    const swizzleSets = [
      ['x', 'y', 'z', 'w'],
      ['r', 'g', 'b', 'a'],
      ['s', 't', 'p', 'q']
    ].map(s => s.slice(0, dimension));
    const trap = {
      get(target, property, receiver) {
        if (property in target) {
          return Reflect.get(...arguments);
        } else {
          for (const set of swizzleSets) {
            if ([...property.toString()].every(char => set.includes(char))) {
              const swizzle = [...property].map(char => {
                const index = set.indexOf(char);
                return swizzleSets[0][index];
              }).join('');
              const node = swizzleNode(strandsContext, target, swizzle);
              return createStrandsNode(node.id, node.dimension, strandsContext);
            }
          }
        }
    },
  set(target, property, value, receiver) {
    for (const swizzleSet of swizzleSets) {
      const chars = [...property];
      const valid =
        chars.every(c => swizzleSet.includes(c)) &&
        new Set(chars).size === chars.length &&
        target.dimension >= chars.length;
      if (!valid) continue;

      const dim = target.dimension;

      // lanes are the underlying values of the target vector
      //  e.g. lane 0 holds the value aliased by 'x', 'r', and 's'
      // the lanes array is in the 'correct' order
      const lanes = new Array(dim);
      for (let i = 0; i < dim; i++) {
        const { id, dimension } = swizzleNode(strandsContext, target, 'xyzw'[i]);
        lanes[i] = createStrandsNode(id, dimension, strandsContext);
      }

      // The scalars array contains the individual components of the users values.
      // This may not be the most efficient way, as we swizzle each component individually,
      // so that .xyz becomes .x, .y, .z
      let scalars = [];
      if (value?.isStrandsNode) {
        if (value.dimension === 1) {
          scalars = Array(chars.length).fill(value);
        } else if (value.dimension === chars.length) {
          for (let k = 0; k < chars.length; k++) {
            const { id, dimension } = swizzleNode(strandsContext, value, 'xyzw'[k]);
            scalars.push(createStrandsNode(id, dimension, strandsContext));
          }
        } else {
          userError('type error', `Swizzle assignment: RHS vector does not match LHS vector (need ${chars.length}, got ${value.dimension}).`);
        }
      } else if (Array.isArray(value)) {
        const flat = value.flat(Infinity);
        if (flat.length === 1) {
          scalars = Array(chars.length).fill(flat[0]);
        } else if (flat.length === chars.length) {
          scalars = flat;
        } else {
          userError('type error', `Swizzle assignment: RHS length ${flat.length} does not match ${chars.length}.`);
        }
      } else if (typeof value === 'number') {
        scalars = Array(chars.length).fill(value);
      } else {
        userError('type error', `Unsupported RHS for swizzle assignment: ${value}`);
      }

      // The canonical index refers to the actual value's position in the vector lanes
      // i.e. we are finding (3,2,1) from .zyx
      // We set the correct value in the lanes array
      for (let j = 0; j < chars.length; j++) {
        const canonicalIndex = swizzleSet.indexOf(chars[j]);
        lanes[canonicalIndex] = scalars[j];
      }

      const orig = getNodeDataFromID(strandsContext.dag, target.id);
      const baseType = orig?.baseType ?? BaseType.FLOAT;
      const { id: newID } = primitiveConstructorNode(
        strandsContext,
        { baseType, dimension: dim },
        lanes
      );

      target.id = newID;

      // If we swizzle assign on a struct component i.e.
      //   inputs.position.rg = [1, 2]
      // The onRebind callback will update the structs components so that it refers to the new values,
      // and make a new ID for the struct with these new values
      if (typeof onRebind === 'function') {
        onRebind(newID);
      }
      return true;
    }
    return Reflect.set(...arguments);
  }
  };
  return trap;
}

export { StrandsNode as S, structInstanceNode as a, structConstructorNode as b, createStrandsNode as c, binaryOpNode as d, scalarLiteralNode as e, functionCallNode as f, constructTypeFromIDs as g, castToFloat as h, swizzleNode as i, swizzleTrap as j, memberAccessNode as m, primitiveConstructorNode as p, statementNode as s, unaryOpNode as u, variableNode as v };
