# copyright 2015-2016 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
# contact https://www.logilab.fr/ -- mailto:contact@logilab.fr
#
# This file is part of CubicWeb.
#
# CubicWeb is free software: you can redistribute it and/or modify it under the
# terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 2.1 of the License, or (at your option)
# any later version.
#
# CubicWeb 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 Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with CubicWeb.  If not, see <https://www.gnu.org/licenses/>.
"""
.. autoclass:: cubicweb.dataimport.massive_store.MassiveObjectStore
"""
import logging
from collections import defaultdict
from itertools import chain
from uuid import uuid4

from cubicweb.dataimport import stores, pgstore
from cubicweb.server.schema2sql import eschema_sql_def


class MassiveObjectStore(stores.RQLObjectStore):
    """Store for massive import of data, with delayed insertion of meta data.

    WARNINGS:

    - This store may only be used with PostgreSQL for now, as it relies
      on the COPY FROM method, and on specific PostgreSQL tables to get all
      the indexes.

    - This store can only insert relations that are not inlined (i.e.,
      which do *not* have inlined=True in their definition in the schema),
      unless they are specified as entity attributes.

    It should be used as follows:

       store = MassiveObjectStore(cnx)
       eid_p = store.prepare_insert_entity('Person',
                                           cwuri=u'http://dbpedia.org/toto',
                                           name=u'Toto')
       eid_loc = store.prepare_insert_entity('Location',
                                             cwuri=u'http://geonames.org/11111',
                                             name=u'Somewhere')
       store.prepare_insert_relation(eid_p, 'lives_in', eid_loc)
       store.flush()
       ...
       store.commit()
       store.finish()

    Full-text indexation is not handled, you'll have to reindex the proper entity types by yourself
    if desired.
    """

    def __init__(
        self, cnx, slave_mode=False, eids_seq_range=10000, metagen=None, drop=True
    ):
        """Create a MassiveObject store, with the following arguments:

        - `cnx`, a connection to the repository
        - `metagen`, optional :class:`MetadataGenerator` instance
        - `eids_seq_range`: size of eid range reserved by the store for each batch
        """
        super(MassiveObjectStore, self).__init__(cnx)

        self.uuid = str(uuid4()).replace("-", "")
        self.slave_mode = slave_mode
        if metagen is None:
            metagen = stores.MetadataGenerator(cnx)
        self.metagen = metagen

        self.logger = logging.getLogger("dataimport.massive_store")
        self.sql = cnx.system_sql
        self.schema = cnx.vreg.schema
        self.default_values = get_default_values(self.schema)
        self.get_next_eid = lambda g=self._get_eid_gen(eids_seq_range): next(g)
        self._source_dbhelper = cnx.repo.system_source.dbhelper
        if drop:
            self._dbh = PGHelper(cnx)
        else:
            self._dbh = NoDropPGHelper(cnx)

        self._data_entities = defaultdict(list)
        self._data_relations = defaultdict(list)
        self._data_entities_to_update = defaultdict(list)
        self._initialized = {}

    def _get_eid_gen(self, eids_seq_range):
        """Function getting the next eid. This is done by preselecting
        a given number of eids from the 'entities_id_seq', and then
        storing them"""
        while True:
            last_eid = self._cnx.repo.system_source.create_eid(
                self._cnx, eids_seq_range
            )
            for eid in range(last_eid - eids_seq_range + 1, last_eid + 1):
                yield eid

    # master/slaves specific API

    def master_init(self, commit=True):
        """Initialize database for massive insertion.

        This is expected to be called once, by the master store in master/slaves configuration.
        """
        assert not self.slave_mode
        if self not in self._initialized:
            self.sql("DROP TABLE IF EXISTS cwmassive_initialized")
            self.sql(
                "CREATE TABLE cwmassive_initialized"
                "(retype text, type varchar(128), uuid varchar(32))"
            )
            self._initialized[self] = None
            if commit:
                self.commit()

    # SQL utilities #########################################################

    def _drop_metadata_constraints(self):
        """Drop constraints and indexes for the metadata tables.

        They will be recreated by the `finish` method.
        """
        rtypes = [
            rtype
            for rtype in self.metagen.meta_relations
            if not self.schema.relation_schema_for(rtype).final
        ]
        rtypes += ("is_instance_of", "is", "cw_source")
        for rtype in rtypes:
            self._dbh.drop_constraints(rtype + "_relation")
            self._dbh.drop_indexes(rtype + "_relation")
        # don't drop constraints for the entities table, the only one is the primary key's index on
        # eid and we want to keep it
        self._dbh.drop_indexes("entities")

    def restart_eid_sequence(self, start_eid):
        self.sql(
            self._cnx.repo.system_source.dbhelper.sql_restart_numrange(
                "entities_id_seq", initial_value=start_eid
            )
        )
        self._cnx.commit()

    # store api ################################################################

    def prepare_insert_entity(self, etype, **data):
        """Given an entity type, attributes and inlined relations, returns the inserted entity's
        eid.
        """
        if etype not in self._initialized:
            if not self.slave_mode:
                self.master_init(commit=False)
            tablename = f"cw_{etype.lower()}"
            tmp_tablename = f"{tablename}_{self.uuid}"
            self.sql(
                "INSERT INTO cwmassive_initialized VALUES (%(e)s, 'etype', %(uuid)s)",
                {"e": etype, "uuid": self.uuid},
            )
            attr_defs = eschema_sql_def(self._source_dbhelper, self.schema[etype])
            self.sql(
                "CREATE TABLE %s(%s);"
                % (
                    tmp_tablename,
                    ", ".join(
                        "cw_%s %s" % (column, sqltype) for column, sqltype in attr_defs
                    ),
                )
            )
            self._initialized[etype] = [attr for attr, _ in attr_defs]

        if "eid" not in data:
            # If eid is not given and the eids sequence is set, use the value from the sequence
            eid = self.get_next_eid()
            data["eid"] = eid
        self._data_entities[etype].append(data)
        return data["eid"]

    def prepare_insert_relation(self, eid_from, rtype, eid_to, **kwargs):
        """Insert into the database a  relation ``rtype`` between entities with eids ``eid_from``
        and ``eid_to``.

        Relation must not be inlined.
        """
        if rtype not in self._initialized:
            if not self.slave_mode:
                self.master_init(commit=False)
            assert not self._cnx.vreg.schema.relation_schema_for(rtype).inlined
            self._initialized[rtype] = None
            tablename = f"{rtype.lower()}_relation"
            tmp_tablename = f"{tablename}_{self.uuid}"
            self.sql(
                "INSERT INTO cwmassive_initialized VALUES (%(r)s, 'rtype', %(uuid)s)",
                {"r": rtype, "uuid": self.uuid},
            )
            self.sql(f"CREATE TABLE {tmp_tablename}(eid_from integer, eid_to integer)")
        self._data_relations[rtype].append({"eid_from": eid_from, "eid_to": eid_to})

    def prepare_update_entity(self, etype, eid, **data):
        """Given an entity type, eid, attributes and inlined relations, returns
        the updated entity's eid.
        """
        entity = self._cnx.entity_from_eid(eid)
        assert entity.cw_etype == etype, f"Trying to update with wrong type {etype}"
        initialized_key = f"update_{etype}"
        if initialized_key not in self._initialized:
            if not self.slave_mode:
                self.master_init(commit=False)
            tablename = f"cw_{etype.lower()}_to_update"
            tmp_tablename = f"{tablename}_{self.uuid}"
            self.sql(
                "INSERT INTO cwmassive_initialized VALUES (%(e)s, 'update_etype', %(uuid)s)",
                {"e": etype, "uuid": self.uuid},
            )
            attr_defs = eschema_sql_def(self._source_dbhelper, self.schema[etype])
            self.sql(
                "CREATE TABLE %s(%s);"
                % (
                    tmp_tablename,
                    ", ".join(
                        "cw_%s %s" % (column, sqltype) for column, sqltype in attr_defs
                    ),
                )
            )
            self._initialized[initialized_key] = [attr for attr, _ in attr_defs]

        # we do that to avoid a lot of warning :)
        completed_data = {attr: None for attr in self._initialized[initialized_key]}
        self._data_entities_to_update[etype].append(
            {**completed_data, **data, "eid": eid}
        )
        return eid

    def flush(self):
        """Flush the data"""
        self.flush_entities()
        self.flush_entities_to_update()
        self.flush_relations()

    def finish(self):
        """Remove temporary tables and columns."""
        try:
            self._finish()
            self._cnx.commit()
        except Exception:
            self._cnx.rollback()
            raise
        finally:
            # delete the meta data table
            self.sql("DROP TABLE IF EXISTS cwmassive_initialized")
            self.commit()
            self._initialized = {}

    def _finish(self):
        """Remove temporary tables and columns."""
        assert (
            not self.slave_mode
        ), "finish method should only be called by the master store"
        self.logger.info("Start cleaning")
        # Get all the initialized etypes/rtypes
        if self._dbh.table_exists("cwmassive_initialized"):
            cu = self.sql("SELECT retype, type, uuid FROM cwmassive_initialized")
            entities = defaultdict(list)
            relations = defaultdict(list)
            entities_to_update = defaultdict(list)
            for retype, _type, uuid in cu.fetchall():
                if _type == "rtype":
                    relations[retype].append(uuid)
                elif _type == "update_etype":
                    entities_to_update[retype].append(uuid)
                else:  # _type = 'etype'
                    entities[retype].append(uuid)
            # if there is some entities to insert, delete constraint on metadata tables once for all
            if entities:
                self._drop_metadata_constraints()
            # get back entity to update data from the temporary tables
            for etype, uuids in entities_to_update.items():
                tablename = f"cw_{etype.lower()}"
                attr_defs = eschema_sql_def(self._source_dbhelper, self.schema[etype])
                set_parts = []
                for attr, _ in attr_defs:
                    if attr == "eid":
                        continue
                    set_parts.append(
                        f"cw_{attr} = COALESCE(TMP.cw_{attr}, {tablename}.cw_{attr})"
                    )

                self._dbh.drop_constraints(tablename)
                for uuid in uuids:
                    tmp_tablename = f"{tablename}_to_update_{uuid}"
                    self.sql(
                        f"""
                        UPDATE {tablename}
                        SET {",".join(set_parts)}
                        FROM {tmp_tablename} as TMP
                        WHERE {tablename}.cw_eid = TMP.cw_eid
                        """
                    )
                    self._tmp_data_cleanup(tmp_tablename, f"update_{etype}", uuid)
            # get back entity data from the temporary tables
            for etype, uuids in entities.items():
                tablename = f"cw_{etype.lower()}"
                attr_defs = eschema_sql_def(self._source_dbhelper, self.schema[etype])
                columns = ",".join(f"cw_{attr}" for attr, _ in attr_defs)
                self._dbh.drop_constraints(tablename)
                self._dbh.drop_indexes(tablename)
                for uuid in uuids:
                    tmp_tablename = f"{tablename}_{uuid}"
                    self.sql(
                        "INSERT INTO %(table)s(%(columns)s) "
                        "SELECT %(columns)s FROM %(tmp_table)s"
                        % {
                            "table": tablename,
                            "tmp_table": tmp_tablename,
                            "columns": columns,
                        }
                    )
                    self._insert_etype_metadata(etype, tmp_tablename)
                    self._tmp_data_cleanup(tmp_tablename, etype, uuid)
            # get back relation data from the temporary tables
            for rtype, uuids in relations.items():
                tablename = f"{rtype.lower()}_relation"
                self._dbh.drop_constraints(tablename)
                self._dbh.drop_indexes(tablename)
                for uuid in uuids:
                    tmp_tablename = f"{tablename}_{uuid}"
                    self.fill_relation_table(tablename, tmp_tablename)
                    self._tmp_data_cleanup(tmp_tablename, rtype, uuid)
        # restore all deleted indexes and constraints
        self._dbh.restore_indexes_and_constraints()

    def _insert_etype_metadata(self, etype, tmp_tablename):
        """Massive insertion of meta data for `etype`, with new entities in `tmp_tablename`."""
        # insert standard metadata relations
        for rtype, eid in self.metagen.base_etype_rels(etype).items():
            self.fill_meta_relation_table(tmp_tablename, rtype, eid)
        # insert cw_source, is and is_instance_of relations (normally handled by the system source)
        self.fill_meta_relation_table(
            tmp_tablename, "cw_source", self.metagen.source.eid
        )
        eschema = self.schema[etype]
        self.fill_meta_relation_table(tmp_tablename, "is", eschema.eid)
        for parent_eschema in chain(eschema.ancestors(), [eschema]):
            self.fill_meta_relation_table(
                tmp_tablename, "is_instance_of", parent_eschema.eid
            )
        self.fill_entities_table(etype, tmp_tablename)

    def fill_entities_table(self, etype, tmp_tablename):
        # finally insert records into the entities table
        self.sql(
            "INSERT INTO entities(eid, type) "
            "SELECT cw_eid, '%s' FROM %s "
            "WHERE NOT EXISTS (SELECT 1 FROM entities WHERE eid=cw_eid)"
            % (etype, tmp_tablename)
        )

    def fill_relation_table(self, tablename, tmp_tablename):
        # XXX no index on the original relation table, EXISTS subquery may be sloooow
        self.sql(
            "INSERT INTO %(table)s(eid_from, eid_to) SELECT DISTINCT "
            "T.eid_from, T.eid_to FROM %(tmp_table)s AS T "
            "WHERE NOT EXISTS (SELECT 1 FROM %(table)s AS TT WHERE "
            "TT.eid_from=T.eid_from AND TT.eid_to=T.eid_to);"
            % {"table": tablename, "tmp_table": tmp_tablename}
        )

    def fill_meta_relation_table(self, tmp_tablename, rtype, eid_to):
        self.sql(
            "INSERT INTO %s_relation(eid_from, eid_to) SELECT cw_eid, %s FROM %s "
            "WHERE NOT EXISTS (SELECT 1 FROM entities WHERE eid=cw_eid)"
            % (rtype, eid_to, tmp_tablename)
        )

    def _tmp_data_cleanup(self, tmp_tablename, ertype, uuid):
        """Drop temporary relation table and record from cwmassive_initialized."""
        self.sql(f"DROP TABLE {tmp_tablename}")
        self.sql(
            "DELETE FROM cwmassive_initialized "
            "WHERE retype = %(rtype)s AND uuid = %(uuid)s",
            {"rtype": ertype, "uuid": uuid},
        )

    # FLUSH #################################################################

    def flush_relations(self):
        """Flush the relations data from in-memory structures to a temporary table."""
        for rtype, data in self._data_relations.items():
            if not data:
                # There is no data for these etype for this flush round.
                continue
            buf = pgstore._create_copyfrom_buffer(data, ("eid_from", "eid_to"))
            cursor = self._cnx.cnxset.cu
            tablename = f"{rtype.lower()}_relation"
            tmp_tablename = f"{tablename}_{self.uuid}"
            cursor.copy_from(
                buf, tmp_tablename, null="NULL", columns=("eid_from", "eid_to")
            )
            # Clear data cache
            self._data_relations[rtype] = []

    def flush_entities(self):
        """Flush the entities data from in-memory structures to a temporary table."""
        metagen = self.metagen
        for etype, data in self._data_entities.items():
            if not data:
                # There is no data for these etype for this flush round.
                continue
            attrs = self._initialized[etype]
            _base_data = dict.fromkeys(attrs)
            _base_data.update(self.default_values[etype])
            _base_data.update(metagen.base_etype_attrs(etype))
            _data = []
            for d in data:
                # do this first on `d`, because it won't fill keys associated to None as provided by
                # `_base_data`
                metagen.init_entity_attrs(etype, d["eid"], d)
                # XXX warn/raise if there is some key not in attrs?
                _d = _base_data.copy()
                _d.update(d)
                _data.append(_d)
            buf = pgstore._create_copyfrom_buffer(_data, attrs)
            tablename = f"cw_{etype.lower()}"
            tmp_tablename = f"{tablename}_{self.uuid}"
            columns = [f"cw_{attr}" for attr in attrs]
            cursor = self._cnx.cnxset.cu
            cursor.copy_from(buf, tmp_tablename, null="NULL", columns=columns)
            # Clear data cache
            self._data_entities[etype] = []

    def flush_entities_to_update(self):
        """Flush the entities to update data from in-memory structures to a temporary table."""
        for etype, data in self._data_entities_to_update.items():
            if not data:
                # There is no data for these etype for this flush round.
                continue
            attrs = self._initialized[f"update_{etype}"]
            buf = pgstore._create_copyfrom_buffer(data, attrs)
            tablename = f"cw_{etype.lower()}_to_update"
            tmp_tablename = f"{tablename}_{self.uuid}"
            columns = [f"cw_{attr}" for attr in attrs]
            cursor = self._cnx.cnxset.cu
            cursor.copy_from(buf, tmp_tablename, null="NULL", columns=columns)
            # Clear data cache
            self._data_entities_to_update[etype] = []


