aethex-studio/src/components/visual-scripting/nodes/CustomNode.tsx

'use client';

import { memo, useState } from 'react';
import { Handle, Position, NodeProps } from 'reactflow';
import { Input } from '@/components/ui/input';
import { Switch } from '@/components/ui/switch';
import { Label } from '@/components/ui/label';
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from '@/components/ui/select';
import { Trash2 } from 'lucide-react';
import { Button } from '@/components/ui/button';

import {
  getNodeDefinition,
  CATEGORY_COLORS,
  PORT_COLORS,
  NodePort,
} from '@/lib/visual-scripting/node-definitions';
import { NodeData } from '@/lib/visual-scripting/code-generator';
import { useVisualScriptStore } from '@/stores/visual-script-store';

export const CustomNode = memo(({ id, data, selected }: NodeProps<NodeData>) => {
  const definition = getNodeDefinition(data.type);
  const { updateNodeValue, removeNode } = useVisualScriptStore();
  const [isHovered, setIsHovered] = useState(false);

  if (!definition) {
    return (
      <div className="bg-destructive/20 border border-destructive rounded p-2">
        <span className="text-xs text-destructive">Unknown node: {data.type}</span>
      </div>
    );
  }

  const handleValueChange = (key: string, value: any) => {
    updateNodeValue(id, key, value);
  };

  // Filter out flow inputs for the input section
  const editableInputs = definition.inputs.filter(
    (input) => input.type !== 'flow'
  );

  return (
    <div
      className={`bg-card rounded-lg border-2 shadow-lg min-w-[180px] max-w-[280px] transition-all ${
        selected ? 'border-primary ring-2 ring-primary/20' : 'border-border'
      }`}
      style={{
        borderTopColor: definition.color,
        borderTopWidth: 3,
      }}
      onMouseEnter={() => setIsHovered(true)}
      onMouseLeave={() => setIsHovered(false)}
    >
      {/* Header */}
      <div
        className="px-3 py-2 rounded-t-md flex items-center justify-between"
        style={{ backgroundColor: `${definition.color}15` }}
      >
        <div className="flex items-center gap-2">
          <div
            className="w-5 h-5 rounded flex items-center justify-center text-white text-xs"
            style={{ backgroundColor: definition.color }}
          >
            {definition.category.charAt(0).toUpperCase()}
          </div>
          <span className="font-medium text-sm">{definition.label}</span>
        </div>
        {isHovered && (
          <Button
            variant="ghost"
            size="icon"
            className="h-5 w-5 text-muted-foreground hover:text-destructive"
            onClick={() => removeNode(id)}
          >
            <Trash2 className="h-3 w-3" />
          </Button>
        )}
      </div>

      {/* Input Handles (Left Side) */}
      <div className="relative">
        {definition.inputs.map((input, index) => (
          <div key={input.id} className="relative">
            <Handle
              type="target"
              position={Position.Left}
              id={input.id}
              className="!w-3 !h-3 !border-2 !border-background"
              style={{
                backgroundColor: PORT_COLORS[input.type],
                top: `${
                  (editableInputs.length > 0 ? 45 : 20) +
                  index * 28 +
                  (input.type === 'flow' ? 0 : editableInputs.indexOf(input) * 36)
                }px`,
              }}
              isConnectable={true}
            />
            {input.type === 'flow' && (
              <div
                className="absolute left-4 text-xs text-muted-foreground"
                style={{
                  top: `${20 + index * 28}px`,
                }}
              >
                {input.name}
              </div>
            )}
          </div>
        ))}
      </div>

      {/* Editable Inputs */}
      {editableInputs.length > 0 && (
        <div className="px-3 py-2 space-y-2 border-t border-border/50">
          {editableInputs.map((input) => (
            <NodeInput
              key={input.id}
              input={input}
              value={data.values?.[input.id] ?? input.defaultValue}
              onChange={(value) => handleValueChange(input.id, value)}
            />
          ))}
        </div>
      )}

      {/* Output Handles (Right Side) */}
      <div className="relative min-h-[20px]">
        {definition.outputs.map((output, index) => (
          <div key={output.id} className="relative">
            <Handle
              type="source"
              position={Position.Right}
              id={output.id}
              className="!w-3 !h-3 !border-2 !border-background"
              style={{
                backgroundColor: PORT_COLORS[output.type],
                top: `${10 + index * 24}px`,
              }}
              isConnectable={true}
            />
            <div
              className="absolute right-4 text-xs text-muted-foreground text-right"
              style={{
                top: `${6 + index * 24}px`,
              }}
            >
              {output.name}
            </div>
          </div>
        ))}
      </div>

      {/* Spacer for outputs */}
      <div
        style={{
          height: Math.max(definition.outputs.length * 24, 20),
        }}
      />
    </div>
  );
});

