import gradio as gr
import os
import threading
from urllib.parse import urlparse, parse_qs
import discord
from discord.ext import commands
import secrets

# 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)

# 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}')

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

@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}")

def run_bot():
    bot.run(DISCORD_TOKEN)

threading.Thread(target=run_bot).start()

# Gradio ------------------------------------------------------------------------------------------------------------
def validate_link(request: gr.Request) -> bool:
    referer = request.headers.get('Referer', '')
    url_str = str(referer)
    query_params = parse_qs(urlparse(url_str).query)
    user_id = query_params.get('user_id', [None])[0]
    token = query_params.get('token', [None])[0]

    if user_id is None or token is None:
        return False

    if int(user_id) not in user_tokens or user_tokens[int(user_id)] != token:
        return False

    return True

def hello(profile: gr.OAuthProfile | None, request: gr.Request) -> str:
    if not validate_link(request):
        return "❌ Invalid link. Please use the link provided by the bot."

    referer = request.headers.get('Referer', '')
    url_str = str(referer)
    query_params = parse_qs(urlparse(url_str).query)
    user_id = query_params.get('user_id', [None])[0]

    if profile is None:
        return f"❌ Not logged in. User ID: {user_id}"

    return f"✅ Successfully logged in as {profile.username}. User ID: {user_id}"

with gr.Blocks() as demo:
    gr.Markdown("# Gradio OAuth Space")
    login_button = gr.LoginButton()
    m1 = gr.Markdown()

    def check_link_and_load():
        if validate_link(gr.Request()):
            demo.load(hello, inputs=None, outputs=m1)
        else:
            m1.update("❌ Invalid link. Please use the link provided by the bot.")

    check_link_and_load()

    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()