discord-bot-game/Bot/InviteTracker.fs

module Degenz.InviteTracker

open System.Text
open System.Threading.Tasks
open DSharpPlus
open DSharpPlus.Entities
open DSharpPlus.EventArgs
open DSharpPlus.SlashCommands
open Degenz.Messaging
open Npgsql.FSharp

let connStr = GuildEnvironment.connectionString
let InviteRewardAmount = 100<GBT>
let WhitelistInviteRequirement = 5

type Invite = {
    Code : string
    Inviter : uint64
    Count : int
}

let getInvites () = async {
    let! invites =
        connStr
        |> Sql.connect
        // TODO: Invites shouldn't expire anymore
        |> Sql.query """
                SELECT code, inviter, count FROM invite
                WHERE created_at > (current_timestamp at time zone 'utc') - interval '1 day'
           """
        |> Sql.executeAsync (fun read -> {
                Code = read.string "code"
                Inviter = read.string "inviter" |> uint64
                Count = read.int "count"
            })
        |> Async.AwaitTask
    return
        invites
        |> List.map (fun inv -> (inv.Code , (inv.Inviter , inv.Count)))
        |> Map.ofList
}

let createInvite inviter code =
    connStr
    |> Sql.connect
    |> Sql.parameters [ "code" , Sql.string code ; "inviter" , Sql.string (string inviter) ]
    |> Sql.query "INSERT INTO invite (code, inviter) VALUES (@code, @inviter)"
    |> Sql.executeNonQueryAsync
    |> Async.AwaitTask

let addInvitedUser did code count =
    try
        connStr
        |> Sql.connect
        |> Sql.executeTransactionAsync [
            """
                INSERT INTO invited_user (discord_id, invite_id)
                VALUES (@did, (SELECT id FROM invite WHERE code = @code));
            """ , [ [ "@code" , Sql.string code ; "@did" , Sql.string (string did) ] ]
            "UPDATE invite SET count = @count WHERE code = @code" , [ [ "count" , Sql.int count ; "code" , Sql.string code ] ]
        ]
        |> Async.AwaitTask
        |> Async.Ignore
    with _ -> async.Zero ()

let acceptInvite did =
    connStr
    |> Sql.connect
    |> Sql.parameters [ "did" , Sql.string (string did) ]
    |> Sql.query "UPDATE invited_user SET accepted = true WHERE discord_id = @did"
    |> Sql.executeNonQueryAsync
    |> Async.AwaitTask

let removeInvitedUser did =
    try
        connStr
        |> Sql.connect
        |> Sql.parameters [ "did" , Sql.string (string did) ]
        |> Sql.query "DELETE FROM invited_user WHERE discord_id = @did"
        |> Sql.executeNonQueryAsync
        |> Async.AwaitTask
        |> Async.Ignore
    with _ -> async.Zero ()

let checkUserInvited userId = async {
    let! result =
        connStr
        |> Sql.connect
        |> Sql.parameters [ "did" , Sql.string (string userId) ]
        |> Sql.query """
                SELECT id FROM invited_user WHERE discord_id = @did
            """
        |> Sql.executeAsync (fun read -> read.int "id")
        |> Async.AwaitTask

    return List.isEmpty result |> not
}

let getInviteAttributions userId =
    connStr
    |> Sql.connect
    |> Sql.parameters [ "did" , Sql.string (string userId) ]
    |> Sql.query """
            SELECT count(*) FROM invited_user
            JOIN invite ON invite.id = invited_user.invite_id
            WHERE invite.inviter = @did AND invited_user.accepted = true;
        """
    |> Sql.executeRowAsync (fun read -> read.int "count")
    |> Async.AwaitTask

let getInvitedUsers userId =
    connStr
    |> Sql.connect
    |> Sql.parameters [ "did" , Sql.string (string userId) ]
    |> Sql.query """
            WITH invite AS (SELECT id FROM invite WHERE inviter = @did)
            SELECT discord_id FROM invited_user, invite
            WHERE invite.id = invited_user.invite_id;
        """
    |> Sql.executeAsync (fun read -> read.string "discord_id" |> uint64)
    |> Async.AwaitTask

