DataMate/frontend/src/pages/DataAnnotation/Annotate/LabelStudioTextEditor.tsx

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { App, Button, Card, List, Spin, Typography, Tag, Switch } from "antd";
import { LeftOutlined, ReloadOutlined, SaveOutlined, MenuFoldOutlined, MenuUnfoldOutlined, CheckOutlined } from "@ant-design/icons";
import { useNavigate, useParams } from "react-router";

import {
  getEditorProjectInfoUsingGet,
  getEditorTaskUsingGet,
  listEditorTasksUsingGet,
  upsertEditorAnnotationUsingPut,
} from "../annotation.api";

type EditorProjectInfo = {
  projectId: string;
  datasetId: string;
  templateId?: string | null;
  labelConfig?: string | null;
  supported: boolean;
  unsupportedReason?: string | null;
};

type EditorTaskListItem = {
  fileId: string;
  fileName: string;
  fileType?: string | null;
  hasAnnotation: boolean;
  annotationUpdatedAt?: string | null;
};

type LsfMessage = {
  type?: string;
  payload?: unknown;
};

type SegmentInfo = {
  idx: number;
  text: string;
  start: number;
  end: number;
  hasAnnotation: boolean;
};

type ApiResponse<T> = {
  code?: number;
  message?: string;
  data?: T;
};

type EditorTaskPayload = {
  id?: number | string;
  data?: Record<string, unknown>;
  annotations?: unknown[];
};

type EditorTaskResponse = {
  task?: EditorTaskPayload;
  segmented?: boolean;
  segments?: SegmentInfo[];
  currentSegmentIndex?: number;
};

type EditorTaskListResponse = {
  content?: EditorTaskListItem[];
};

type ExportPayload = {
  taskId?: number | string | null;
  fileId?: string | null;
  segmentIndex?: number | string | null;
  annotation?: Record<string, unknown>;
  requestId?: string | null;
};

type SwitchDecision = "save" | "discard" | "cancel";

const LSF_IFRAME_SRC = "/lsf/lsf.html";

const resolveSegmentIndex = (value: unknown) => {
  if (value === null || value === undefined) return undefined;
  const parsed = Number(value);
  return Number.isFinite(parsed) ? parsed : undefined;
};

const normalizePayload = (payload: unknown): ExportPayload | undefined => {
  if (!payload || typeof payload !== "object") return undefined;
  return payload as ExportPayload;
};

const resolvePayloadMessage = (payload: unknown) => {
  if (!payload || typeof payload !== "object") return undefined;
  if ("message" in payload && typeof (payload as { message?: unknown }).message === "string") {
    return (payload as { message?: string }).message;
  }
  return undefined;
};

const isRecord = (value: unknown): value is Record<string, unknown> =>
  !!value && typeof value === "object" && !Array.isArray(value);

const normalizeSnapshotValue = (value: unknown, seen: WeakSet<object>): unknown => {
  if (!value || typeof value !== "object") return value;
  const obj = value as object;
  if (seen.has(obj)) return undefined;
  seen.add(obj);
  if (Array.isArray(value)) {
    return value.map((item) => normalizeSnapshotValue(item, seen));
  }
  const record = value as Record<string, unknown>;
  const sorted: Record<string, unknown> = {};
  Object.keys(record)
    .sort()
    .forEach((key) => {
      sorted[key] = normalizeSnapshotValue(record[key], seen);
    });
  return sorted;
};

const stableStringify = (value: unknown) => {
  const normalized = normalizeSnapshotValue(value, new WeakSet<object>());
  return JSON.stringify(normalized);
};

const buildAnnotationSnapshot = (annotation?: Record<string, unknown>) => {
  if (!annotation) return "";
  const cleaned: Record<string, unknown> = { ...annotation };
  delete cleaned.updated_at;
  delete cleaned.updatedAt;
  delete cleaned.created_at;
  delete cleaned.createdAt;
  return stableStringify(cleaned);
};

const buildSnapshotKey = (fileId: string, segmentIndex?: number) =>
  `${fileId}::${segmentIndex ?? "full"}`;