def get_default_values(schema):
    """analyzes yams ``schema`` and returns the list of default values.

    The returned value is a dictionary mapping entity types to a
    sub-dictionnaries mapping attribute names -> default values.
    """
    default_values = {}
    # iterates on all entity types
    for eschema in schema.entities():
        # for each entity type, iterates on attribute definitions
        default_values[eschema.type] = eschema_constraints = {}
        for rschema, _ in eschema.attribute_definitions():
            # for each attribute, if a size constraint is found,
            # append it to the size constraint list
            if eschema.default(rschema.type) is not None:
                eschema_constraints[rschema.type] = eschema.default(rschema.type)
    return default_values


class PGHelper:
    """This class provides some helper methods to manipulate a postgres database metadata (index and
    constraints).
    """

    def __init__(self, cnx):
        self.sql = cnx.system_sql
        # Deals with pg schema, see #3216686
        pg_schema = cnx.repo.config.system_source_config.get("db-namespace") or "public"
        self.pg_schema = pg_schema

    def drop_indexes(self, tablename):
        """Drop indexes and constraints, storing them in a table for later restore."""
        # Create a table to save the constraints, it allows reloading even after crash
        self.sql(
            "CREATE TABLE IF NOT EXISTS cwmassive_constraints(sql TEXT, insert_order SERIAL)"
        )
        indexes = self.table_indexes(tablename)
        for name, query in indexes.items():
            self.sql(
                "INSERT INTO cwmassive_constraints(sql) VALUES (%(sql)s)",
                {"sql": query},
            )
            self.sql(f"DROP INDEX {name}")

    def drop_constraints(self, tablename):
        self.sql(
            "CREATE TABLE IF NOT EXISTS cwmassive_constraints(sql TEXT, insert_order SERIAL)"
        )
        constraints = self.table_constraints(tablename)
        for name, query in constraints.items():
            self.sql(
                "INSERT INTO cwmassive_constraints(sql) VALUES (%(sql)s)",
                {"sql": query},
            )
            self.sql(f"ALTER TABLE {tablename} DROP CONSTRAINT {name}")

    def restore_indexes_and_constraints(self):
        """Restore indexes and constraints."""
        if not self.table_exists("cwmassive_constraints"):
            return
        cu = self.sql(
            "SELECT sql, insert_order FROM cwmassive_constraints "
            "ORDER BY insert_order DESC"
        )
        for query, order in cu.fetchall():
            self.sql(query)
            self.sql(
                "DELETE FROM cwmassive_constraints WHERE insert_order=%(order)s",
                {"order": order},
            )
        self.sql("DROP TABLE cwmassive_constraints")

    def table_exists(self, tablename):
        """Return True if the given table already exists in the database."""
        cu = self.sql(
            "SELECT 1 from information_schema.tables "
            "WHERE table_name=%(t)s AND table_schema=%(s)s",
            {"t": tablename, "s": self.pg_schema},
        )
        return bool(cu.fetchone())

    def table_indexes(self, tablename):
        """Return a dictionary of indexes {index name: index sql}, constraints included."""
        indexes = {}
        for name in self._index_names(tablename):
            indexes[name] = self._index_sql(name)
        return indexes

    def table_constraints(self, tablename):
        """Return a dictionary of constraints {constraint name: constraint sql}."""
        constraints = {}
        for name in self._constraint_names(tablename):
            query = self._constraint_sql(name)
            constraints[name] = f"ALTER TABLE {tablename} ADD CONSTRAINT {name} {query}"
        return constraints

    def _index_names(self, tablename):
        """Return the names of all indexes in the given table (including constraints.)"""
        cu = self.sql(
            "SELECT c.relname FROM pg_catalog.pg_class c "
            "JOIN pg_catalog.pg_index i ON i.indexrelid = c.oid "
            "JOIN pg_catalog.pg_class c2 ON i.indrelid = c2.oid "
            "LEFT JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner "
            "LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "
            "WHERE c.relkind IN ('i','') "
            " AND c2.relname = %(t)s "
            " AND i.indisprimary = FALSE "
            " AND n.nspname NOT IN ('pg_catalog', 'pg_toast') "
            " AND pg_catalog.pg_table_is_visible(c.oid);",
            {"t": tablename},
        )
        return [name for name, in cu.fetchall()]

    def _constraint_names(self, tablename):
        """Return the names of all constraints in the given table."""
        cu = self.sql(
            "SELECT i.conname FROM pg_catalog.pg_class c "
            "JOIN pg_catalog.pg_constraint i ON i.conrelid = c.oid "
            "JOIN pg_catalog.pg_class c2 ON i.conrelid=c2.oid "
            "LEFT JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner "
            "LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "
            "WHERE c2.relname = %(t)s "
            "AND n.nspname NOT IN ('pg_catalog', 'pg_toast') "
            "AND pg_catalog.pg_table_is_visible(c.oid)",
            {"t": tablename},
        )
        return [name for name, in cu.fetchall()]

    def _index_sql(self, name):
        """Return the SQL to be used to recreate the index of the given name."""
        return self.sql(
            "SELECT pg_get_indexdef(c.oid) FROM pg_catalog.pg_class c "
            "LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "
            "WHERE c.relname = %(r)s AND n.nspname=%(n)s",
            {"r": name, "n": self.pg_schema},
        ).fetchone()[0]

    def _constraint_sql(self, name):
        """Return the SQL to be used to recreate the constraint."""
        return self.sql(
            "SELECT pg_get_constraintdef(c.oid) FROM pg_catalog.pg_constraint c "
            "LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.connamespace "
            "WHERE c.conname = %(r)s AND n.nspname=%(n)s",
            {"r": name, "n": self.pg_schema},
        ).fetchone()[0]


class NoDropPGHelper(PGHelper):
    """Custom PGHelper that doesn't drop any index nor constraint.
    Instead of dropping the constraints, they are disabled, and restored at
    the end of the import.

    It's usually faster when a table already contains data that you want to
    keep.
    """

    def drop_indexes(self, tablename):
        pass

    def drop_constraints(self, tablename):
        """Disable constraints on the given table."""
        self.sql(
            "CREATE TABLE IF NOT EXISTS cwmassive_constraints(sql TEXT, insert_order SERIAL)"
        )

        self.sql(f"ALTER TABLE {tablename} DISABLE TRIGGER ALL;")

        self.sql(
            "INSERT INTO cwmassive_constraints(sql) VALUES (%(sql)s)",
            {"sql": f"ALTER TABLE {tablename} ENABLE TRIGGER ALL;"},
        )