let createGuildInvite (ctx : IDiscordContext) =
    task {
        do! ctx.Respond(InteractionResponseType.DeferredChannelMessageWithSource)

        let channel = ctx.GetGuild().Channels.[GuildEnvironment.channelWelcome]
        let! invite = channel.CreateInviteAsync(max_age = 0, unique = true)

        // When a player generates an invite code but it hasn't expired, it generates the same code, creating a duplicate entry
        // so catch the exception thrown because the code column is unique
        try
            let! _ = createInvite (ctx.GetDiscordMember().Id) invite.Code
            return ()
        with ex ->
            printfn "%A" ex.Message
            ()

        let header =
            DiscordEmbedBuilder()
                .WithDescription($"**⏀ | Your Mission:**\nCopy the link & share it with Degenz you want to `/recruit`.\n\n" +
                                 $"**⌼ | Your Reward:**\n`Earn {InviteRewardAmount} $GBT` 💰  for every Degen you've `/recruited`.")
                .WithImageUrl("https://pbs.twimg.com/profile_banners/1449270642340089856/1640071520/1500x500")
//                .WithImageUrl("https://s7.gifyu.com/images/Resistance_Poster_Final_1.jpg")
                .WithTitle("Recruitment")

        let msg =
            DiscordFollowupMessageBuilder()
                .AddEmbed(header.Build())
                .AsEphemeral(true)

        do! ctx.FollowUp(msg)

        let msg =
            DiscordFollowupMessageBuilder()
                .WithContent($"https://discord.gg/{invite.Code}")
                .AsEphemeral(true)

        do! ctx.FollowUp(msg)
    }

let listServerInvites (ctx : IDiscordContext) = task {
    let! invites = ctx.GetGuild().GetInvitesAsync()
    let sb = StringBuilder()
    for invite in invites do
        sb.AppendLine($"{invite.Inviter.Username} - {invite.Code}") |> ignore
    let msg =
        DiscordInteractionResponseBuilder()
            .AsEphemeral(true)
            .WithContent("Server Invites\n" + sb.ToString())
    do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, msg)
}

let getAttributions (ctx : IDiscordContext) userId = task {
    let! total = getInviteAttributions(userId)
    let msg =
        DiscordInteractionResponseBuilder()
            .AsEphemeral(true)
            .WithContent($"<@{userId}> has invited {total} people")
    do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, msg)
}

let getInvitedUsersForId (ctx : IDiscordContext) = task {
    let! users = getInvitedUsers(ctx.GetDiscordMember().Id)
    let sb = StringBuilder()
    let mutable count = 0
    for user in users do
        count <- count + 1
        sb.AppendLine($"{count}.) <@{user}>") |> ignore
    let msg =
        let str =
            if users.Length > 0 then
                $"**Total Recruited:** `{users.Length} Degenz`\n**Total Earned:** `{users.Length * InviteRewardAmount} 💰$GBT`\n\n**Last 10 users recruited:**\n{sb}"
            else
                $"You haven't recruited anyone yet, use the `/recruit` command to get the recruitment link"
        DiscordInteractionResponseBuilder()
            .AsEphemeral(true)
            .WithContent(str)
    do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, msg)
}

let clearInvites (ctx : IDiscordContext) = task {
    let! invites = ctx.GetGuild().GetInvitesAsync()
    do!
        invites
        |> Seq.map (fun invite -> invite.DeleteAsync() |> Async.AwaitTask)
        |> Async.Parallel
        |> Async.Ignore
}

let processNewUser (eventArgs : GuildMemberAddEventArgs) =
    task {
         let! guildInvites = eventArgs.Guild.GetInvitesAsync()
         let! cachedInvites = getInvites()
         for invite in guildInvites do
             let result = cachedInvites.TryFind(invite.Code)
             match result with
             | Some (inviter,count) ->
                 if invite.Uses > count then
                     do! addInvitedUser eventArgs.Member.Id invite.Code invite.Uses |> Async.Ignore
                     let! _ = acceptInvite eventArgs.Member.Id
                     let! player = DbService.tryFindPlayer inviter
                     match player with
                     | Some player ->
                         do! DbService.updatePlayerCurrency (int InviteRewardAmount) player |> Async.Ignore
                         let builder = DiscordMessageBuilder()
                         builder.WithContent($"{eventArgs.Member.DisplayName} was recruited to the server. <@{player.DiscordId}> just earned {InviteRewardAmount} 💰$GBT for their efforts!") |> ignore
                         let channel = eventArgs.Guild.GetChannel(GuildEnvironment.channelEventsHackerBattle)
                         do! channel.SendMessageAsync(builder)
                             |> Async.AwaitTask
                             |> Async.Ignore
                     | None -> return ()
             | None -> ()
    } :> Task

