radical-bot/ext/libcal.py

import traceback
from hikari.files import Resource

import miru
import hikari
import lightbulb

from typing import Optional

from lib.config import load_config
from lib.gcal import GoogleCalendarAPI
from lib.room import Room
from lib.ocr import get_room_data, NoMatchException

plugin = lightbulb.Plugin("LibCal")
config = load_config()


def add_rooms_to_embed(room_list: list[Room], embed: hikari.Embed):
    for room in room_list:
        embed.add_field(room.get_time_str(), room.number)


class BookingConfirmationView(miru.View):
    def __init__(self, *, timeout: Optional[float] = 120, autodefer: bool = True,
                 room_list: list[Room] = [], image: Resource | None = None, private: bool = False
                 ) -> None:
        self.room_list = room_list
        self.image = image
        self.private = private
        super().__init__(timeout=timeout, autodefer=autodefer)

    @miru.button(label="Yes", style=hikari.ButtonStyle.SUCCESS)
    async def confirm_button(self, button: miru.Button, ctx: miru.Context) -> None:
        gcal = GoogleCalendarAPI()
        for room in self.room_list:
            await gcal.create_event(room.number, room.start_time, room.end_time)
        embed = hikari.Embed(title="Rooms Booked!")
        add_rooms_to_embed(self.room_list, embed)
        embed.set_image(self.image)
        await ctx.edit_response(embed=embed, components=[])

        thread_embed = hikari.Embed(title="Rooms were booked!")
        thread_embed.set_author(
            name=ctx.user.username, icon=ctx.user.avatar_url)
        if not self.private:
            thread_embed.set_image(self.image)
        add_rooms_to_embed(self.room_list, thread_embed)
        await ctx.app.rest.create_message(config.discord_study_room_thread, embed=thread_embed)
        self.stop()

    @miru.button(label="No", style=hikari.ButtonStyle.DANGER)
    async def cancel_button(self, button: miru.Button, ctx: miru.Context) -> None:
        embed = hikari.Embed(title="Aborted!")
        add_rooms_to_embed(self.room_list, embed)
        embed.set_image(self.image)
        await ctx.edit_response(embed=embed, components=[])
        self.stop()


@plugin.command
@lightbulb.option("img", "image to check for events", type=hikari.Attachment, required=True)
@lightbulb.option("private", "if private, the bot will not share confirmation e-mails.", type=bool, default=False)
@lightbulb.command("book", "add room bookings to Google Calendar", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommand)
async def book(ctx: lightbulb.Context) -> None:
    response = hikari.Embed()
    image = await ctx.options.img.read()
    room_list = get_room_data(image)
    add_rooms_to_embed(room_list, response)
    response.title = "Room Bookings"
    response.description = "Are the booking(s) correct?"
    response.set_image(ctx.options.img)
    view = BookingConfirmationView(
        room_list=room_list, image=ctx.options.img, private=ctx.options.private)
    resp = await ctx.respond(response, components=view.build())
    view.start(await resp.message())


@book.set_error_handler
async def book_error_handler(event: lightbulb.CommandErrorEvent):
    print("ERROR TYPE:", type(event.exception))
    if isinstance(event.exception.__cause__, NoMatchException):
        embed = hikari.Embed(title="Huh, I can't read that",
                             description="Could you try another picture?")
        await event.context.respond(embed=embed)
    else:
        embed = hikari.Embed(
            title="Booking Error", description="Whelp, better luck next time I guess... the images used are attached.")
        embed.add_field("Traceback", traceback.format_exc())
        await event.context.respond(embed=embed, attachments=event.context.attachments)


@plugin.command
@lightbulb.command("rooms", "View today's rooms in Google Calendar.", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommand)
async def rooms(ctx: lightbulb.Context) -> None:
    gcal = GoogleCalendarAPI()
    events = await gcal.get_events()

    room_list = [Room.from_event_dict(event) for event in events]
    embed = hikari.Embed(title="Today's Rooms!")
    if len(room_list) == 0:
        embed.description = "Looks like there's no rooms booked today. :("
    for room in room_list[::-1]:
        embed.add_field(
            f"{room.start_time.strftime('%H:%M')} - {room.end_time.strftime('%H:%M')}", room.number)
    await ctx.respond(embed)


def load(bot: lightbulb.BotApp):
    bot.add_plugin(plugin)


def unload(bot: lightbulb.BotApp):
    bot.remove_plugin(plugin)