From 51f99127874dba307918a9dc3567a8c55138bcb9 Mon Sep 17 00:00:00 2001
From: lomion <lomion@aon.at>
Date: Mon, 22 Jun 2020 14:41:31 +0200
Subject: [PATCH] Created database handler

---
 maubot.yaml      |  3 ++-
 tmdb/tmdb.py     | 53 ++++++++++++++++++++++++++++++++++++++++++++++++
 tmdb/tmdb_api.py |  9 ++++++--
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/maubot.yaml b/maubot.yaml
index 9dda137..31d9324 100644
--- a/maubot.yaml
+++ b/maubot.yaml
@@ -1,7 +1,8 @@
 maubot: 0.1.0
 id: lomion.tmdb
-version: 0.0.1
+version: 0.0.2
 license: AGPL 3.0
 modules:
 - tmdb
 main_class: TmdbBot
+database: true
diff --git a/tmdb/tmdb.py b/tmdb/tmdb.py
index 229b4e7..cbba73b 100644
--- a/tmdb/tmdb.py
+++ b/tmdb/tmdb.py
@@ -7,8 +7,17 @@ from maubot.handlers import command
 
 from tmdb.tmdb_api import Movie
 
+from sqlalchemy import (Column, String, Integer, ForeignKey, Table, MetaData,
+                        select, and_)
+from sqlalchemy.engine.base import Engine
+
 class TmdbBot(Plugin):
 
+
+    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}'
@@ -48,6 +57,13 @@ class TmdbBot(Plugin):
     @command.argument("message", pass_raw=True, required=True)
     async def movie_search(self, evt: MessageEvent, message: str = "") -> None:
         movie = Movie()
+        self.db = Database(self.database)
+        language = self.db.get_language(evt.sender)
+        if language:
+            movie.set_language(language)
+        else:
+            self.db.set_language(evt.sender, 'de')
+            movie.set_language('de')
         movie.search_title(message)
         if movie.valid:
             await self.send_movie_info(evt, movie)
@@ -56,3 +72,40 @@ class TmdbBot(Plugin):
             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 = Database(self.database)
+        self.db.set_language(evt.sender, message)
+        content = TextMessageEventContent(
+        msgtype=MessageType.NOTICE, format=Format.HTML,
+        body=f"Language set to {language}!")
+        await evt.respond(content)        
+
+class Database:
+    def __init__(self, db: Engine) -> None:
+        self.db = db
+        
+        meta = MetaData()
+        meta.bind = db
+
+        self.language = Table("tmdb_language", meta,
+                              Column("id", Integer, primary_key=True, autoincrement=True),
+                              Column("user_id", String(255), nullable=False),
+                              Column("language", String(255), nullable=False),)
+        meta.create_all()
+        
+    def set_language(self, user_id, language):
+        with self.db.begin() as tx:
+            tx.execute(self.language.delete().where(self.language.c.user_id == user_id))
+            tx.execute(self.language.insert().values(user_id=user_id, language=language))
+            
+    def get_language(self, user_id):
+        rows = self.db.execute(select([self.language.c.language])
+                                   .where(self.language.c.user_id == user_id))
+        if len(rows) == 1:
+            row = rows[0]
+            return row[2]
+        else:
+            return None
diff --git a/tmdb/tmdb_api.py b/tmdb/tmdb_api.py
index f819671..03e9740 100644
--- a/tmdb/tmdb_api.py
+++ b/tmdb/tmdb_api.py
@@ -11,6 +11,8 @@ class Connection():
         result = requests.get(self.base_url + 'configuration', params = self.get_apikey()).json()
         self.base_url_poster = result['images']['base_url'] + result['images']['poster_sizes'][0]
         self.valid = False
+        
+        self.language = 'de'
 
     def get_apikey(self):
         return { 'api_key' : self.api_key }
@@ -18,11 +20,14 @@ class Connection():
     def request(self, request_uri):
         url = self.base_url + request_uri
         payload = self.get_apikey()
-        payload['language'] = 'de'
+        payload['language'] = self.language
         result = requests.get(url, params=payload)
         self.valid = True
         return result
 
+    def set_language(self, language):
+        self.language = language
+
     def get_image_binary(self):
         return requests.get(self.poster_url).content
 
@@ -35,7 +40,7 @@ class Movie(Connection):
     def search_title(self, title):
         url = self.base_url+ 'search/movie'
         payload = self.get_apikey()
-        payload['language'] = 'de'
+        payload['language'] = self.language
         payload['query'] = title
         result = requests.get(url, params=payload)
         json = result.json()