let handleGuildMemberAdded _ (eventArgs : GuildMemberAddEventArgs) =
    task {
        let! exists = checkUserInvited eventArgs.Member.Id
        if not exists then
            do! processNewUser eventArgs
    } :> Task

// If we do it like this then there's an obvious exploit where the user can come and go as many times and it will keep
// rewarding GBT.
//let handleGuildMemberRemoved _ (eventArgs : GuildMemberRemoveEventArgs) =
//    task {
//        do! removeInvitedUser eventArgs.Member.Id
//    } :> Task

let sendInitialEmbed (client : DiscordClient)  =
    async {
        try
            let! channel = client.GetChannelAsync(GuildEnvironment.channelWhitelist) |> Async.AwaitTask
            let builder = DiscordMessageBuilder()
            let embed = DiscordEmbedBuilder()
            embed.ImageUrl <- "https://securitygladiators.com/wp-content/uploads/2020/09/Whitelist-Website-Featured-Image.jpg"
            builder.AddEmbed embed |> ignore
            builder.Content <- "Click on the button to get whitelist!"
            let button = DiscordButtonComponent(ButtonStyle.Success, $"Trainer-1", $"Gimme") :> DiscordComponent
            builder.AddComponents [| button |] |> ignore
            do! channel.SendMessageAsync(builder)
                |> Async.AwaitTask
                |> Async.Ignore
        with e ->
            printfn $"Error trying to get channel {GuildEnvironment.channelTraining}\n\n{e.Message}"
    } |> Async.RunSynchronously

type WhitelistResult =
    | NotEnoughInvites of currentAmount : int
    | Granted of DiscordRole
    | AlreadyWhitelisted

let tryGrantWhitelist (ctx : IDiscordContext) =
    async {
        let user = ctx.GetDiscordMember()
        let role = ctx.GetGuild().GetRole(GuildEnvironment.roleWhitelist)
        if Seq.contains role user.Roles
            then return AlreadyWhitelisted
            else
                let! total = getInviteAttributions user.Id
                if total >= WhitelistInviteRequirement then
                    return Granted role
                else
                    return NotEnoughInvites total
    }

let handleWhitelist (_ : DiscordClient) (event : ComponentInteractionCreateEventArgs) =
    task {
        let ctx = DiscordEventContext event :> IDiscordContext
        let builder = DiscordInteractionResponseBuilder().AsEphemeral(true)
        match! tryGrantWhitelist ctx with
        | AlreadyWhitelisted ->
            builder.Content <- "You are already whitelisted"
        | NotEnoughInvites total ->
            builder.Content <- $"You need to invite more people into the server. Please invite {WhitelistInviteRequirement - total} more people to get Whitelist"
        | Granted role ->
            do! ctx.GetDiscordMember().GrantRoleAsync(role)
            builder.Content <- "You have been granted whitelist"
        do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, builder)
    } :> Task

type Inviter() =
    inherit ApplicationCommandModule ()

    [<SlashCommand("recruit", "Recruit another user to this discord and earn rewards")>]
    member this.CreateInvite (ctx : InteractionContext) =
        createGuildInvite (DiscordInteractionContext ctx)

    [<SlashCommand("recruited", "Get total invites from a specific user")>]
    member this.ListInvitedPeople (ctx : InteractionContext) =
        getInvitedUsersForId (DiscordInteractionContext ctx)

//    [<SlashCommand("invites-list", "List all the invites")>]
    member this.ListServerInvites (ctx : InteractionContext) =
        listServerInvites (DiscordInteractionContext ctx)

//    [<SlashCommand("invites-attributions", "Get total invites from a specific user")>]
    member this.getAttributions (ctx : InteractionContext, [<Option("player", "The player you want to check")>] user : DiscordUser) =
        getAttributions (DiscordInteractionContext ctx) user.Id

//    [<SlashCommand("invites-clear", "Get total invites from a specific user")>]
    member this.ClearInvites (ctx : InteractionContext) =
        clearInvites (DiscordInteractionContext ctx)