link-stack/apps/bridge-worker/tasks/import-leafcutter.ts

/* eslint-disable camelcase */
import fetch from "node-fetch";
import { URLSearchParams } from "url";
import { withDb, AppDatabase } from "../db";
import { loadConfig } from "@digiresilience/bridge-config";

type LabelStudioTicket = {
  id: string;
  is_labeled: boolean;
  annotations: Record<string, unknown>[];
  data: Record<string, unknown>;
  updated_at: string;
};

type LeafcutterTicket = {
  id: string;
  incident: string[];
  technology: string[];
  targeted_group: string[];
  country: string[];
  region: string[];
  continent: string[];
  date: Date;
  origin: string;
  origin_id: string;
  source_created_at: string;
  source_updated_at: string;
};

const getLabelStudioTickets = async (
  page: number,
): Promise<LabelStudioTicket[]> => {
  const {
    leafcutter: { labelStudioApiUrl, labelStudioApiKey },
  } = await loadConfig();

  const headers = {
    Authorization: `Token ${labelStudioApiKey}`,
    Accept: "application/json",
  };
  const ticketsQuery = new URLSearchParams({
    page_size: "50",
    page: `${page}`,
  });
  console.log({ url: `${labelStudioApiUrl}/projects/1/tasks?${ticketsQuery}` });
  const res = await fetch(
    `${labelStudioApiUrl}/projects/1/tasks?${ticketsQuery}`,
    { headers },
  );
  console.log({ res });
  const tasksResult: any = await res.json();
  console.log({ tasksResult });

  return tasksResult;
};

const fetchFromLabelStudio = async (
  minUpdatedTimestamp: Date,
): Promise<LabelStudioTicket[]> => {
  const pages = [...Array.from({ length: 10000 }).keys()];
  const allDocs: LabelStudioTicket[] = [];

  for await (const page of pages) {
    const docs = await getLabelStudioTickets(page + 1);
    console.log({ page, docs });

    if (docs && docs.length > 0) {
      for (const doc of docs) {
        const updatedAt = new Date(doc.updated_at);
        console.log({ updatedAt, minUpdatedTimestamp });
        if (updatedAt > minUpdatedTimestamp) {
          console.log(`Adding doc`, { doc });
          allDocs.push(doc);
        }
      }
    } else {
      break;
    }
  }

  console.log({ allDocs });
  return allDocs;
};

const sendToLeafcutter = async (tickets: LabelStudioTicket[]) => {
  const {
    leafcutter: {
      contributorId,
      opensearchApiUrl,
      opensearchUsername,
      opensearchPassword,
    },
  } = await loadConfig();

  console.log({ tickets });
  const filteredTickets = tickets.filter((ticket) => ticket.is_labeled);
  console.log({ filteredTickets });
  const finalTickets: LeafcutterTicket[] = filteredTickets.map((ticket) => {
    const {
      id,
      annotations,
      data: { source_id, source_created_at, source_updated_at },
    } = ticket;

    const getTags = (tags: Record<string, any>[], name: string) =>
      tags
        .filter((tag) => tag.from_name === name)
        .map((tag) => tag.value.choices)
        .flat();

    const allTags = annotations.map(({ result }) => result).flat();
    const incident = getTags(allTags, "incidentType tag");
    const technology = getTags(allTags, "platform tag");
    const country = getTags(allTags, "country tag");
    const targetedGroup = getTags(allTags, "targetedGroup tag");

    return {
      id,
      incident,
      technology,
      targeted_group: targetedGroup,
      country,
      region: [],
      continent: [],
      date: new Date(source_created_at as string),
      origin: contributorId,
      origin_id: source_id as string,
      source_created_at: source_created_at as string,
      source_updated_at: source_updated_at as string,
    };
  });

  console.log("Sending to Leafcutter");
  console.log({ finalTickets });

  const result = await fetch(opensearchApiUrl, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Basic ${Buffer.from(`${opensearchUsername}:${opensearchPassword}`).toString("base64")}`,
    },
    body: JSON.stringify({ tickets: finalTickets }),
  });
  console.log({ result });
};

const importLeafcutterTask = async (): Promise<void> => {
  withDb(async (db: AppDatabase) => {
    const {
      leafcutter: { contributorName },
    } = await loadConfig();
    const settingName = `${contributorName}ImportLeafcutterTask`;
    const res: any = await db.settings.findByName(settingName);
    const startTimestamp = res?.value?.minUpdatedTimestamp
      ? new Date(res.value.minUpdatedTimestamp as string)
      : new Date("2023-03-01");
    const newLastTimestamp = new Date();
    console.log({
      contributorName,
      settingName,
      res,
      startTimestamp,
      newLastTimestamp,
    });
    const tickets = await fetchFromLabelStudio(startTimestamp);
    console.log({ tickets });
    await sendToLeafcutter(tickets);
    await db.settings.upsert(settingName, {
      minUpdatedTimestamp: newLastTimestamp,
    });
  });
};

export default importLeafcutterTask;