import json
import sqlite3
import os
from hashlib import md5

class SQLite():
    def __init__(self, db_name: str, schema:dict = None) -> None:
        self.conn: sqlite3.Connection|None = None
        self.cur = None
        self.db_name = db_name
        self.__backup = None
        self.schema = {} if schema is None else schema
        self.DB_VER = md5(json.dumps(self.schema).encode()).hexdigest()

    def connect(self) -> None:
        try:
            self.conn = sqlite3.connect(self.db_name, check_same_thread = False)
        except Exception:
            path_name = os.path.dirname(self.db_name)
            if not os.path.exists(path_name):
                os.makedirs(path_name)
            with open(self.db_name, 'x'):
                pass
            self.conn = sqlite3.connect(self.db_name, check_same_thread = False)
        def dict_factory(cursor, row):
            d = {}
            for idx, col in enumerate(cursor.description):
                d[col[0]] = row[idx]
            return d
        self.conn.row_factory = dict_factory

    def backup(self):
        with open(self.db_name, "rb") as f:
            self.__backup = f.read()
    
    def restore(self):
        were_active = True if self.conn else False
        self.disconnect()
        if self.__backup:
            with open(self.db_name, "wb") as f:
                f.write(self.__backup)
            self.__backup = None
        if were_active:
            self.connect()
            
    def delete_backup(self):
        self.__backup = None
    
    def disconnect(self) -> None:
        if self.conn:
            self.conn.close()
        self.conn = None

    def create_schema(self, tables = {}) -> None:
        if self.conn:
            cur = self.conn.cursor()
            cur.execute("CREATE TABLE IF NOT EXISTS main.keys_values(key VARCHAR(100) PRIMARY KEY, value VARCHAR(100) NOT NULL);")
            for t in tables:
                if t == "QUERY":
                    continue
                cur.execute('CREATE TABLE IF NOT EXISTS main.{}({});'.format(t, ''.join([(c + ' ' + tables[t][c] + ', ') for c in tables[t]])[:-2]))
            if "QUERY" in tables:
                [cur.execute(qry) for qry in tables["QUERY"]]
            cur.close()
    
    def query(self, query, *values):
        return self.queries([(query, *values)])[0]

    def queries(self, queries: list[tuple[str, ...]]):
        return list(self.queries_iter(queries))

    def queries_iter(self, queries: list[tuple[str, ...]]):
        cur = self.conn.cursor()
        try:
            cur.execute("BEGIN")
            for query_data in queries:
                values = []
                str_query = None
                if isinstance(query_data, str):
                    str_query = query_data
                elif (isinstance(query_data, tuple) or isinstance(query_data, list)) and len(query_data) > 0 and isinstance(query_data[0], str):
                    str_query = query_data[0]
                    values = query_data[1:]
                if str_query:
                    cur.execute(str_query, values)
                yield cur.fetchall()
            cur.execute("COMMIT")
        except Exception as e:
            cur.execute("ROLLBACK")
            raise e
        finally:
            cur.close()
            try:
                self.conn.commit()
            except Exception:
                pass
    
    def delete(self):
        self.disconnect()
        os.remove(self.db_name)
    
    def init(self):
        self.connect()
        try:
            if self.get('DB_VERSION') != self.DB_VER:
                raise Exception("DB_VERSION is not correct")
        except Exception:
            self.delete()
            self.connect()
            self.create_schema(self.schema)
            self.put('DB_VERSION', self.DB_VER)

    def get(self, key, default = None):
        q = self.query('SELECT value FROM keys_values WHERE key = ?', key)
        if len(q) == 0:
            return default
        else:
            return q[0]["value"]

    def put(self, key, value):
        if self.get(key) is None:
            self.query('INSERT INTO keys_values (key, value) VALUES (?, ?);', key, str(value))
        else:
            self.query('UPDATE keys_values SET value=? WHERE key = ?;', str(value), key)
    
    def set(self, key, value):
        return self.put(key, value)
        
    def dump(self) -> dict:
        self.connect()
        res = {}
        tables = self.query("SELECT name FROM sqlite_master WHERE type='table';")
        for table in tables:
            tname = table["name"]
            if tname == 'keys_values':
                # Three keys never travel in a backup. `password` and `secret` are this
                # instance's credentials, and `auth_disabled` is *whether it asks for
                # them* — a property of where this firegex is deployed, not of the
                # configuration being backed up. A backup that could carry it would be a
                # backup that can turn a firewall's authentication off on restore.
                res[tname] = self.query(
                    f"SELECT * FROM {tname} WHERE key NOT IN ('password', 'secret', 'auth_disabled');"
                )
            else:
                res[tname] = self.query(f"SELECT * FROM {tname};")
        return res
        
    def load(self, data: dict):
        self.connect()
        cur = self.conn.cursor()
        try:
            cur.execute("BEGIN")
            for table_name, rows in data.items():
                # Verify the table exists to prevent SQL injection or schema mismatches
                cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?;", (table_name,))
                if len(cur.fetchall()) == 0:
                    continue
                # Ignore anything that isn't a proper list of row objects
                if not isinstance(rows, list) or not all(isinstance(r, dict) for r in rows):
                    continue
                cur.execute(f"DELETE FROM {table_name};")
                if rows:
                    # Column names are identifiers and can't be parameterized, so only
                    # keep the ones that actually exist on the table - this stops a
                    # crafted backup from injecting arbitrary identifiers into the query.
                    cur.execute(f'PRAGMA table_info("{table_name}");')
                    valid_cols = {r["name"] for r in cur.fetchall()}
                    cols = [c for c in rows[0].keys() if c in valid_cols]
                    if not cols:
                        continue
                    placeholders = ", ".join(["?"] * len(cols))
                    col_names = ", ".join(cols)
                    query = f"INSERT INTO {table_name} ({col_names}) VALUES ({placeholders})"
                    for row in rows:
                        cur.execute(query, [row.get(c) for c in cols])
            cur.execute("COMMIT")
        except Exception as e:
            cur.execute("ROLLBACK")
            raise e
        finally:
            cur.close()
            try:
                self.conn.commit()
            except Exception:
                pass