DataMate/frontend/src/pages/DataAnnotation/Create/CreateTask.tsx

import { useEffect, useMemo, useState } from "react";
import { Button, Input, Select, Form, message, Radio } from "antd";
import type { RadioChangeEvent } from "antd";
import TextArea from "antd/es/input/TextArea";
import { DatabaseOutlined } from "@ant-design/icons";
import { Link, useNavigate } from "react-router";
import { ArrowLeft } from "lucide-react";
import { queryDatasetsUsingGet } from "../../DataManagement/dataset.api";
import { mapDataset } from "@/pages/DataManagement/dataset.const";
import { DatasetType, type Dataset } from "@/pages/DataManagement/dataset.model";
import {
  createAnnotationTaskUsingPost,
  queryAnnotationTemplatesUsingGet,
} from "../annotation.api";
import { DataType, type AnnotationTemplate } from "../annotation.model";
import TemplateConfigurationTreeEditor from "../components/TemplateConfigurationTreeEditor";

const DEFAULT_SEGMENTATION_ENABLED = true;
const SEGMENTATION_OPTIONS = [
  { label: "需要切片段", value: true },
  { label: "不需要切片段", value: false },
];
const resolveTemplateDataType = (datasetType?: DatasetType) => {
  switch (datasetType) {
    case DatasetType.TEXT:
      return DataType.TEXT;
    case DatasetType.IMAGE:
      return DataType.IMAGE;
    case DatasetType.AUDIO:
      return DataType.AUDIO;
    case DatasetType.VIDEO:
      return DataType.VIDEO;
    default:
      return undefined;
  }
};
const resolveTemplateTimestamp = (template: AnnotationTemplate) => {
  const timestamp = template.updatedAt || template.createdAt;
  const parsed = Date.parse(timestamp);
  return Number.isNaN(parsed) ? 0 : parsed;
};
const resolveDefaultTemplate = (items: AnnotationTemplate[]) =>
  items.reduce<AnnotationTemplate | undefined>((latest, current) => {
    if (!latest) {
      return current;
    }
    return resolveTemplateTimestamp(current) > resolveTemplateTimestamp(latest)
      ? current
      : latest;
  }, undefined);

