import "dotenv/config";
import { Jetstream } from "@skyware/jetstream";
import { LABELS } from "./const";
import { pipeline, TextClassificationOutput } from "@huggingface/transformers";
import pThrottle from "p-throttle";
import { Agent, CredentialSession } from "@atproto/api";
import slugify from "slugify";

const session = new CredentialSession(new URL("https://bsky.social"));
await session.login({
  identifier: "news.aendra.dev",
  password: process.env.BSKY_PASS!,
});
const agent = new Agent(session);

const throttle = pThrottle({
  limit: 1,
  interval: 5000,
});

const infer = await pipeline("text-classification", "./models/bert-news-class");

const MODEL_NAME = "cssupport/bert-news-class_v1";

const createLabel = async (
  post: { uri: string; cid: string },
  label: string,
  score: number | string,
  model: string = MODEL_NAME,
) => {
  try {
    const { uri, cid } = post;

    return agent.tools.ozone.moderation.emitEvent(
      {
        // specify the label event
        event: {
          $type: "tools.ozone.moderation.defs#modEventLabel",
          createLabelVals: [`class-${slugify(label)}`.toLowerCase()],
          negateLabelVals: [],
          comment: `Inferred by model ${model} (${score})`,
        },
        // specify the labeled post by strongRef
        subject: {
          $type: "com.atproto.repo.strongRef",
          uri,
          cid,
        },
        createdBy: session.did!,
        subjectBlobCids: [],
      },
      {
        encoding: "application/json",
        headers: {
          "atproto-proxy": `${session.did}#atproto_labeler`,
        },
      },
    );
  } catch (e) {
    console.error(e);
  }
};

const FEED_URI = "https://bsky.app/profile/aendra.com/feed/verified-news";
const CONTRAILS_ENDPOINT = `wss://api.graze.social/app/contrail?feed=${FEED_URI}`;

const jetstream = new Jetstream({
  wantedCollections: ["app.bsky.feed.post"], // omit to receive all collections
  endpoint: CONTRAILS_ENDPOINT, // Uncomment to get just the News feed
});

jetstream.start();

jetstream.onCreate("app.bsky.feed.post", async (event) => {
  const fetchDescription = throttle(async (uri: string) =>
    fetch(`https://cardyb.bsky.app/v1/extract?url=${uri}`),
  );
  if (
    event.commit.record.embed?.$type === "app.bsky.embed.external" &&
    event.commit.record.langs?.includes("en")
  ) {
    try {
      const { description, title }: { description: string; title: string } =
        await (
          await fetchDescription(event.commit.record.embed.external.uri)
        ).json();

      if (description) {
        const [result] = (
          (await infer(`${title}: ${description}`, {
            top_k: 5,
          })) as TextClassificationOutput
        )
          ?.filter(
            (d) =>
              !Object.values(LABELS)
                .filter((d) => d.includes("UNUSED"))
                .includes(d.label as LABELS),
          )
          .sort((a, b) => b.score - a.score)
          .map(({ label, score }: { score: number; label: string }) => ({
            label: LABELS[label as keyof typeof LABELS],
            score,
          }));

        // if (result.score > 0.5) {
        await createLabel(
          {
            uri: `at://${event.did}/${event.commit.collection}/${event.commit.rkey}`,
            cid: event.commit.cid,
          },
          result.label,
          result.score,
        );
        console.log(
          `https://bsky.app/profile/${event.did}/post/${event.commit.rkey}`,
          result.label,
          result.score,
        );
        // }
      }
    } catch (e) {
      console.error(e);
    }
  }
});