import { message } from "antd";
import { Cell, Graph } from "@antv/x6";
import { handleParseFowlAIData } from "@/utils";
import { aiRequest } from "./aiRequest";

type LabelMap = Record<
  string,
  { cell: string; key: string; cellType: string }[]
>;

// 从元素中提取文本 当前无选中元素时，提取所有元素
const getLabels = (graph?: Graph, cell?: Cell) => {
  const labelMap: LabelMap = {};
  let cells: Cell[] | undefined;

  if(cell) {
    cells = [cell];
  } else {
    cells = graph?.getSelectedCells();
  }
  
  if (!cells || cells.length === 0) {
    cells = graph?.getCells();
  }

  if (!cells || !cells.length) return;

  cells.forEach((cell) => {
    const data = cell.getData();
    // 从label中提取文本
    if (data?.label?.trim()) {
      if (!labelMap[data.label.trim()]) {
        labelMap[data.label.trim()] = [
          { cell: cell.id, key: "label", cellType: cell.shape },
        ];
      } else {
        labelMap[data.label.trim()].push({
          cell: cell.id,
          key: "label",
          cellType: cell.shape,
        });
      }
    }
    // 从name中提取文本
    if (data?.name?.trim()) {
      if (!labelMap[data.name.trim()]) {
        labelMap[data.name.trim()] = [
          { cell: cell.id, key: "name", cellType: cell.shape },
        ];
      } else {
        labelMap[data.name.trim()].push({
          cell: cell.id,
          key: "name",
          cellType: cell.shape,
        });
      }
    }
    // 从边线中提取文本
    if (cell.isEdge()) {
      (cell.labels || []).forEach((label) => {
        const labelText = (label?.attrs?.label?.text as string)?.trim();
        if (labelText) {
          if (!labelMap[labelText]) {
            labelMap[labelText] = [
              { cell: cell.id, key: "label", cellType: cell.shape },
            ];
          } else {
            labelMap[labelText].push({
              cell: cell.id,
              key: "label",
              cellType: cell.shape,
            });
          }
        }
      });
    }
  });
  return labelMap;
};

// 替换节点文本内容
const handleReplace = (
  labelMap: LabelMap,
  data: Record<string, string>,
  graph?: Graph
) => {
  const keyMap: Record<string, string> = {};
  if (Array.isArray(data)) {
    data.forEach((item) => {
      keyMap[item.original] = item.value;
    });
  }

  Object.keys(keyMap).forEach((key) => {
    if (labelMap[key]) {
      labelMap[key].forEach((item) => {
        const cell = graph?.getCellById(item.cell);
        console.log(cell, item, graph);
        if (cell && cell.shape !== "edge") {
          cell.setData({
            [item.key]: keyMap[key],
          });
          graph?.select(cell);
        } else if (cell?.isEdge()) {
          // 设置边线文本
          const labels = cell.getLabels();
          cell.setLabels(
            labels.map((item) => {
              return {
                ...item,
                attrs: {
                  label: {
                    text:
                      keyMap?.[(item.attrs?.label?.text as string)?.trim()] ||
                      item.attrs?.label?.text,
                  },
                },
              };
            })
          );
        }
      });
    }
  });
};

// 语法修复
export const handleGrammar = (graph?: Graph, cell?: Cell) => {
  const labelMap = getLabels(graph, cell);
  if (!labelMap) {
    message.open({
      key: "ai-handler-message",
      type: "info",
      content: "无可修复的数据",
    });
    return;
  }

  message.open({
    key: "ai-handler-message",
    type: "loading",
    content: "AI正在修复语法...",
    duration: 0,
  });

  const data = JSON.stringify(Object.keys(labelMap));
  let result = "";
  aiRequest(
    `修复语法错误，需要修复的数据：${data}`,
    {
      onUpdate: (data) => {
        result += data.answer;
      },
      onSuccess: (data) => {
        message.open({
          key: "ai-handler-message",
          type: "success",
          content: "AI修复语法完成",
          duration: 2,
          style: {
            marginTop: 300,
          },
        });
        handleParseFowlAIData({
          content: result,
          message: {
            key: "ai-handler-message",
            instance: message,
          },
          onSuccess: (data) => {
            handleReplace(labelMap, data, graph);
          },
        });
      },
      onError: (err) => {
        console.error(err);
        message.open({
          key: "ai-handler-message",
          type: "error",
          content: err.message || "AI创作失败",
          duration: 2,
          style: {
            marginTop: 300,
          },
        });
      },
    },
    "语法修复"
  );
};

// 翻译
export const handleTranslation = (lang: "en" | "zh", graph?: Graph, cell?: Cell) => {
  const labelMap = getLabels(graph, cell);
  if (!labelMap) {
    message.open({
      key: "ai-handler-message",
      type: "info",
      content: "无可翻译的数据",
    });
    return;
  }

  message.open({
    key: "ai-handler-message",
    type: "loading",
    content: "AI正在翻译...",
    duration: 0,
  });

  const data = JSON.stringify(Object.keys(labelMap));
  let result = "";
  aiRequest(
    `翻译成${lang === "en" ? "英文" : "中文"}，需要翻译的数据：${data}`,
    {
      onUpdate: (data) => {
        result += data.answer;
      },
      onSuccess: (data) => {
        message.open({
          key: "ai-handler-message",
          type: "success",
          content: "AI翻译完成",
          duration: 2,
          style: {
            marginTop: 300,
          },
        });
        handleParseFowlAIData({
          content: result,
          message: {
            key: "ai-handler-message",
            instance: message,
          },
          onSuccess: (data) => {
            handleReplace(labelMap, data, graph);
          },
        });
      },
      onError: (err) => {
        console.error(err);
        message.open({
          key: "ai-handler-message",
          type: "error",
          content: err.message || "AI创作失败",
          duration: 2,
          style: {
            marginTop: 300,
          },
        });
      },
    },
    "翻译内容"
  );
};