export default function AnnotationTaskCreate() {
  const navigate = useNavigate();
  const [form] = Form.useForm();
  const [datasets, setDatasets] = useState<Dataset[]>([]);
  const [templates, setTemplates] = useState<AnnotationTemplate[]>([]);
  const [selectedDatasetId, setSelectedDatasetId] = useState<string | null>(null);
  const [labelConfig, setLabelConfig] = useState("");
  const [configMode, setConfigMode] = useState<"template" | "custom">("template");
  const [submitting, setSubmitting] = useState(false);

  const selectedDataset = useMemo(
    () => datasets.find((dataset) => dataset.id === selectedDatasetId),
    [datasets, selectedDatasetId]
  );
  const isTextDataset = selectedDataset?.datasetType === DatasetType.TEXT;

  const fetchDatasets = async () => {
    try {
      const { data } = await queryDatasetsUsingGet({ page: 0, pageSize: 1000 });
      const list = data?.content || [];
      setDatasets(list.map((item) => mapDataset(item)) || []);
    } catch (error) {
      console.error("加载数据集失败:", error);
      message.error("加载数据集失败");
    }
  };

  const fetchTemplates = async (dataType?: string) => {
    if (!dataType) {
      setTemplates([]);
      return;
    }
    try {
      const response = await queryAnnotationTemplatesUsingGet({
        page: 1,
        size: 200,
        dataType,
      });
      if (response.code === 200 && response.data) {
        setTemplates(response.data.content || []);
      } else {
        message.error(response.message || "加载模板失败");
      }
    } catch (error) {
      console.error("加载模板失败:", error);
      message.error("加载模板失败");
    }
  };

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

  useEffect(() => {
    if (!selectedDataset) {
      setTemplates([]);
      form.setFieldsValue({ templateId: undefined });
      setLabelConfig("");
      return;
    }
    const dataType = resolveTemplateDataType(selectedDataset.datasetType);
    fetchTemplates(dataType);
  }, [form, selectedDataset]);

  useEffect(() => {
    if (configMode !== "template" || !selectedDataset) {
      return;
    }
    if (templates.length === 0) {
      form.setFieldsValue({ templateId: undefined });
      setLabelConfig("");
      return;
    }
    const currentTemplateId = form.getFieldValue("templateId");
    const currentTemplate = templates.find((template) => template.id === currentTemplateId);
    if (currentTemplate) {
      return;
    }
    const defaultTemplate = resolveDefaultTemplate(templates);
    if (defaultTemplate) {
      form.setFieldsValue({ templateId: defaultTemplate.id });
      setLabelConfig(defaultTemplate.labelConfig || "");
    }
  }, [configMode, form, selectedDataset, templates]);
  const handleTemplateSelect = (value?: string) => {
    if (!value) {
      setLabelConfig("");
      return;
    }
    const selectedTemplate = templates.find((template) => template.id === value);
    setLabelConfig(selectedTemplate?.labelConfig || "");
  };

  const handleConfigModeChange = (e: RadioChangeEvent) => {
    const mode = e.target.value;
    setConfigMode(mode);
    if (mode === "custom") {
      form.setFieldsValue({ templateId: undefined });
    }
  };

  const handleSubmit = async () => {
    try {
      const values = await form.validateFields();
      if (!labelConfig.trim()) {
        message.error("请配置标注模板");
        return;
      }

      setSubmitting(true);
      const requestData: Record<string, unknown> = {
        name: values.name,
        description: values.description,
        datasetId: values.datasetId,
        templateId: configMode === "template" ? values.templateId : undefined,
        labelConfig: labelConfig.trim(),
      };
      if (isTextDataset) {
        requestData.segmentationEnabled =
          values.segmentationEnabled ?? DEFAULT_SEGMENTATION_ENABLED;
      }
      await createAnnotationTaskUsingPost(requestData);
      message.success("标注任务创建成功");
      navigate("/data/annotation");
    } catch (error: unknown) {
      const err = error as { errorFields?: unknown; message?: string; data?: { message?: string } };
      if (err?.errorFields) {
        message.error("请完善必填信息");
      } else {
        const msg = err?.message || err?.data?.message || "创建失败，请稍后重试";
        message.error(msg);
        console.error(error);
      }
    } finally {
      setSubmitting(false);
    }
  };

  return (
    <div className="flex-overflow-auto">
      <div className="flex items-center mb-2">
        <Link to="/data/annotation">
          <Button type="text">
            <ArrowLeft className="w-4 h-4 mr-1" />
          </Button>
        </Link>
        <h1 className="text-xl font-bold bg-clip-text">创建标注任务</h1>
      </div>

      <div className="flex-overflow-auto bg-white rounded-lg shadow-sm">
        <div className="flex-1 overflow-y-auto p-6">
          <Form form={form} layout="vertical">
            <h2 className="font-medium text-gray-900 text-lg mb-2">基本信息</h2>
            <Form.Item
              label="任务名称"
              name="name"
              rules={[{ required: true, message: "请输入任务名称" }]}
            >
              <Input placeholder="输入任务名称" />
            </Form.Item>
            <Form.Item label="任务描述" name="description">
              <TextArea placeholder="详细描述标注任务的要求和目标" rows={3} />
            </Form.Item>
            <Form.Item
              label="选择数据集"
              name="datasetId"
              rules={[{ required: true, message: "请选择数据集" }]}
            >
              <Select
                optionFilterProp="children"
                placeholder="请选择数据集"
                size="large"
                options={datasets.map((dataset) => ({
                  label: (
                    <div className="flex items-center justify-between gap-3 py-2">
                      <div className="font-medium text-gray-900">
                        {dataset?.icon || <DatabaseOutlined className="mr-2" />}
                        {dataset.name}
                      </div>
                      <div className="text-xs text-gray-500">
                        {dataset?.fileCount} 文件 • {dataset.size}
                      </div>
                    </div>
                  ),
                  value: dataset.id,
                }))}
                onChange={(value) => {
                  setSelectedDatasetId(value);
                  form.setFieldsValue({ templateId: undefined });
                  setLabelConfig("");
                  const dataset = datasets.find((item) => item.id === value);
                  if (dataset?.datasetType === DatasetType.TEXT) {
                    const currentValue = form.getFieldValue("segmentationEnabled");
                    if (currentValue === undefined) {
                      form.setFieldsValue({
                        segmentationEnabled: DEFAULT_SEGMENTATION_ENABLED,
                      });
                    }
                  } else if (dataset) {
                    form.setFieldsValue({ segmentationEnabled: false });
                  }
                }}
              />
            </Form.Item>

            {selectedDatasetId && isTextDataset && (
              <Form.Item
                label="段落切片"
                name="segmentationEnabled"
                initialValue={DEFAULT_SEGMENTATION_ENABLED}
                extra="仅文本数据集可配置该项"
              >
                <Radio.Group
                  options={SEGMENTATION_OPTIONS}
                  optionType="button"
                  buttonStyle="solid"
                />
              </Form.Item>
            )}

            <div className="flex items-center justify-between mt-6 mb-2">
              <h2 className="font-medium text-gray-900 text-lg flex items-center gap-2">
                标注配置
              </h2>
              <Radio.Group value={configMode} onChange={handleConfigModeChange} buttonStyle="solid">
                <Radio.Button value="template">选择现有模板</Radio.Button>
                {/*<Radio.Button value="custom">自定义配置</Radio.Button>*/}
              </Radio.Group>
            </div>

            {configMode === "template" && (
              <Form.Item
                label="加载现有模板"
                name="templateId"
                rules={[{ required: true, message: "请选择标注模板" }]}
              >
                <Select
                  placeholder="选择一个模板作为基础"
                  showSearch
                  allowClear
                  optionFilterProp="label"
                  options={templates.map((template) => ({
                    label: template.name,
                    value: template.id,
                    title: template.description,
                    config: template.labelConfig,
                  }))}
                  onChange={handleTemplateSelect}
                  optionRender={(option) => (
                    <div>
                      <div style={{ fontWeight: 500 }}>{option.label}</div>
                      {option.data.title && (
                        <div style={{ fontSize: 12, color: "#999", marginTop: 2 }}>
                          {option.data.title}
                        </div>
                      )}
                    </div>
                  )}
                />
              </Form.Item>
            )}

            <div className="bg-gray-50 p-4 rounded-md border border-gray-200">
              <TemplateConfigurationTreeEditor
                value={labelConfig}
                onChange={setLabelConfig}
                height={420}
              />
            </div>
          </Form>
        </div>
        <div className="flex gap-2 justify-end border-t border-gray-200 p-6">
          <Button onClick={() => navigate("/data/annotation")} disabled={submitting}>
            取消
          </Button>
          <Button type="primary" onClick={handleSubmit} loading={submitting}>
            创建任务
          </Button>
        </div>
      </div>
    </div>
  );
}