import os

import pymysql


class DB:
    def __init__(self, host=None, port=None, user=None, password=None, db_name=None, ssl_ca=None):
        self.host = host or os.getenv('DB_HOST')
        self.port = port or int(os.getenv('DB_PORT'))
        self.user = user or os.getenv('DB_USERNAME')
        self.password = password or os.getenv('DB_PASSWORD')
        self.db_name = db_name or os.getenv('DB_DATABASE')
        self.ssl_ca = ssl_ca or os.getenv('DB_SSL_CA')
        self.connection = None
        # print all vars in one line
        print(f"DB: {self.host} {self.port} {self.user} {self.password} {self.db_name} {self.ssl_ca}")

    def connect(self):
        if self.connection is not None:
            self.connection.close()
        self.connection = pymysql.connect(host=self.host,
                                          port=self.port,
                                          user=self.user,
                                          password=self.password,
                                          db=self.db_name,
                                          ssl_verify_cert=True,
                                          ssl_verify_identity=True,
                                          ssl_ca=self.ssl_ca,
                                          client_flag=pymysql.constants.CLIENT.MULTI_STATEMENTS)

    def query(self, sql):
        with self.connection.cursor() as cursor:
            cursor.execute(sql)
            result = cursor.fetchall()
        return result

    def query_file(self, filename):
        with open(filename, 'r') as file:
            sql = file.read()
        # set tidb_multiple_statements to ON to allow multiple statements in one query
        self.query("SET SESSION tidb_multi_statement_mode = 'ON'")
        return self.query(sql)