from peewee import MySQLDatabase, PostgresqlDatabase
from settings.db import (
    DB_ENGINE,
    DB_HOST,
    DB_PORT,
    DB_USER,
    DB_PASSWORD
)


def guess_database(db_name):
    """
    function checks DB_ENGINE and set correct connection class
    with connection params. Expected database name in input.

    Args:
        db_name (_str_): database name

    Returns:
        (_MySQLDatabase_or_PostgresqlDatabase_ class):
        contains db connection params
    """
    if DB_ENGINE == 'mysql':
        return MySQLDatabase(
            db_name, **{
                'charset': 'utf8',
                'sql_mode': 'PIPES_AS_CONCAT',
                'use_unicode': True,
                'host': DB_HOST,
                'port': DB_PORT,
                'user': DB_USER,
                'password': DB_PASSWORD
            }
        )
    if DB_ENGINE == 'psql':
        return PostgresqlDatabase(
            db_name, **{
                'charset': 'utf8',
                'sql_mode': 'PIPES_AS_CONCAT',
                'use_unicode': True,
                'host': DB_HOST,
                'port': DB_PORT,
                'user': DB_USER,
                'password': DB_PASSWORD
            }
        )