From b8d511e3bffeeadff877eb9864cbb688a18b1701 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcos=20Susa=C3=B1a?= <marcosjgs03@gmail.com>
Date: Tue, 12 Mar 2024 00:15:05 -0400
Subject: [PATCH] fix: button builder types

---
 src/builders/Button.ts | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/builders/Button.ts b/src/builders/Button.ts
index d24fae2..bd192c8 100644
--- a/src/builders/Button.ts
+++ b/src/builders/Button.ts
@@ -1,7 +1,6 @@
 import { throwError } from '..';
 import {
 	ComponentType,
-	type APIButtonComponent,
 	type APIButtonComponentWithCustomId,
 	type APIButtonComponentWithURL,
 	type APIMessageComponentEmoji,
@@ -13,6 +12,14 @@ import { resolvePartialEmoji } from '../structures/extra/functions';
 
 export type ButtonStylesForID = Exclude<ButtonStyle, ButtonStyle.Link>;
 
+export type ButtonSelect<T extends ButtonStyle> = T extends ButtonStyle.Link
+	? Omit<Button<false>, 'setCustomId' | 'setURL' | 'setStyle'> & {
+			setStyle(style: ButtonStyle.Link): Omit<Button<false>, 'setCustomId' | 'setURL' | 'setStyle'>;
+	  }
+	: Omit<Button<true>, 'setCustomId' | 'setURL' | 'setStyle'> & {
+			setStyle(style: ButtonStylesForID): Omit<Button<true>, 'setCustomId' | 'setURL' | 'setStyle'>;
+	  };
+
 /**
  * Represents a button component.
  * @template Type - The type of the button component.
@@ -31,10 +38,10 @@ export class Button<Type extends boolean = boolean> {
 	 * @param id - The custom ID to set.
 	 * @returns The modified Button instance.
 	 */
-	setCustomId(id: string) {
+	setCustomId(id: string): ButtonSelect<ButtonStylesForID> {
 		// @ts-expect-error
 		this.data.custom_id = id;
-		return this;
+		return this as ButtonSelect<ButtonStylesForID>;
 	}
 
 	/**
@@ -42,10 +49,10 @@ export class Button<Type extends boolean = boolean> {
 	 * @param url - The URL to set.
 	 * @returns The modified Button instance.
 	 */
-	setURL(url: string) {
+	setURL(url: string): ButtonSelect<ButtonStyle.Link> {
 		// @ts-expect-error
 		this.data.url = url;
-		return this;
+		return this as ButtonSelect<ButtonStyle.Link>;
 	}
 
 	/**
@@ -80,6 +87,8 @@ export class Button<Type extends boolean = boolean> {
 		return this;
 	}
 
+	setStyle(style: ButtonStyle.Link): Omit<Button, 'setCustomId'>;
+	setStyle(style: ButtonStylesForID): Omit<Button, 'setURL'>;
 	setStyle(style: ButtonStyle) {
 		this.data.style = style;
 		return this;
@@ -90,6 +99,6 @@ export class Button<Type extends boolean = boolean> {
 	 * @returns The JSON representation of the Button instance.
 	 */
 	toJSON() {
-		return { ...this.data } as unknown as APIButtonComponent;
+		return { ...this.data } as When<Type, APIButtonComponentWithCustomId, APIButtonComponentWithURL>;
 	}
 }