CustomNode.displayName = 'CustomNode';

// Input field component for node properties
function NodeInput({
  input,
  value,
  onChange,
}: {
  input: NodePort;
  value: any;
  onChange: (value: any) => void;
}) {
  switch (input.type) {
    case 'number':
      return (
        <div className="space-y-1">
          <Label className="text-xs text-muted-foreground">{input.name}</Label>
          <Input
            type="number"
            value={value ?? ''}
            onChange={(e) => onChange(parseFloat(e.target.value) || 0)}
            className="h-7 text-xs"
          />
        </div>
      );

    case 'string':
      if (input.id === 'operation') {
        return (
          <div className="space-y-1">
            <Label className="text-xs text-muted-foreground">{input.name}</Label>
            <Select value={value || 'add'} onValueChange={onChange}>
              <SelectTrigger className="h-7 text-xs">
                <SelectValue />
              </SelectTrigger>
              <SelectContent>
                <SelectItem value="add">+ Add</SelectItem>
                <SelectItem value="subtract">- Subtract</SelectItem>
                <SelectItem value="multiply">× Multiply</SelectItem>
                <SelectItem value="divide">÷ Divide</SelectItem>
                <SelectItem value="modulo">% Modulo</SelectItem>
                <SelectItem value="power">^ Power</SelectItem>
              </SelectContent>
            </Select>
          </div>
        );
      }
      if (input.id === 'comparison') {
        return (
          <div className="space-y-1">
            <Label className="text-xs text-muted-foreground">{input.name}</Label>
            <Select value={value || 'equals'} onValueChange={onChange}>
              <SelectTrigger className="h-7 text-xs">
                <SelectValue />
              </SelectTrigger>
              <SelectContent>
                <SelectItem value="equals">= Equals</SelectItem>
                <SelectItem value="notEquals">≠ Not Equals</SelectItem>
                <SelectItem value="greater">&gt; Greater</SelectItem>
                <SelectItem value="less">&lt; Less</SelectItem>
                <SelectItem value="greaterEqual">≥ Greater or Equal</SelectItem>
                <SelectItem value="lessEqual">≤ Less or Equal</SelectItem>
              </SelectContent>
            </Select>
          </div>
        );
      }
      if (input.id === 'key') {
        return (
          <div className="space-y-1">
            <Label className="text-xs text-muted-foreground">{input.name}</Label>
            <Select value={value || 'E'} onValueChange={onChange}>
              <SelectTrigger className="h-7 text-xs">
                <SelectValue />
              </SelectTrigger>
              <SelectContent>
                {['E', 'F', 'G', 'Q', 'R', 'Space', 'Shift', 'One', 'Two', 'Three'].map(
                  (key) => (
                    <SelectItem key={key} value={key}>
                      {key}
                    </SelectItem>
                  )
                )}
              </SelectContent>
            </Select>
          </div>
        );
      }
      if (input.id === 'service') {
        return (
          <div className="space-y-1">
            <Label className="text-xs text-muted-foreground">{input.name}</Label>
            <Select value={value || 'Players'} onValueChange={onChange}>
              <SelectTrigger className="h-7 text-xs">
                <SelectValue />
              </SelectTrigger>
              <SelectContent>
                {[
                  'Players',
                  'ReplicatedStorage',
                  'ServerStorage',
                  'Lighting',
                  'TweenService',
                  'UserInputService',
                  'RunService',
                  'SoundService',
                  'DataStoreService',
                ].map((svc) => (
                  <SelectItem key={svc} value={svc}>
                    {svc}
                  </SelectItem>
                ))}
              </SelectContent>
            </Select>
          </div>
        );
      }
      return (
        <div className="space-y-1">
          <Label className="text-xs text-muted-foreground">{input.name}</Label>
          <Input
            type="text"
            value={value ?? ''}
            onChange={(e) => onChange(e.target.value)}
            className="h-7 text-xs"
            placeholder={input.name}
          />
        </div>
      );

    case 'boolean':
      return (
        <div className="flex items-center justify-between">
          <Label className="text-xs text-muted-foreground">{input.name}</Label>
          <Switch
            checked={value ?? input.defaultValue ?? false}
            onCheckedChange={onChange}
          />
        </div>
      );

    default:
      return (
        <div className="space-y-1">
          <Label className="text-xs text-muted-foreground">{input.name}</Label>
          <Input
            type="text"
            value={value ?? ''}
            onChange={(e) => onChange(e.target.value)}
            className="h-7 text-xs"
            placeholder={`${input.name} (${input.type})`}
          />
        </div>
      );
  }
}