from nio import *


latex_regex = r'((?:\$[^\$]+\$)|(?:\$\$[^\$]+\$\$)|(?:\\\[[^\]]+\\\]))'

# Our message callback. It should be passed through a router.
async def msg_cb(room: MatrixRoom, event: RoomMessageText) -> None:
    import latex
    import utils
    import shlex
    import re

    client = utils.get_client()
    if event.sender == client.user_id:
        return
    
    if event.body.startswith("tex!"):
        import router
        args = shlex.split(event.body[4:])
        if len(args) >= 1:
            await router.handle_command(args[0], room.room_id, event, args[1:])
    
    for tex in re.findall(latex_regex, event.body, re.M):
        try:
            filename = latex.render(event.sender, tex)
            await utils.send_png(room, filename)
        except FileNotFoundError as e:
            content = {
                "msgtype": "m.text",
                "body": f"Couldn't parse LaTeX correctly.\n```{e.args[0]}\n```",
                "formatted_body": f"Couldn't parse LaTeX correctly.<br><code><pre>{e.args[0]}</pre></code>",
            }
            await client.room_send(room.room_id, message_type="m.room.message", content=content)
        except OSError:
            content = {
                "msgtype": "m.text",
                "body": "???",
                "formatted_body": "???"
            }

            await client.room_send(room.room_id, message_type="m.room.message", content=content)

# Our file callback.
async def file_cb(room: MatrixRoom, event: RoomMessageFile) -> None:
    import utils
    import os

    url = event.url
    filename = event.body

    if filename.endswith(".sty"):
        # Download file and save it for user.
        client = utils.get_client()
        directory = utils.create_user_dir(event.sender)
        response = await client.download(url, filename)
        
        if not isinstance(response, DownloadResponse):
            return

        with open(os.path.join(directory, filename), 'wb') as f:
            f.write(response.body)