import mysql.connector
from config import DATABASE

class Database:
    def __init__(self):
        self.connection   = mysql.connector.connect(**DATABASE)
        self.db = self.connection.cursor(buffered=True)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.db.close()
        self.connection.close()

    def add_user(self, chat_id: int):
        sql = "INSERT INTO users (chat_id) VALUES (%s)"
        self.db.execute(sql, (chat_id,))
        self.connection.commit()
        return

    def get_user(self, chat_id: int, get: str = '*'):
        sql = "SELECT {} FROM users WHERE chat_id=%s".format(get)
        self.db.execute(sql, (chat_id,))
        return self.db.fetchall()

    def update_user(self, chat_id: int, column: str, entry):
        sql = "UPDATE users SET {} = %s WHERE chat_id = %s".format(column)
        self.db.execute(sql, (entry, chat_id,))
        self.connection.commit()
        return

    def list_admins(self):
        sql = "SELECT chat_id FROM users"
        self.db.execute(sql)
        user_ids = self.db.fetchall()
        return [user_id[0] for user_id in user_ids]

    def add_subtitle(self, file_name: str):
        sql = "INSERT INTO subtitles (file_name) VALUES (%s)"
        self.db.execute(sql, (file_name,))
        last_id = self.db.lastrowid
        self.connection.commit()
        return last_id

    def get_subtitle(self, id: int, get: str = '*'):
        sql = "SELECT {} FROM subtitles WHERE id=%s".format(get)
        self.db.execute(sql, (id,))
        return self.db.fetchall()

    def add_media_task(self, package_id: int, message_id: int, file: str, file_name: str, file_size: int):
        sql = "INSERT INTO media_tasks (package_id,message_id,file,file_name,file_size) VALUES (%s,%s,%s,%s,%s)"
        self.db.execute(sql, (package_id,message_id,file,file_name,file_size))
        last_id = self.db.lastrowid
        self.connection.commit()
        return last_id

    def get_media_task(self, id: int, search_column: str, get: str = '*', order: str = 'ORDER BY id ASC', op: str = '='):
        sql = "SELECT {} FROM media_tasks WHERE {}{}%s {}".format(get,search_column,op,order)
        self.db.execute(sql, (id,))
        return self.db.fetchall()

    def update_media_task(self, id: int, column: str, entry):
        sql = "UPDATE media_tasks SET {} = %s WHERE id = %s".format(column)
        self.db.execute(sql, (entry, id,))
        self.connection.commit()
        return

    def delete_media_task(self,id: int, column: str, order = str()):
        sql = "DELETE FROM `media_tasks` WHERE {} = %s {}".format(column,order)
        self.db.execute(sql, (id,))
        self.connection.commit()
        return

    def add_package(self, chat_id: int, type: int):
        sql = "INSERT INTO task_packages (chat_id,type) VALUES (%s,%s)"
        self.db.execute(sql, (chat_id,type))
        last_id = self.db.lastrowid
        self.connection.commit()
        return last_id

    def get_package(self, id: int, search_column: str, get: str = '*', order: str = 'ORDER BY id ASC'):
        sql = "SELECT {} FROM task_packages WHERE {}=%s {}".format(get,search_column,order)
        self.db.execute(sql, (id,))
        return self.db.fetchall()

    def update_package(self, id: int, column: str, entry):
        sql = "UPDATE task_packages SET {} = %s WHERE id = %s".format(column)
        self.db.execute(sql, (entry, id,))
        self.connection.commit()
        return

    def add_sub(self, user_id: int, file_name: str):
        sql = "INSERT INTO subtitles (chat_id,file_name) VALUES (%s,%s)"
        self.db.execute(sql, (user_id,file_name))
        last_id = self.db.lastrowid
        self.connection.commit()
        return last_id

    def get_sub(self, id: int):
        sql = "SELECT * FROM subtitles WHERE id=%s"
        self.db.execute(sql, (id,))
        return self.db.fetchall()