export default function LabelStudioTextEditor() {
  const { projectId = "" } = useParams();
  const navigate = useNavigate();
  const { message, modal } = App.useApp();

  const origin = useMemo(() => window.location.origin, []);
  const iframeRef = useRef<HTMLIFrameElement | null>(null);
  const initSeqRef = useRef(0);
  const expectedTaskIdRef = useRef<number | null>(null);
  const exportCheckRef = useRef<{
    requestId: string;
    resolve: (payload?: ExportPayload) => void;
    timer?: number;
  } | null>(null);
  const exportCheckSeqRef = useRef(0);
  const savedSnapshotsRef = useRef<Record<string, string>>({});

  const [loadingProject, setLoadingProject] = useState(true);
  const [loadingTasks, setLoadingTasks] = useState(false);
  const [loadingTaskDetail, setLoadingTaskDetail] = useState(false);
  const [saving, setSaving] = useState(false);
  const [segmentSwitching, setSegmentSwitching] = useState(false);

  const [iframeReady, setIframeReady] = useState(false);
  const [lsReady, setLsReady] = useState(false);
  const [project, setProject] = useState<EditorProjectInfo | null>(null);
  const [tasks, setTasks] = useState<EditorTaskListItem[]>([]);
  const [selectedFileId, setSelectedFileId] = useState<string>("");
  const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
  const [autoSaveOnSwitch, setAutoSaveOnSwitch] = useState(false);

  // 分段相关状态
  const [segmented, setSegmented] = useState(false);
  const [segments, setSegments] = useState<SegmentInfo[]>([]);
  const [currentSegmentIndex, setCurrentSegmentIndex] = useState(0);

  const postToIframe = useCallback((type: string, payload?: unknown) => {
    const win = iframeRef.current?.contentWindow;
    if (!win) return;
    win.postMessage({ type, payload }, origin);
  }, [origin]);

  const loadProject = useCallback(async () => {
    setLoadingProject(true);
    try {
      const resp = (await getEditorProjectInfoUsingGet(projectId)) as ApiResponse<EditorProjectInfo>;
      const data = resp?.data;
      if (!data?.projectId) {
        message.error("获取标注项目信息失败");
        setProject(null);
        return;
      }
      setProject(data);
    } catch (e) {
      console.error(e);
      message.error("获取标注项目信息失败");
      setProject(null);
    } finally {
      setLoadingProject(false);
    }
  }, [message, projectId]);

  const loadTasks = useCallback(async (silent = false) => {
    if (!projectId) return;
    if (!silent) setLoadingTasks(true);
    try {
      const resp = (await listEditorTasksUsingGet(projectId, {
        page: 0,
        size: 200,
      })) as ApiResponse<EditorTaskListResponse>;
      const content = resp?.data?.content || [];
      const items = Array.isArray(content) ? content : [];
      setTasks(items);
      if (items.length > 0) {
        setSelectedFileId((prev) => prev || (items[0]?.fileId ?? ""));
      }
    } catch (e) {
      console.error(e);
      if (!silent) message.error("获取文件列表失败");
      setTasks([]);
    } finally {
      if (!silent) setLoadingTasks(false);
    }
  }, [message, projectId]);

  const initEditorForFile = useCallback(async (fileId: string, segmentIdx?: number) => {
    if (!project?.supported) return;
    if (!project?.labelConfig) {
      message.error("该项目未绑定标注模板，无法加载编辑器");
      return;
    }
    if (!iframeReady) return;

    const seq = ++initSeqRef.current;
    setLoadingTaskDetail(true);
    setLsReady(false);
    expectedTaskIdRef.current = null;

    try {
      const resp = (await getEditorTaskUsingGet(projectId, fileId, {
        segmentIndex: segmentIdx,
      })) as ApiResponse<EditorTaskResponse>;
      const data = resp?.data;
      const task = data?.task;
      if (!task) {
        message.error("获取任务详情失败");
        return;
      }
      if (seq !== initSeqRef.current) return;

      // 更新分段状态
      const segmentIndex = data?.segmented
        ? resolveSegmentIndex(data.currentSegmentIndex) ?? 0
        : undefined;
      if (data?.segmented) {
        setSegmented(true);
        setSegments(data.segments || []);
        setCurrentSegmentIndex(segmentIndex ?? 0);
      } else {
        setSegmented(false);
        setSegments([]);
        setCurrentSegmentIndex(0);
      }

      const taskData = {
        ...(task?.data || {}),
        file_id: fileId,
        fileId: fileId,
      };
      if (data?.segmented) {
        const normalizedIndex = segmentIndex ?? 0;
        taskData.segment_index = normalizedIndex;
        taskData.segmentIndex = normalizedIndex;
      }
      const taskForIframe = {
        ...task,
        data: taskData,
      };

      const annotations = Array.isArray(taskForIframe.annotations) ? taskForIframe.annotations : [];
      const initialAnnotation = annotations.length > 0 && isRecord(annotations[0]) ? annotations[0] : undefined;
      savedSnapshotsRef.current[buildSnapshotKey(fileId, segmentIndex)] =
        buildAnnotationSnapshot(initialAnnotation);

      expectedTaskIdRef.current = Number(taskForIframe?.id) || null;
      postToIframe("LS_INIT", {
        labelConfig: project.labelConfig,
        task: taskForIframe,
        user: { id: "datamate" },
        // 完整的 Label Studio 原生界面配置
        interfaces: [
          // 核心面板
          "panel",              // 导航面板（undo/redo/reset）
          "update",             // 更新按钮
          "submit",             // 提交按钮
          "controls",           // 控制面板
          // 侧边栏（包含 Outliner 和 Details）
          "side-column",
          // 标注管理
          "annotations:tabs",
          "annotations:menu",
          "annotations:current",
          "annotations:add-new",
          "annotations:delete",
          "annotations:view-all",
          // 预测
          "predictions:tabs",
          "predictions:menu",
          // 其他
          "auto-annotation",
          "edit-history",
        ],
        selectedAnnotationIndex: 0,
        allowCreateEmptyAnnotation: true,
      });
    } catch (e) {
      console.error(e);
      message.error("加载编辑器失败");
    } finally {
      if (seq === initSeqRef.current) setLoadingTaskDetail(false);
    }
  }, [iframeReady, message, postToIframe, project, projectId]);

  const saveFromExport = useCallback(async (payload?: ExportPayload | null) => {
    const payloadTaskId = payload?.taskId;
    if (expectedTaskIdRef.current && payloadTaskId) {
      if (Number(payloadTaskId) !== expectedTaskIdRef.current) {
        message.warning("已忽略过期的保存请求");
        return false;
      }
    }
    const fileId = payload?.fileId || selectedFileId;
    const annotation = payload?.annotation;
    if (!fileId || !annotation || typeof annotation !== "object") {
      message.error("导出标注失败：缺少 fileId/annotation");
      return false;
    }
    const payloadSegmentIndex = resolveSegmentIndex(payload?.segmentIndex);
    const segmentIndex =
      payloadSegmentIndex !== undefined
        ? payloadSegmentIndex
        : segmented
          ? currentSegmentIndex
          : undefined;

    setSaving(true);
    try {
      await upsertEditorAnnotationUsingPut(projectId, String(fileId), {
        annotation,
        segmentIndex,
      });
      message.success("标注已保存");
      await loadTasks(true);

      const snapshotKey = buildSnapshotKey(String(fileId), segmentIndex);
      const snapshot = buildAnnotationSnapshot(isRecord(annotation) ? annotation : undefined);
      savedSnapshotsRef.current[snapshotKey] = snapshot;

      // 分段模式下更新当前段落的标注状态
      if (segmented && segmentIndex !== undefined) {
        setSegments((prev) =>
          prev.map((seg) =>
            seg.idx === segmentIndex
              ? { ...seg, hasAnnotation: true }
              : seg
          )
        );
      }
      return true;
    } catch (e) {
      console.error(e);
      message.error("保存失败");
      return false;
    } finally {
      setSaving(false);
    }
  }, [
    currentSegmentIndex,
    loadTasks,
    message,
    projectId,
    segmented,
    selectedFileId,
  ]);

  const requestExportForCheck = useCallback(() => {
    if (!iframeReady || !lsReady) return Promise.resolve(undefined);
    if (exportCheckRef.current) {
      if (exportCheckRef.current.timer) {
        window.clearTimeout(exportCheckRef.current.timer);
      }
      exportCheckRef.current.resolve(undefined);
      exportCheckRef.current = null;
    }
    const requestId = `check_${Date.now()}_${++exportCheckSeqRef.current}`;
    return new Promise<ExportPayload | undefined>((resolve) => {
      const timer = window.setTimeout(() => {
        if (exportCheckRef.current?.requestId === requestId) {
          exportCheckRef.current = null;
        }
        resolve(undefined);
      }, 3000);
      exportCheckRef.current = {
        requestId,
        resolve,
        timer,
      };
      postToIframe("LS_EXPORT_CHECK", { requestId });
    });
  }, [iframeReady, lsReady, postToIframe]);

  const confirmSaveBeforeSwitch = useCallback(() => {
    return new Promise<SwitchDecision>((resolve) => {
      let resolved = false;
      let modalInstance: { destroy: () => void } | null = null;
      const settle = (decision: SwitchDecision) => {
        if (resolved) return;
        resolved = true;
        resolve(decision);
      };
      const handleDiscard = () => {
        if (modalInstance) modalInstance.destroy();
        settle("discard");
      };
      modalInstance = modal.confirm({
        title: "当前段落有未保存标注",
        content: (
          <div className="flex flex-col gap-2">
            <Typography.Text>切换段落前请先保存当前标注。</Typography.Text>
            <Button type="link" danger style={{ padding: 0, height: "auto" }} onClick={handleDiscard}>
              放弃未保存并切换
            </Button>
          </div>
        ),
        okText: "保存并切换",
        cancelText: "取消",
        onOk: () => settle("save"),
        onCancel: () => settle("cancel"),
      });
    });
  }, [modal]);

  const requestExport = () => {
    if (!selectedFileId) {
      message.warning("请先选择文件");
      return;
    }
    postToIframe("LS_EXPORT", {});
  };

  // 段落切换处理
  const handleSegmentChange = async (newIndex: number) => {
    if (newIndex === currentSegmentIndex) return;
    if (segmentSwitching || saving || loadingTaskDetail) return;
    if (!iframeReady || !lsReady) {
      message.warning("编辑器未就绪，无法切换段落");
      return;
    }

    setSegmentSwitching(true);
    try {
      const payload = await requestExportForCheck();
      if (!payload) {
        message.warning("无法读取当前标注，已取消切换");
        return;
      }

      const payloadTaskId = payload.taskId;
      if (expectedTaskIdRef.current && payloadTaskId) {
        if (Number(payloadTaskId) !== expectedTaskIdRef.current) {
          message.warning("已忽略过期的标注数据");
          return;
        }
      }

      const payloadFileId = payload.fileId || selectedFileId;
      const payloadSegmentIndex = resolveSegmentIndex(payload.segmentIndex);
      const resolvedSegmentIndex =
        payloadSegmentIndex !== undefined
          ? payloadSegmentIndex
          : segmented
            ? currentSegmentIndex
            : undefined;
      const annotation = isRecord(payload.annotation) ? payload.annotation : undefined;
      const snapshotKey = payloadFileId
        ? buildSnapshotKey(String(payloadFileId), resolvedSegmentIndex)
        : undefined;
      const latestSnapshot = buildAnnotationSnapshot(annotation);
      const lastSnapshot = snapshotKey ? savedSnapshotsRef.current[snapshotKey] : undefined;
      const hasUnsavedChange = snapshotKey !== undefined && lastSnapshot !== undefined && latestSnapshot !== lastSnapshot;

      if (hasUnsavedChange) {
        if (autoSaveOnSwitch) {
          const saved = await saveFromExport(payload);
          if (!saved) return;
        } else {
          const decision = await confirmSaveBeforeSwitch();
          if (decision === "cancel") return;
          if (decision === "save") {
            const saved = await saveFromExport(payload);
            if (!saved) return;
          }
        }
      }

      await initEditorForFile(selectedFileId, newIndex);
    } finally {
      setSegmentSwitching(false);
    }
  };

  useEffect(() => {
    setIframeReady(false);
    setProject(null);
    setTasks([]);
    setSelectedFileId("");
    initSeqRef.current = 0;
    setLsReady(false);
    expectedTaskIdRef.current = null;
    // 重置分段状态
    setSegmented(false);
    setSegments([]);
    setCurrentSegmentIndex(0);
    savedSnapshotsRef.current = {};
    if (exportCheckRef.current?.timer) {
      window.clearTimeout(exportCheckRef.current.timer);
    }
    exportCheckRef.current = null;

    if (projectId) loadProject();
  }, [projectId, loadProject]);

  useEffect(() => {
    if (!project?.supported) return;
    loadTasks();
  }, [project?.supported, loadTasks]);

  useEffect(() => {
    if (!selectedFileId) return;
    initEditorForFile(selectedFileId);
  }, [selectedFileId, iframeReady, initEditorForFile]);

  useEffect(() => {
    const handler = (event: MessageEvent<LsfMessage>) => {
      if (event.origin !== origin) return;
      const msg = event.data || {};
      if (!msg?.type) return;

      if (msg.type === "LS_IFRAME_READY") {
        setIframeReady(true);
        return;
      }

      const payload = normalizePayload(msg.payload);

      if (msg.type === "LS_READY") {
        const readyTaskId = payload?.taskId;
        if (expectedTaskIdRef.current && readyTaskId) {
          if (Number(readyTaskId) !== expectedTaskIdRef.current) return;
        }
        setLsReady(true);
        return;
      }

      if (msg.type === "LS_EXPORT_RESULT") {
        saveFromExport(payload);
        return;
      }

      if (msg.type === "LS_EXPORT_CHECK_RESULT") {
        const pending = exportCheckRef.current;
        if (!pending) return;
        const requestId = payload?.requestId;
        if (requestId && requestId !== pending.requestId) return;
        if (pending.timer) {
          window.clearTimeout(pending.timer);
        }
        exportCheckRef.current = null;
        pending.resolve(payload);
        return;
      }

      // 兼容 iframe 内部在 submit 时直接上报（若启用）
      if (msg.type === "LS_SUBMIT") {
        saveFromExport(payload);
        return;
      }

      if (msg.type === "LS_ERROR") {
        const payloadMessage = resolvePayloadMessage(msg.payload);
        message.error(payloadMessage || "编辑器发生错误");
        setLsReady(false);
      }
    };

    window.addEventListener("message", handler);
    return () => window.removeEventListener("message", handler);
  }, [message, origin, saveFromExport]);

  if (loadingProject) {
    return (
      <div className="h-full flex items-center justify-center">
        <Spin />
      </div>
    );
  }

  if (!project) {
    return (
      <div className="h-full flex items-center justify-center">
        <Card>
          <Typography.Text>未找到标注项目</Typography.Text>
          <div className="mt-4 flex justify-end">
            <Button onClick={() => navigate("/data/annotation")}>返回</Button>
          </div>
        </Card>
      </div>
    );
  }

  if (!project.supported) {
    return (
      <div className="h-full flex items-center justify-center">
        <Card style={{ maxWidth: 640 }}>
          <Typography.Title level={4}>暂不支持该数据类型</Typography.Title>
          <Typography.Paragraph type="secondary">
            {project.unsupportedReason || "当前仅支持文本（TEXT）项目的内嵌编辑器。"}
          </Typography.Paragraph>
          <div className="flex justify-end gap-2">
            <Button onClick={() => navigate("/data/annotation")}>返回</Button>
          </div>
        </Card>
      </div>
    );
  }

  return (
    <div className="h-full flex flex-col">
      {/* 顶部工具栏 */}
      <div className="flex items-center justify-between px-3 py-2 border-b border-gray-200 bg-white">
        <div className="flex items-center gap-2">
          <Button icon={<LeftOutlined />} onClick={() => navigate("/data/annotation")}>
            返回
          </Button>
          <Button
            icon={sidebarCollapsed ? <MenuUnfoldOutlined /> : <MenuFoldOutlined />}
            onClick={() => setSidebarCollapsed(!sidebarCollapsed)}
            title={sidebarCollapsed ? "展开文件列表" : "收起文件列表"}
          />
          <Typography.Title level={5} style={{ margin: 0 }}>
            标注编辑器
          </Typography.Title>
        </div>
        <div className="flex items-center gap-2">
          <Button icon={<ReloadOutlined />} loading={loadingTasks} onClick={() => loadTasks()}>
            刷新
          </Button>
          <Button
            type="primary"
            icon={<SaveOutlined />}
            loading={saving}
            disabled={!iframeReady || !selectedFileId}
            onClick={requestExport}
          >
            保存
          </Button>
        </div>
      </div>

      {/* 主体区域 */}
      <div className="flex flex-1 min-h-0">
        {/* 左侧文件列表 - 可折叠 */}
        <div
          className="border-r border-gray-200 bg-gray-50 flex flex-col transition-all duration-200"
          style={{ width: sidebarCollapsed ? 0 : 240, overflow: "hidden" }}
        >
          <div className="px-3 py-2 border-b border-gray-200 bg-white font-medium text-sm">
            文件列表
          </div>
          <div className="flex-1 overflow-auto">
            <List
              loading={loadingTasks}
              size="small"
              dataSource={tasks}
              renderItem={(item) => (
                <List.Item
                  key={item.fileId}
                  className="cursor-pointer hover:bg-blue-50"
                  style={{
                    background: item.fileId === selectedFileId ? "#e6f4ff" : undefined,
                    padding: "8px 12px",
                    borderBottom: "1px solid #f0f0f0",
                  }}
                  onClick={() => setSelectedFileId(item.fileId)}
                >
                  <div className="flex flex-col w-full gap-1">
                    <Typography.Text ellipsis style={{ fontSize: 13 }}>
                      {item.fileName}
                    </Typography.Text>
                    <div className="flex items-center justify-between">
                      <Typography.Text
                        type={item.hasAnnotation ? "success" : "secondary"}
                        style={{ fontSize: 11 }}
                      >
                        {item.hasAnnotation ? "已标注" : "未标注"}
                      </Typography.Text>
                      {item.annotationUpdatedAt && (
                        <Typography.Text type="secondary" style={{ fontSize: 10 }}>
                          {item.annotationUpdatedAt}
                        </Typography.Text>
                      )}
                    </div>
                  </div>
                </List.Item>
              )}
            />
          </div>
        </div>

        {/* 右侧编辑器 - Label Studio iframe */}
        <div className="flex-1 flex flex-col min-h-0">
          {/* 段落导航栏 */}
          {segmented && segments.length > 0 && (
            <div className="flex items-center gap-2 px-3 py-2 bg-gray-50 border-b border-gray-200">
              <Typography.Text style={{ fontSize: 12 }}>段落：</Typography.Text>
              <div className="flex gap-1 flex-wrap">
                {segments.map((seg) => (
                  <Button
                    key={seg.idx}
                    size="small"
                    type={seg.idx === currentSegmentIndex ? "primary" : "default"}
                    onClick={() => handleSegmentChange(seg.idx)}
                    disabled={segmentSwitching || saving || loadingTaskDetail || !lsReady}
                    style={{ minWidth: 32, padding: "0 8px" }}
                  >
                    {seg.idx + 1}
                    {seg.hasAnnotation && (
                      <CheckOutlined style={{ marginLeft: 2, fontSize: 10 }} />
                    )}
                  </Button>
                ))}
              </div>
              <div className="flex items-center gap-2 ml-auto">
                <Tag color="blue">{currentSegmentIndex + 1} / {segments.length}</Tag>
                <Typography.Text style={{ fontSize: 12 }}>切段自动保存</Typography.Text>
                <Switch
                  size="small"
                  checked={autoSaveOnSwitch}
                  onChange={(checked) => setAutoSaveOnSwitch(checked)}
                  disabled={segmentSwitching || saving || loadingTaskDetail || !lsReady}
                />
              </div>
            </div>
          )}

          {/* 编辑器区域 */}
          <div className="flex-1 relative">
            {(!iframeReady || loadingTaskDetail || (selectedFileId && !lsReady)) && (
              <div className="absolute inset-0 z-10 flex items-center justify-center bg-white/80">
                <Spin
                  tip={
                    !iframeReady
                      ? "编辑器资源加载中..."
                      : loadingTaskDetail
                        ? "任务数据加载中..."
                        : "编辑器初始化中..."
                  }
                />
              </div>
            )}
            <iframe
              ref={iframeRef}
              title="Label Studio Frontend"
              src={LSF_IFRAME_SRC}
              className="w-full h-full border-0"
            />
          </div>
        </div>
      </div>
    </div>
  );
}