From e5039e7d7bb7d248deedafd3d35bf00786379d50 Mon Sep 17 00:00:00 2001
From: Atridad Lahiji <88056492+atridadl@users.noreply.github.com>
Date: Fri, 7 Jul 2023 12:06:55 -0600
Subject: [PATCH] Quick change to commands

---
 src/commands/ask.ts     |  2 +-
 src/commands/credits.ts | 50 +++++++++++++++++++++++++++++++++++++++++
 src/commands/pic.ts     | 22 +++++++++++++-----
 3 files changed, 68 insertions(+), 6 deletions(-)
 create mode 100644 src/commands/credits.ts

diff --git a/src/commands/ask.ts b/src/commands/ask.ts
index 0c29d90..46ff6c5 100644
--- a/src/commands/ask.ts
+++ b/src/commands/ask.ts
@@ -27,7 +27,7 @@ export class UserCommand extends Command {
 
 	// Message command
 	public async messageRun(message: Message, args: Args) {
-		return this.ask(message, args.getOption('prompt') || message.content.split('!wryna ')[1]);
+		return this.ask(message, args.getOption('prompt') || message.content.split('!ask ')[1]);
 	}
 
 	// Chat Input (slash) command
diff --git a/src/commands/credits.ts b/src/commands/credits.ts
new file mode 100644
index 0000000..a3fa0bf
--- /dev/null
+++ b/src/commands/credits.ts
@@ -0,0 +1,50 @@
+import { ApplyOptions } from '@sapphire/decorators';
+import { Command } from '@sapphire/framework';
+import { Message } from 'discord.js';
+
+@ApplyOptions<Command.Options>({
+	description: 'Make a picture!'
+})
+export class UserCommand extends Command {
+	// Register Chat Input and Context Menu command
+	public override registerApplicationCommands(registry: Command.Registry) {
+		registry.registerChatInputCommand((builder) => builder.setName(this.name).setDescription(this.description));
+	}
+
+	// Message command
+	public async messageRun(message: Message) {
+		return this.credits(message);
+	}
+
+	// Chat Input (slash) command
+	public async chatInputRun(interaction: Command.ChatInputCommandInteraction) {
+		return this.credits(interaction);
+	}
+
+	private async credits(interactionOrMessage: Message | Command.ChatInputCommandInteraction | Command.ContextMenuCommandInteraction) {
+		const askMessage =
+			interactionOrMessage instanceof Message
+				? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' })
+				: await interactionOrMessage.reply({ content: '🤔 Thinking... 🤔', fetchReply: true });
+
+		const creditCountResponse = await fetch(`https://api.stability.ai/v1/user/balance`, {
+			method: 'GET',
+			headers: {
+				'Content-Type': 'application/json',
+				Authorization: `Bearer ${process.env.STABILITY_API_KEY}`
+			}
+		});
+
+		const balance = (await creditCountResponse.json()).credits || 0;
+
+		const content = `I have ${balance} credits remaining for image generation!`;
+
+		if (interactionOrMessage instanceof Message) {
+			return askMessage.edit({ content });
+		}
+
+		return interactionOrMessage.editReply({
+			content: content
+		});
+	}
+}
diff --git a/src/commands/pic.ts b/src/commands/pic.ts
index 34def63..4bb802a 100644
--- a/src/commands/pic.ts
+++ b/src/commands/pic.ts
@@ -2,6 +2,8 @@ import { ApplyOptions } from '@sapphire/decorators';
 import { Args, BucketScope, Command } from '@sapphire/framework';
 import { AttachmentBuilder, Message } from 'discord.js';
 
+// This is literally the worlds messiest TS code. Please don't judge me...
+
 @ApplyOptions<Command.Options>({
 	description: 'Make a picture!',
 	options: ['prompt'],
@@ -25,7 +27,7 @@ export class UserCommand extends Command {
 
 	// Message command
 	public async messageRun(message: Message, args: Args) {
-		return this.pic(message, args.getOption('prompt') || message.content.split('!wryna ')[1]);
+		return this.pic(message, args.getOption('prompt') || message.content.split('!pic ')[1]);
 	}
 
 	// Chat Input (slash) command
@@ -94,11 +96,21 @@ export class UserCommand extends Command {
 				const responseJSON = (await imageGenResponse.json()) as GenerationResponse;
 				const imageAttachment = new AttachmentBuilder(Buffer.from(responseJSON.artifacts[0].base64, 'base64'));
 
+				const newCreditCountResponse = await fetch(`https://api.stability.ai/v1/user/balance`, {
+					method: 'GET',
+					headers: {
+						'Content-Type': 'application/json',
+						Authorization: `Bearer ${process.env.STABILITY_API_KEY}`
+					}
+				});
+
+				const newBalance = (await newCreditCountResponse.json()).credits || 0;
+
 				const content =
-					`Prompt: ${prompt}${
-						balance &&
-						balance <= 300 &&
-						`\n\n⚠️I am now at ${balance} credits. If you'd like to help fund this command, please type "/support" for details!`
+					`Credits Used: ${balance - newBalance}\nPrompt: ${prompt}${
+						balance <= 300
+							? `\n\n⚠️I am now at ${balance} credits. If you'd like to help fund this command, please type "/support" for details!`
+							: ''
 					}` || 'ERROR!';
 
 				if (interactionOrMessage instanceof Message) {