# -*- coding: utf-8 -*-
#
# This file is part of Invenio.
# Copyright (C) 2015-2018 CERN.
# Copyright (C) 2022 RERO.
# Copyright (C) 2022 Graz University of Technology.
#
# Invenio is free software; you can redistribute it and/or modify it
# under the terms of the MIT License; see LICENSE file for more details.

"""Database management for Invenio."""

import os
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as package_version
from importlib.resources import files

import sqlalchemy as sa
from flask_alembic import Alembic
from invenio_base.utils import entry_points
from sqlalchemy_utils.functions import get_class_by_table

from .cli import db as db_cmd
from .shared import db
from .utils import versioning_models_registered


class InvenioDB(object):
    """Invenio database extension."""

    def __init__(self, app=None, **kwargs):
        """Extension initialization."""
        self.alembic = Alembic(run_mkdir=False, command_name="alembic")
        if app:
            self.init_app(app, **kwargs)

    def init_app(self, app, **kwargs):
        """Initialize application object."""
        self.init_db(app, **kwargs)

        def pathify(base_entry):
            return str(files(base_entry.module) / os.path.join(base_entry.attr))

        alembic_entry_points = entry_points(group="invenio_db.alembic")
        version_locations = [
            (base_entry.name, pathify(base_entry))
            for base_entry in alembic_entry_points
        ]
        script_location = str(files("invenio_db") / "alembic")
        app.config.setdefault(
            "ALEMBIC",
            {
                "script_location": script_location,
                "version_locations": version_locations,
            },
        )
        app.config.setdefault(
            "ALEMBIC_CONTEXT",
            {
                "transaction_per_migration": True,
                "compare_type": True,  # Allows to detect change of column type, accuracy depends on backend
            },
        )

        self.alembic.init_app(app)
        app.extensions["invenio-db"] = self
        app.cli.add_command(db_cmd)

    def init_db(self, app, entry_point_group="invenio_db.models", **kwargs):
        """Initialize Flask-SQLAlchemy extension."""
        # Setup SQLAlchemy
        app.config.setdefault(
            "SQLALCHEMY_DATABASE_URI",
            "sqlite:///" + os.path.join(app.instance_path, app.name + ".db"),
        )
        app.config.setdefault("SQLALCHEMY_ECHO", False)
        # Needed for before/after_flush/commit/rollback events
        app.config.setdefault("SQLALCHEMY_TRACK_MODIFICATIONS", True)

        # Initialize Flask-SQLAlchemy extension.
        database = kwargs.get("db", db)
        database.init_app(app)

        # Initialize versioning support.
        self.init_versioning(app, database, kwargs.get("versioning_manager"))

        # Initialize model bases
        if entry_point_group:
            for base_entry in entry_points(group=entry_point_group):
                base_entry.load()

        # All models should be loaded by now.
        sa.orm.configure_mappers()
        # Ensure that versioning classes have been built.
        if app.config["DB_VERSIONING"]:
            manager = self.versioning_manager
            if manager.pending_classes:
                if not versioning_models_registered(manager, database.Model):
                    manager.builder.configure_versioned_classes()
            elif "transaction" not in database.metadata.tables:
                manager.declarative_base = database.Model
                manager.create_transaction_model()
                manager.plugins.after_build_tx_class(manager)

    def init_versioning(self, app, database, versioning_manager=None):
        """Initialize the versioning support using SQLAlchemy-Continuum."""
        try:
            package_version("sqlalchemy_continuum")
        except PackageNotFoundError:  # pragma: no cover
            default_versioning = False
        else:
            default_versioning = True

        app.config.setdefault("DB_VERSIONING", default_versioning)

        if not app.config["DB_VERSIONING"]:
            return

        if not default_versioning:  # pragma: no cover
            raise RuntimeError(
                "Please install extra versioning support first by running "
                "pip install invenio-db[versioning]."
            )

        # Now we can import SQLAlchemy-Continuum.
        from sqlalchemy_continuum import make_versioned
        from sqlalchemy_continuum import versioning_manager as default_vm
        from sqlalchemy_continuum.plugins import FlaskPlugin

        # Try to guess user model class:
        if "DB_VERSIONING_USER_MODEL" not in app.config:  # pragma: no cover
            try:
                package_version("invenio_accounts")
            except PackageNotFoundError:
                user_cls = None
            else:
                user_cls = "User"
        else:
            user_cls = app.config.get("DB_VERSIONING_USER_MODEL")

        plugins = [FlaskPlugin()] if user_cls else []

        # Call make_versioned() before your models are defined.
        self.versioning_manager = versioning_manager or default_vm
        make_versioned(
            user_cls=user_cls,
            manager=self.versioning_manager,
            plugins=plugins,
        )

        # Register models that have been loaded beforehand.
        builder = self.versioning_manager.builder

        for tbl in database.metadata.tables.values():
            builder.instrument_versioned_classes(
                database.mapper, get_class_by_table(database.Model, tbl)
            )
