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

# project
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


def create_vector_layer(
    connection: Connection,
    layer_name: str,
    table_name: str,
    schema: str,
    geom_col: str = "",
    key_col: str = "",
) -> QgsVectorLayer:
    """Create a QgsVectorLayer from connection / table / schema

    Args:
        connection (Connection): connection
        layer_name (str): layer name
        table_name (str): table name
        schema (str): schema
        geom_col (str): column for geometry (default "")
        key_col (str): key column for geometry (default "")

    Returns:
        QgsVectorLayer: create QgsVectorLayer
    """
    uri = create_datasource_uri(connection)
    uri.setDataSource(
        aSchema=schema,
        aTable=table_name,
        aGeometryColumn=geom_col,
        aKeyColumn=key_col,
    )
    if not key_col:
        uri.setParam("checkPrimaryKeyUnicity", "0")
    uri_str = uri.uri(False)

    layer = QgsVectorLayer(uri_str, layer_name, "postgres")
    return layer


def check_if_table_exist(connection: Connection, table_name: str, schema: str) -> bool:
    """Check if a table exists in a postgres connection

    Args:
        connection (Connection): connection
        table_name (str): table name_
        schema (str): schema

    Returns:
        bool: True if the table exists, False otherwise
    """
    layer = create_vector_layer(
        connection=connection,
        layer_name="temp_layer",
        table_name=table_name,
        schema=schema,
    )
    return layer.isValid()
