from __future__ import annotations

from contextlib import contextmanager
from time import sleep
from typing import Generator

import psycopg2
import psycopg2.extensions
import psycopg2.pool

from ajusta_bling.common import Args


class Database:
    pool: psycopg2.pool.SimpleConnectionPool | None = None

    def __init__(self, args: Args):
        self.user = args.dbuser
        self.pswd = args.dbpassword
        self.host = args.dbhost
        self.port = args.dbport
        self.db = args.dbname
    
    def connect(self):
        self.pool = psycopg2.pool.SimpleConnectionPool(
            1, 10,
            user = self.user,
            password = self.pswd,
            host = self.host,
            port = self.port,
            database = self.db,
        )

    @contextmanager
    def get_cur(self) -> Generator[psycopg2.extensions.cursor]:
        if self.pool is None:
            self.connect()
        conn = self.pool.getconn()
        try:
            yield conn.cursor()
            conn.commit()
        finally:
            self.pool.putconn(conn)