maubot-tmdb/tmdb/tmdb.py

'''
This file is part of tmdb-bot.

tmdb-bot is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General  Public License version 3 as published by
the Free Software Foundation.

tmdb-bot is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General  Public License for more details.

You should have received a copy of the GNU Affero General  Public License
along with tmdb-bot.  If not, see <https://www.gnu.org/licenses/>.
'''
from html import escape
import re

from mautrix.types import TextMessageEventContent, MediaMessageEventContent, MessageType, Format

from maubot import Plugin, MessageEvent
from maubot.handlers import command

from tmdb.tmdb_api import Movie, TvShow
from tmdb.database import Database


class TmdbBot(Plugin):
    db: Database

    async def start(self) -> None:
        await super().start()
        self.db = Database(self.database)

    async def send_movie_info(self, evt: MessageEvent, movie) -> None:
        mxc_uri = await self.client.upload_media(data=movie.get_image_binary())
        text_message = f'{movie.title}'
        if len(movie.overview) > 200:
            three_dotts = " [...]"
        else:
            three_dotts = ""
        cast = "Acting: "
        for actor in movie.cast[:3]:
            cast += f'{actor}, '
        cast = cast[:-2]
        html_message = f"""<p><b>{escape(movie.title)}</b></p>
        <p>{escape(movie.overview)[:200]}{three_dotts}</p>
        <p>{cast}</p>
        <p>taken from www.themoviedb.org</p>"""

    async def send_html_message(self, evt: MessageEvent, text_message, html_message) -> None:
        content = TextMessageEventContent(
            msgtype=MessageType.TEXT, format=Format.HTML,
            body=f"{text_message}",
            formatted_body=f"{html_message}")
        await evt.respond(content)

    async def send_help(self, evt: MessageEvent) -> None:
        html = """Use <b>!movie-id {tmdb id}</b> to get movie detail for tmdb-id.</br>
                  Use <b>!movie-search {title} [y:{release year}]</b> to get movie detail based on the given title.</br>
                  Use <b>!movie-language {language}</b> to set your prefered language.</br>
                  Use <b>!tvshow-search {title}</b> to get detail about a tv show based on the given title.</br>"""
        content = TextMessageEventContent(
            msgtype=MessageType.TEXT, format=Format.HTML,
            body=f"Help for TMDB Bot",
            formatted_body=f"{html}")
        await evt.respond(content)

    async def send_image(self, evt: MessageEvent, title, image) -> None:
        mxc_uri = await self.client.upload_media(image)
        content = MediaMessageEventContent(
            msgtype=MessageType.IMAGE,
            body=f"Image {title}",
            url=f"{mxc_uri}")
        await evt.respond(content)

    def construct_html_message(self, movie, overview_length = 200, cast_length = 3) -> str:
        if len(movie.overview) > overview_length:
            three_dotts = " [...]"
        else:
            three_dotts = ""

        cast = "Acting: "
        for actor in movie.cast[:cast_length]:
            cast+= f'{actor}, '
        cast = cast[:-2]
        html_message = f"""<p><a href="{movie.web_url}"><b>{escape(movie.title)}</b></a></p>
                           <p>{escape(movie.overview)[:overview_length]}{three_dotts}</p>
                           <p>{cast}</p>
                           <p>Taken from www.themoviedb.org</p>"""
        return html_message

    def split_title_year(self, message : str) -> (str, int):
        m = re.search(r'^(.*) (y:\d\d\d\d)', message)
        if m:
            title = m.group(1)
            year = int(m.group(2)[2:])
            return (title, year)
        return (message, None)

    async def send_movie_info(self, evt: MessageEvent, movie) -> None:
        html_message = self.construct_html_message(movie)
        await self.send_html_message(evt, f'{movie.title}', html_message)
        await self.send_image(evt, movie.title, movie.get_image_binary())

    @command.new("movie-id", help="Movie lookup by id")
    @command.argument("message", pass_raw=True, required=True)
    async def movie_id(self, evt: MessageEvent, message: str = "") -> None:
        movie = Movie()
        language = self.db.get_language(evt.sender)
        if language:
            movie.set_language(language)
        movie.query_details(message)
        await self.send_movie_info(evt, movie)

    @command.new("movie-search", help="Movie lookup by Title")
    @command.argument("message", pass_raw=True, required=True)
    async def movie_search(self, evt: MessageEvent, message: str = "") -> None:
        movie = Movie()
        language = self.db.get_language(evt.sender)
        if language:
            movie.set_language(language)
        title, year = self.split_title_year(message)
        movie.search_title(title, year)
        if movie.valid:
            await self.send_movie_info(evt, movie)
        else:
            content = TextMessageEventContent(
            msgtype=MessageType.NOTICE, format=Format.HTML,
            body=f"No movie found!")
        await evt.respond(content)

    @command.new("movie-language", help="Set language for lookup")
    @command.argument("message", pass_raw=True, required=True)
    async def movie_language(self, evt: MessageEvent, message: str = "") -> None:
        self.db.set_language(evt.sender, message)
        content = TextMessageEventContent(
        msgtype=MessageType.NOTICE, format=Format.HTML,
        body=f"Language set to {message}!")
        await evt.respond(content)        

    @command.new("movie-help", help="Help for TMDB Bot")
    async def movie_help(self, evt: MessageEvent, message: str = "") -> None:
        await self.send_help(evt)

    @command.new("tvshow-help", help="Help for TMDB Bot")
    async def tvshow_help(self, evt: MessageEvent, message: str = "") -> None:
        await self.send_help(evt)

    @command.new("tvshow-search", help="TV Show lookup by Title")
    @command.argument("message", pass_raw=True, required=True)
    async def tvshow_search(self, evt: MessageEvent, message: str = "") -> None:
        movie = TvShow()
        language = self.db.get_language(evt.sender)
        if language:
            movie.set_language(language)
        movie.search_title(message)
        if movie.valid:
            await self.send_movie_info(evt, movie)
        else:
            content = TextMessageEventContent(
            msgtype=MessageType.NOTICE, format=Format.HTML,
            body=f"No tv show found!")
        await evt.respond(content)