import psycopg2
from qgis.core import QgsDataSourceUri
from sqlalchemy import create_engine, event
from sqlalchemy.engine import URL, Engine
from sqlalchemy.orm import Session, sessionmaker

from openlog.datamodel.connection.openlog_connection import Connection


def create_session(connection: Connection) -> (Session, Engine):
    """
    Create a sqlalchemy session for postgres database

    Args:
        connection: connection model with database connection parameters

    Returns: sqlachemy session

    """
    engine = create_engine("postgresql+psycopg2://user:pass@hostname/dbname")
    args = connection.get_postgres_args()
    # Add sqlalchemy event to use postgres args so service file can be use with additionnal arguments
    @event.listens_for(engine, "do_connect")
    def receive_do_connect(dialect, conn_rec, cargs, cparams):
        # return the new DBAPI connection with postgres args
        return psycopg2.connect(**args)

    postgres_session = sessionmaker(bind=engine)
    session = postgres_session()

    return session, engine


def create_datasource_uri(connection: Connection) -> QgsDataSourceUri:
    """
    Returns a QgsDataSourceUri from connection parameters

    Args:
        connection: connection model with database connection parameters

    Returns: QgsDataSourceUri for connection parameters

    """
    uri = QgsDataSourceUri()
    if not connection.service:
        uri.setConnection(
            aHost=connection.host,
            aPort=connection.port,
            aDatabase=connection.database,
            aUsername=connection.user,
            aPassword=connection.password,
        )
    else:
        uri.setConnection(
            aService=connection.service,
            aDatabase=connection.database,
            aPassword=connection.password,
            aUsername=connection.user,
        )
    return uri
