"""
The latest version of this package is available at:
<http://github.com/jantman/biweeklybudget>

################################################################################
Copyright 2016-2024 Jason Antman <http://www.jasonantman.com>

    This file is part of biweeklybudget, also known as biweeklybudget.

    biweeklybudget is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    biweeklybudget is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with biweeklybudget.  If not, see <http://www.gnu.org/licenses/>.

The Copyright and Authors attributions contained herein may not be removed or
otherwise altered, except to add the Author attribution of a contributor to
this work. (Additional Terms pursuant to Section 7b of the AGPL v3)
################################################################################
While not legally required, I sincerely request that anyone who finds
bugs please submit them at <https://github.com/jantman/biweeklybudget> or
to me via email, and that you send any contributions or improvements
either as a pull request on GitHub, or to me via email.
################################################################################

AUTHORS:
Jason Antman <jason@jasonantman.com> <http://www.jasonantman.com>
################################################################################

This is mostly based on http://flask.pocoo.org/docs/0.12/patterns/sqlalchemy/
"""

from biweeklybudget import settings
import logging
import os
from copy import deepcopy
import warnings
import pkg_resources

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
from pymysql.err import Warning
from alembic.config import Config
from alembic.script import ScriptDirectory
from alembic.runtime.environment import EnvironmentContext
from alembic import command

from biweeklybudget.utils import in_directory

logger = logging.getLogger(__name__)

logger.debug('Creating DB engine with connection: %s',
             settings.DB_CONNSTRING)

echo = False
if os.environ.get('SQL_ECHO', '') == 'true':
    echo = True

# for Py37 warning 3135:
# "'NO_ZERO_DATE', 'NO_ZERO_IN_DATE' and
# 'ERROR_FOR_DIVISION_BY_ZERO' sql modes should be used with strict
#  mode. They will be merged with strict mode in a future release."
warnings.filterwarnings('ignore', '.*merged with strict mode.*')

# For some reason, with PyMySQL, even setting sql_mode to TRADITIONAL isn't
# raising an Exception when data is truncated. So we need to explicitly convert
# ``pymysql.err.Warning`` to an exception...
warnings.simplefilter('error', category=Warning)

engine_args = {
    'echo': echo,
    'pool_recycle': 3600
}

if settings.DB_CONNSTRING.startswith('mysql'):
    engine_args['connect_args'] = {'sql_mode': 'TRADITIONAL'}

if 'SQL_POOL_PRE_PING' in os.environ:
    engine_args['pool_pre_ping'] = True

#: The database engine object; return value of
#: :py:func:`sqlalchemy.create_engine`.
engine = create_engine(settings.DB_CONNSTRING, **engine_args)

logger.debug('Creating DB session')

#: :py:class:`sqlalchemy.orm.scoping.scoped_session` session
db_session = scoped_session(
    sessionmaker(autocommit=False, autoflush=False, bind=engine)
)

logger.debug('Setting up Base and query')

from biweeklybudget.models.base import Base  # noqa
Base.query = db_session.query_property()
from biweeklybudget.db_event_handlers import init_event_listeners  # noqa


def _alembic_get_current_rev(config, script):
    """
    Works sorta like alembic.command.current

    :param config: alembic Config
    :return: current revision
    :rtype: str
    """
    config._curr_rev = None

    def display_version(rev, _):
        for rev in script.get_all_current(rev):
            config._curr_rev = rev.cmd_format(False)
        return []

    with EnvironmentContext(
        config,
        script,
        fn=display_version
    ):
        script.run_env()
    if config._curr_rev is not None and ' ' in config._curr_rev:
        config._curr_rev = config._curr_rev.strip().split(' ')[0]
    return config._curr_rev


