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 getInvitesFromUser discordId = async {
    let! invites =
        connStr
        |> Sql.connect
        |> Sql.parameters [ "did" , Sql.string (string discordId) ]
        |> Sql.query """
                SELECT code, count FROM invite
                WHERE inviter = @did
           """
        |> Sql.executeAsync (fun read -> {
                Code = read.string "code"
                Inviter = discordId
                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) showWhitelistReward =
    task {
        let invitesRequired = 5
        let rewardMsg =
            if showWhitelistReward then
                $"**⏀ | Your Mission:**\nUse the link to **recruit** {invitesRequired} Degenz into Beautopia©.\n\n" +
                $"**⌼ | Your Reward:**\nGet awarded **Whitelist**."
            else
                $"**⏀ | Your Mission:**\nUse the link to share it with Degenz you want to `/recruit`.\n\n" +
                $"**⌼ | Your Reward:**\n`Earn {InviteRewardAmount} $GBT` 💰 for every Degen you've `/recruited`."
        let embed =
            DiscordEmbedBuilder()
                .WithDescription(rewardMsg)
                .WithImageUrl("https://pbs.twimg.com/profile_banners/1449270642340089856/1640071520/1500x500")
                .WithTitle("Recruitment")

        let builder =
            DiscordInteractionResponseBuilder()
                .AddEmbed(embed)
                .AsEphemeral(true)

        let button = DiscordButtonComponent(ButtonStyle.Success, $"CreateGuildInvite", $"Give it to me") :> DiscordComponent
        builder.AddComponents [| button |] |> ignore
        do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, builder)

    } :> Task

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

// 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, $"GimmeWhitelist", $"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 (ctx : IDiscordContext) =
    task {
        let builder = DiscordInteractionResponseBuilder().AsEphemeral(true)
        match! tryGrantWhitelist ctx with
        | AlreadyWhitelisted ->
            builder.Content <- "You are already whitelisted"
            do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, builder)
        | NotEnoughInvites total ->
            do! createGuildInvite ctx true
//            builder.Content <- "Testing"
//            do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, builder)
//            do! async.Zero()
        | Granted role ->
            do! ctx.GetDiscordMember().GrantRoleAsync(role)
            builder.Content <- "You have been granted whitelist"
            do! ctx.Respond(InteractionResponseType.ChannelMessageWithSource, builder)
    } :> Task

let handleCreateInvite (ctx : IDiscordContext) =
    task {
        let builder = DiscordInteractionResponseBuilder().AsEphemeral(true)
        do! ctx.Respond(InteractionResponseType.DeferredChannelMessageWithSource, builder)

        let user = ctx.GetDiscordMember()

        let! code =
            task {
                let! invites = getInvitesFromUser user.Id
                match invites |> Map.toList with
                | [] ->
                    let channel = ctx.GetGuild().Channels.[GuildEnvironment.channelWelcome]
                    let! invite = channel.CreateInviteAsync(max_age = 0, unique = true)

                    try do! createInvite (ctx.GetDiscordMember().Id) invite.Code |> Async.Ignore
                    with ex -> printfn "%A" ex.Message

                    return invite.Code
                | invite::_ ->
                    return invite |> fst
            }

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

        do! ctx.FollowUp(msg)
    }

let handleButtonEvent  (_ : DiscordClient) (event : ComponentInteractionCreateEventArgs) =
    let eventCtx = DiscordEventContext event :> IDiscordContext
    match event.Id with
    | id when id.StartsWith("GimmeWhitelist")  -> handleWhitelist eventCtx
    | id when id.StartsWith("CreateGuildInvite")  -> handleCreateInvite eventCtx
    | _ ->
        task {
            let builder = DiscordInteractionResponseBuilder()
            builder.IsEphemeral <- true
            builder.Content <- $"Incorrect Action identifier {eventCtx.GetInteractionId()}"
            do! eventCtx.Respond(InteractionResponseType.ChannelMessageWithSource, builder) |> Async.AwaitTask
        }

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

type Inviter() =
    inherit ApplicationCommandModule ()

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

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