link-stack/apps/metamigo-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/metamigo-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 = 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;