def init_db():
    """
    Initialize the database; call
    :py:meth:`sqlalchemy.schema.MetaData.create_all` on the metadata object.
    """
    logger.debug(
        'Initializing database with connstring: %s', settings.DB_CONNSTRING
    )
    # import all modules here that might define models so that
    # they will be registered properly on the metadata.  Otherwise
    # you will have to import them first before calling init_db()
    alembic_ini = pkg_resources.resource_filename(
        pkg_resources.Requirement.parse('biweeklybudget'),
        'biweeklybudget/alembic/alembic.ini'
    )
    topdir = os.path.abspath(
        os.path.join(os.path.dirname(alembic_ini), '..', '..')
    )
    logger.debug('Alembic configuration: %s', alembic_ini)
    with in_directory(topdir):
        alembic_config = Config(alembic_ini)
        script = ScriptDirectory.from_config(alembic_config)
        curr_rev = _alembic_get_current_rev(alembic_config, script)
        head_rev = script.get_revision("head").revision
        if curr_rev is None:
            # alembic not initialized at all; stamp with current version
            logger.warning(
                'Alembic not setup; creating all models and stamping'
            )
            logger.debug('Creating all models')
            with warnings.catch_warnings():
                # Ignore warning from MariaDB MDEV-17544 when creating tables;
                # SQLAlchemy 1.3.13 names the primary keys, but MariaDB ignores
                # the names; MariaDB >= 10.4.7 now throws a warning for this.
                warnings.filterwarnings(
                    'ignore',
                    message=r'^\(1280, "Name.*ignored for PRIMARY key\."\)$',
                    category=Warning
                )
                Base.metadata.create_all(engine)
            command.stamp(alembic_config, "head")
            logger.debug("DB stamped at %s", head_rev)
        elif curr_rev != head_rev:
            logger.warning("Alembic head is %s but this DB is at %s; "
                           "running migrations", head_rev, curr_rev)
            command.upgrade(alembic_config, "head")
            logger.info("Migrations complete")
        else:
            logger.debug('Alembic is at the correct head version (%s)',
                         curr_rev)
    logger.debug('Done initializing DB')
    init_event_listeners(db_session, engine)


def cleanup_db():
    """
    This must be called from all scripts, using

        atexit.register(cleanup_db)

    """
    logger.debug('Closing DB session')
    db_session.remove()


def upsert_record(model_class, key_fields, **kwargs):
    """
    Upsert a record in the database.

    ``key_fields`` is either a string primary key field name (a key in the
    ``kwargs`` dict) or a list or tuple of string primary key field names, for
    compound keys.

    If a record can be found matching these keys, it will be updated and
    committed. If not, a new one will be inserted. Either way, the record is
    returned.

    :py:meth:`sqlalchemy.orm.session.Session.commit` is **NOT** called.

    :param model_class: the class of model to insert/update
    :type model_class: biweeklybudget.models.base.ModelAsDict
    :param key_fields: The field name(s) (keys in ``kwargs``) that make up the
      primary key. This can be a single string, or a list or tuple of strings
      for compound keys. The values for these key fields MUST be included in
      ``kwargs``.
    :param kwargs: arguments to provide to the model class constructor, or to
      update if there is an existing record matching the key.
    :type kwargs: dict
    :return: inserted or updated record; type is an instance of ``model_class``
    """
    args = {}
    for k, v in kwargs.items():
        if isinstance(v, Base):
            args[k] = v
        else:
            args[k] = deepcopy(v)
    pkey = None
    if isinstance(key_fields, type('')):
        pkey = args[key_fields]
        del args[key_fields]
    else:
        pkey = []
        for k in key_fields:
            pkey.append(args[k])
            del args[k]
        pkey = tuple(pkey)
    logger.info('Upserting %s key=%s: %s', model_class, pkey, kwargs)
    res = db_session.query(model_class).get(pkey)
    if res is None:
        # not in DB yet; insert
        logger.debug('INSERTing (no existing record found)')
        o = model_class(**kwargs)
        db_session.add(o)
        return o
    logger.debug('Matching record exists; updating')
    for k, v in args.items():
        setattr(res, k, v)
    return res
