From 809409920eb52198428a7a9584acf9ec87642078 Mon Sep 17 00:00:00 2001
From: derpystuff <3515180-derpystuff@users.noreply.gitlab.com>
Date: Thu, 18 Jan 2024 21:34:06 +0100
Subject: [PATCH] screenshot

---
 commands/message/utils/screenshot.js | 21 +++++----------------
 labscore/api/obelisk/endpoints.js    |  4 +++-
 labscore/api/obelisk/index.js        |  8 ++++++++
 3 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/commands/message/utils/screenshot.js b/commands/message/utils/screenshot.js
index 53144eb..050dd99 100644
--- a/commands/message/utils/screenshot.js
+++ b/commands/message/utils/screenshot.js
@@ -1,10 +1,10 @@
-const { screenshot } = require("../../../labscore/api");
 const { createEmbed } = require("../../../labscore/utils/embed");
 const { editOrReply } = require("../../../labscore/utils/message");
 
 const superagent = require('superagent')
 
 const { Permissions } = require("detritus-client/lib/constants");
+const { webshot } = require("../../../labscore/api/obelisk");
 
 async function processJob(jobUrl) {
   let job = await superagent.get(jobUrl)
@@ -28,12 +28,12 @@ module.exports = {
   run: async (context, args) => {
     if (!args.url) return editOrReply(context, createEmbed("warning", context, "No url supplied."))
 
-    let response = await editOrReply(context, createEmbed("loading", context, `Creating website screenshot...`))
+    await editOrReply(context, createEmbed("loading", context, `Creating website screenshot...`))
 
     try {
       const t = Date.now();
 
-      let ss = await screenshot(context, args.url, context.channel.nsfw)
+      let ss = await webshot(context, args.url, context.channel.nsfw)
 
       if (ss.response.body.status && ss.response.body.status !== 3) {
         if (ss.response.body.image) return await editOrReply(context,
@@ -45,28 +45,17 @@ module.exports = {
         return await editOrReply(context, createEmbed("error", context, "Unable to create screenshot."))
       }
 
-      let job = await processJob(ss.response.body.job)
-
-      if (job.status) {
-        if (!job.image) job = await processJob(ss.response.body.job)
-        if (job.image) return await editOrReply(context, createEmbed("image", context, {
-          url: job.image,
-          time: ((Date.now() - t) / 1000).toFixed(2)
-        }))
-        return await editOrReply(context, createEmbed("error", context, "Unable to create screenshot."))
-      }
-
       return await editOrReply(context, {
         embeds: [createEmbed("image", context, {
           url: "screenshot.png",
           time: ((Date.now() - t) / 1000).toFixed(2)
         })],
-        files: [{ filename: "screenshot.png", value: job }]
+        files: [{ filename: "screenshot.png", value: ss.response.body }]
       })
     } catch (e) {
       console.log(e)
       return await editOrReply(context, createEmbed("image", context, {
-        url: "https://bignutty.gitlab.io/webstorage4/v2/assets/screenshot/screenshot_error.png"
+        url: "https://bignutty.gitlab.io/webstorage4/v2/assets/screenshot/brand-update-2024/scr_unavailable.png"
       }))
     }
   }
diff --git a/labscore/api/obelisk/endpoints.js b/labscore/api/obelisk/endpoints.js
index e2e3404..10b9e4f 100644
--- a/labscore/api/obelisk/endpoints.js
+++ b/labscore/api/obelisk/endpoints.js
@@ -16,7 +16,9 @@ const ObeliskApi = Object.freeze({
 
   SUMMARIZE_WEBPAGES: "/flamingo/v1/web:summarize",
   
-  AI_WALLPAPER: "/robin/v1/wallpaper:generate"
+  AI_WALLPAPER: "/robin/v1/wallpaper:generate",
+  
+  WEBSHOT: "/peacock/v1/screenshot",
 })
 
 module.exports = {
diff --git a/labscore/api/obelisk/index.js b/labscore/api/obelisk/index.js
index 4092672..94c0c5a 100644
--- a/labscore/api/obelisk/index.js
+++ b/labscore/api/obelisk/index.js
@@ -94,4 +94,12 @@ module.exports.aiWallpaper = async function(context, prompt, style){
     prompt,
     style
   })
+}
+
+// PEACOCK
+module.exports.webshot = async function(context, url, allowAdultContent = false){
+  return await request(ObeliskApi.WEBSHOT, "POST", {}, {
+    url,
+    allow_adult: allowAdultContent
+  })
 }
\ No newline at end of file