""" Implements database connection, cursor, etc. functionality.
"""
from __future__ import annotations

import psycopg as pg
import psycopg.sql as sql
from psycopg.types.enum import EnumInfo, register_enum

import chemrecon.schema as schema
from chemrecon.database.params import Params

# Initialisation functions
# ----------------------------------------------------------------------------------------------------------------------
def postgres_connect(
    params: Params,
    initialise_enums: bool = True
) -> pg.Connection:
    """ Connect to the database server.
        Connection string given relative to the working directory (chemrecon/).
        If find_directory is set, will automatically try to locate the chemrecon directory
    """
    # Load parameters
    print(f'[ChemRecon] Attempting to connect with string: {params.connection_string()}')
    try:
        conn: pg.Connection = pg.connect(
            params.connection_string()
        )
    except pg.OperationalError as e:
        print(f'[ChemRecon] Could not connect: {e}')
        raise ConnectionError from e

    print(f'[ChemRecon] Connection succesful.')

    # Set correct schemata to search
    schemata = [
        'chemrecon',
        'pg_catalog',
        'meta',
        'public'
    ]
    conn.execute(
        sql.SQL("SET {search_path} TO {schemata}").format(
            search_path = sql.Identifier('search_path'),
            schemata = sql.SQL(', ').join(
                [sql.Identifier(s) for s in schemata]
            )
        )
    )
    conn.commit()

    # Register enum mapping
    if initialise_enums:
        conn.commit()
        for enum in schema.enums.enum_register:
            #  enum_info = EnumInfo.fetch(conn, f'chemrecon.{enum.__name__}')
            enum_info = EnumInfo.fetch(conn, sql.Identifier(enum.__name__))
            if enum_info is None:
                raise RuntimeError(f'Enum info none for enum {enum}.')
            register_enum(enum_info, conn, enum)

    # Done
    conn.commit()
    return conn