gpustack-ui/src/pages/model-routes/forms/targets.tsx

import MetadataList from '@/components/metadata-list';
import SealCascader from '@/components/seal-form/seal-cascader';
import SealInput from '@/components/seal-form/seal-input';
import useAppUtils from '@/hooks/use-app-utils';
import { useIntl } from '@umijs/max';
import { Form } from 'antd';
import _ from 'lodash';
import { forwardRef, useEffect, useImperativeHandle, useState } from 'react';
import { FormData } from '../config/types';
import useTargetSourceModels from '../hooks/use-target-source-models';

const TargetsForm = forwardRef((props, ref) => {
  const intl = useIntl();
  const { getRuleMessage } = useAppUtils();
  const { sourceModels, fetchSourceModels } = useTargetSourceModels();
  const [validTriggered, setValidTriggered] = useState<boolean>(false);
  const form = Form.useFormInstance<FormData>();
  const targets = Form.useWatch('targets', form) || [];
  const [fallbackValues, setFallbackValues] = useState<{ value: any[] }>({
    value: []
  });
  const [dataList, setDataList] = useState<
    {
      weight: number | null;
      value: any[];
    }[]
  >([]);

  useImperativeHandle(ref, () => ({
    initFallbackValues: (values: { value: any[] }) => {
      setFallbackValues(values);
    },
    initDataList: (
      list: {
        weight: number | null;
        value: any[];
      }[]
    ) => {
      setDataList(list);
    }
  }));

  const handleTargetsChange = (value: any[], index: number, options: any[]) => {
    const selectedOption =
      options?.find?.((opt) => opt.value === value[1]) || {};
    const targetList = [...targets];
    targetList[index] = {
      weight: targetList[index]?.weight,
      ...selectedOption?.data
    };

    form.setFieldValue('targets', [...targetList]);

    const newDataList = [...dataList];
    newDataList[index] = {
      weight: newDataList[index]?.weight || null,
      value: value
    };
    form.validateFields(['targets']);
    setDataList(newDataList);
  };

  const handleOnAdd = () => {
    const newDataList = [
      ...dataList,
      {
        weight: 100,
        value: []
      }
    ];
    setDataList(newDataList);
    const newTargets = [
      ...targets,
      {
        weight: 100,
        value: []
      }
    ];
    form.setFieldValue('targets', newTargets);
  };

  const handleOnDelete = (index: number, item: any) => {
    const newDataList = dataList.filter((_, i) => i !== index);
    setDataList(newDataList);

    const targetList = [...targets];
    targetList.splice(index, 1);
    form.setFieldValue('targets', [...targetList]);
  };

  const handleFallbackChange = (value: any[], options?: any[]) => {
    if (!value || value.length === 0) {
      form.setFieldValue('fallback_target', null);
      setFallbackValues({
        value: []
      });
      return;
    }
    const selectedOption =
      options?.find?.((opt) => opt.value === value[1]) || {};

    form.setFieldValue('fallback_target', {
      ...selectedOption?.data
    });
    setFallbackValues({
      value: value
    });
  };

  const handleOnWeightChange = (value: any, index: number) => {
    const targetList = [...targets];
    if (targetList[index]) {
      targetList[index] = {
        ...targetList[index],
        weight: value
      };
      form.setFieldValue('targets', [...targetList]);
    }

    const newDataList = [...dataList];
    newDataList[index] = {
      ...newDataList[index],
      weight: value
    };
    form.validateFields(['targets']);
    setDataList(newDataList);
  };

  const buildKey = (path?: any[]) =>
    Array.isArray(path) ? path.join('/') : '';

  const filterOptions = (currentValue: any[]) => {
    const currKey = buildKey(currentValue);

    const selectedDataList = [...dataList, { value: fallbackValues.value }];

    const selectedKeys = new Set(
      selectedDataList
        .filter((item) => item.value)
        .map((item) => buildKey(item.value))
    );

    return sourceModels
      .map((model) => {
        const children = model.children?.filter((child) => {
          const key = buildKey([child.data?.parentId, child.value]);

          return !selectedKeys.has(key) || key === currKey;
        });
        console.log('children', children);

        return {
          ...model,
          children
        };
      })
      .filter((model) => model.children && model.children.length > 0);
  };

  const displayRender = (labels: any[]) => {
    return (
      <span className="flex-center gap-4">
        {labels[0]}
        <span>/</span>
        <span>{labels[1]}</span>
      </span>
    );
  };

  useEffect(() => {
    fetchSourceModels();
  }, []);

  return (
    <>
      <Form.Item
        name="targets"
        data-field="targets"
        trigger=""
        rules={[
          {
            validator(rule, value) {
              if (value && value?.length > 0) {
                if (_.some(value, (item: any) => !item.weight)) {
                  setValidTriggered(true);
                  return Promise.reject(
                    getRuleMessage('input', 'routes.form.target.weight')
                  );
                }

                if (
                  _.some(
                    dataList,
                    (item: any) => !item.value || item.value.length === 0
                  )
                ) {
                  setValidTriggered(true);
                  return Promise.reject(
                    getRuleMessage('input', 'providers.form.target.placeholder')
                  );
                }
              }
              setValidTriggered(false);
              return Promise.resolve();
            }
          }
        ]}
      >
        <MetadataList
          label={''}
          styles={{
            wrapper: {
              paddingTop: 14
            }
          }}
          dataList={dataList}
          btnText={intl.formatMessage({ id: 'routes.form.target.add' })}
          onAdd={handleOnAdd}
          onDelete={handleOnDelete}
        >
          {(item, index) => (
            <>
              <SealCascader
                required
                showSearch
                status={
                  (!item.value || item.value.length === 0) && validTriggered
                    ? 'error'
                    : 'success'
                }
                expandTrigger="hover"
                multiple={false}
                alwaysFocus={true}
                onChange={(value, options) =>
                  handleTargetsChange(value, index, options)
                }
                classNames={{
                  popup: {
                    root: 'cascader-popup-wrapper gpu-selector'
                  }
                }}
                styles={{
                  popup: {
                    listItem: {
                      padding: '5px 10px'
                    }
                  }
                }}
                maxTagCount={1}
                placeholder={intl.formatMessage({
                  id: 'providers.form.target.placeholder'
                })}
                value={item.value}
                options={filterOptions(item.value)}
                showCheckedStrategy="SHOW_CHILD"
                displayRender={displayRender}
                getPopupContainer={(triggerNode) => triggerNode.parentNode}
              ></SealCascader>
              <span className="seprator">:</span>
              <SealInput.Number
                style={{ flex: 100 }}
                min={0}
                step={1}
                status={
                  item.weight === null && validTriggered ? 'error' : 'success'
                }
                value={item.weight}
                onChange={(value) => handleOnWeightChange(value, index)}
                placeholder={intl.formatMessage({
                  id: 'routes.form.target.weight'
                })}
              ></SealInput.Number>
            </>
          )}
        </MetadataList>
      </Form.Item>
      <Form.Item name="fallback_target">
        <div>
          <SealCascader
            showSearch
            expandTrigger="hover"
            multiple={false}
            alwaysFocus={true}
            classNames={{
              popup: {
                root: 'cascader-popup-wrapper gpu-selector'
              }
            }}
            styles={{
              popup: {
                listItem: {
                  padding: '5px 10px'
                }
              }
            }}
            label={intl.formatMessage({
              id: 'routes.form.target.fallback'
            })}
            placeholder={intl.formatMessage({
              id: 'providers.form.target.placeholder'
            })}
            maxTagCount={1}
            value={fallbackValues.value}
            options={filterOptions(fallbackValues.value)}
            onChange={(value, options) => handleFallbackChange(value, options)}
            showCheckedStrategy="SHOW_CHILD"
            displayRender={displayRender}
            getPopupContainer={(triggerNode) => triggerNode.parentNode}
          ></SealCascader>
        </div>
      </Form.Item>
    </>
  );
});

export default TargetsForm;