import psycopg2
from psycopg2 import pool
from .base_db import BaseDBHandler

class PostgreSQLHandler(BaseDBHandler):
    def connect(self, service, user, password, database=None):
        try:
            self.postgresql_pool = psycopg2.pool.SimpleConnectionPool(
                minconn=1,
                maxconn=5,
                host=service,
                port=5623,
                user=user,
                password=password,
                database=database
            )
            connection = self.postgresql_pool.getconn()
            self.postgresql_pool.putconn(connection)

            return self.make_connection_response("Successfully connected to PostgreSQL")
        except Exception as e:
            return self.make_error_response(500, f"Internal Server Error, Failed to connect to PostgreSQL: {e}")

    def execute_query(self, query, database=None):
        try:
            connection = self.postgresql_pool.getconn()
            cursor = connection.cursor()

            cursor.execute(query)

            if query.strip().lower().startswith(("insert", "update", "delete", "create", "drop", "alter")):
                connection.commit()
                results = f"{cursor.rowcount} rows affected"
                columns_names = []
            else:
                results = cursor.fetchall()
                columns_names = [desc[0] for desc in cursor.description] if cursor.description else []

            cursor.close()
            self.postgresql_pool.putconn(connection)
            return self.make_success_response(
                columns_names=columns_names,
                results=results
            )
        except Exception as e:
            return self.make_error_response(500, f"Internal Server Error, Please check the query and try again: {e}")