app.py

import os
import discord
import secrets
import threading
import gradio as gr

from discord.ext import commands
from discord.ui import Button, View
from urllib.parse import urlparse, parse_qs


# Discord bot ------------------------------------------------------------------------------------------------------
intents = discord.Intents.all()
bot = commands.Bot(command_prefix="!", intents=intents)
GRADIO_APP_URL = "https://huggingface.co/spaces/lunarflu/gradio-oauth2"
DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)
BUTTON_CHANNEL_ID = "1106995001155649680"
BUTTON_MESSAGE_ID = "1266438263531634708"

# Dictionary to store user IDs and their corresponding unique strings
user_tokens = {}

@bot.event
async def on_ready():
    print(f'Logged in as {bot.user}')

    if BUTTON_MESSAGE_ID:
        channel = bot.get_channel(932563860597121054)
        if channel:
            try:
                message = await channel.fetch_message(1269917762327674974)
                if message:
                    button = DMButton(label="Verify Discord Account", style=discord.ButtonStyle.primary)
                    view = View(timeout=None)
                    view.add_item(button)
                    await message.edit(view=view)
                    print("message edited")
            except discord.NotFound:
                print(f"Message with ID {KNOWN_MESSAGE_ID} not found.")
            except discord.HTTPException as e:
                print(f"Failed to fetch message with ID {KNOWN_MESSAGE_ID}: {e}")
                

def generate_unique_string(length=6):
    return secrets.token_hex(length // 2)


def run_bot():
    bot.run(DISCORD_TOKEN)

threading.Thread(target=run_bot).start()
# commands ---------------------------------------------------------------------------------------------------------
@bot.command()
async def sendlink(ctx, user: discord.User):
    if ctx.author.id == 811235357663297546:
        unique_string = generate_unique_string()
        user_tokens[user.id] = unique_string
        unique_link = f"{GRADIO_APP_URL}?user_id={user.id}&token={unique_string}"
        await user.send(f"Click the link to sign in with Hugging Face: {unique_link}")


class DMButton(Button):
    def __init__(self, label, style):
        super().__init__(label=label, style=style)

    async def callback(self, interaction: discord.Interaction):
        # await interaction.user.send(self.message) # this is for DMs, but users may have DMs disabled
        user_id = interaction.user.id
        if int(user_id) in user_tokens:
            del user_tokens[int(user_id)] # always delete all past tokens when creating new link
        unique_string = generate_unique_string()
        user_tokens[user_id] = unique_string
        unique_link = f"<{GRADIO_APP_URL}?user_id={user_id}&token={unique_string}>"
        message = f"Login link generated! To complete the verification process:\n- 1 Visit this link,\n- 2 click the '🤗Sign in with Hugging Face' button\n\n{unique_link}"
        await interaction.response.send_message(message, ephemeral=True)


@bot.command(name='sendbutton')
async def send_button(ctx):
    if ctx.author.id == 811235357663297546:
        button = DMButton(label="Verify Discord Account", style=discord.ButtonStyle.primary)
        view = View(timeout=None)
        view.add_item(button)
        await ctx.send("Click the button below to generate your verification link:",view=view) #  


# Gradio ------------------------------------------------------------------------------------------------------------
def hello(profile: gr.OAuthProfile | None, request: gr.Request) -> str:
    url_str = str(request.url)
    query_params = parse_qs(urlparse(url_str).query)
    user_id = query_params.get('user_id', [None])[0]
    token = query_params.get('token', [None])[0]

    print(f"||| token:{token}||| user_id:{user_id}||| profile:{profile}||| user_tokens:{user_tokens}")

    if user_id is None or token is None:
        return "# ❌ Invalid link. Generate a new link [here](https://discord.com/channels/879548962464493619/900125909984624713) !"
        
    if int(user_id) not in user_tokens or user_tokens[int(user_id)] != token:
        return "# ❌ Invalid or expired link. Generate a new link [here](https://discord.com/channels/879548962464493619/900125909984624713) !"

    if profile is None:
        return f"# ❌ Not logged in with Hugging Face yet."

    # Remove the token after successful verification
    del user_tokens[int(user_id)]

    # profile.username
    return f"# ✅ Verification successful! We have linked your Hugging Face and Discord accounts. You can now earn exp for activity on Hugging Face as well as on Discord!"


with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# Discord Verification Space")
    with gr.Row():
        login_button = gr.LoginButton()

    m1 = gr.Markdown()
    demo.load(hello, inputs=None, outputs=m1)

    def check_login_status():
        try:
            return login_button.get_session().get("oauth_info", None)
        except AttributeError:
            return None

    def check_login_wrapper():
        session = check_login_status()
        if session is None:
            return "Not logged in."
        else:
            return f"Logged in as {session.get('username', 'Unknown')}"

    login_button.click(check_login_wrapper, inputs=None, outputs=m1)

demo.launch()


# nice version! :D