#
#   Copyright © 2021,2022 IsardVDI S.L.
#   Copyright © 2022 Evilham <contact@evilham.com>
#
#   This file is part of DD
#
#   DD is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or (at your
#   option) any later version.
#
#   DD 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 DD. If not, see <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: AGPL-3.0-or-later

import psycopg2
import psycopg2.sql
from psycopg2.extensions import connection, cursor

from typing import Any, List, Tuple, Union

query = Union[str, psycopg2.sql.SQL]

class Postgres:
    # TODO: Fix this whole class
    cur : cursor
    conn : connection
    def __init__(self, host : str, database : str, user : str, password : str) -> None:
        self.conn = psycopg2.connect(
            host=host, database=database, user=user, password=password
        )

    def select(self, sql: query) -> List[Tuple[Any, ...]]:
        self.cur = self.conn.cursor()
        self.cur.execute(sql)
        data = self.cur.fetchall()
        self.cur.close()  # type: ignore  # psycopg2 type hint missing
        return data

    def update(self, sql : query) -> None:
        self.cur = self.conn.cursor()
        self.cur.execute(sql)
        self.conn.commit()
        self.cur.close()  # type: ignore  # psycopg2 type hint missing
        # return self.cur.fetchall()

    def select_with_headers(self, sql : query) -> Tuple[List[Any], List[Tuple[Any, ...]]]:
        self.cur = self.conn.cursor()
        self.cur.execute(sql)
        data = self.cur.fetchall()
        fields = [a.name for a in self.cur.description]
        self.cur.close()  # type: ignore  # psycopg2 type hint missing
        return (fields, data)
