# coding: utf-8

# copyright 2015 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

# contact http://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 <http://www.gnu.org/licenses/>.

import logging

from datetime import datetime

from collections import defaultdict

from io import StringIO

from six.moves import range

from yams.constraints import SizeConstraint

from psycopg2 import ProgrammingError

from cubicweb.dataimport import stores, pgstore

from cubicweb.utils import make_uid

from cubicweb.server.sqlutils import SQL_PREFIX

class MassiveObjectStore(stores.RQLObjectStore):

    """

    Store for massive import of data, with delayed insertion of meta data.

    WARNINGS:

   - This store may be only 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).

   It should be used as follows:

       store = MassiveObjectStore(cnx)

       store.init_rtype_table('Person', 'lives_in', 'Location')

       ...

       store.create_entity('Person', subj_iid_attribute=person_iid, ...)

       store.create_entity('Location', obj_iid_attribute=location_iid, ...)

       ...

       # subj_iid_attribute and obj_iid_attribute are argument names

       # chosen by the user (e.g. "cwuri"). These names can be identical.

       # person_iid and location_iid are unique IDs and depend on the data

       # (e.g URI).

       store.flush()

       store.relate_by_iid(person_iid, 'lives_in', location_iid)

       # For example:

       store.create_entity('Person',

                           cwuri='http://dbpedia.org/toto',

                           name='Toto')

       store.create_entity('Location',

                           uri='http://geonames.org/11111',

                           name='Somewhere')

       store.flush()

       store.relate_by_iid('http://dbpedia.org/toto',

                       'lives_in',

                       'http://geonames.org/11111')

       # Finally

       store.flush_meta_data()

       store.convert_relations('Person', 'lives_in', 'Location',

                               'subj_iid_attribute', 'obj_iid_attribute')

       # For the previous example:

       store.convert_relations('Person', 'lives_in', 'Location', 'cwuri', 'uri')

       ...

       store.cleanup()

    """

    # size of eid range reserved by the store for each batch

    eids_seq_range = 10000

    # initial eid (None means use the value in the db)

    eids_seq_start = None

    def __init__(self, cnx, autoflush_metadata=True,

                 commit_at_flush=True,

                 on_commit_callback=None, on_rollback_callback=None,

                 slave_mode=False,

                 source=None):

        """ Create a MassiveObject store, with the following attributes:

        - cnx: CubicWeb cnx

        - autoflush_metadata: Boolean.

                              Automatically flush the metadata after

                              each flush()

        - commit_at_flush: Boolean. Commit after each flush().

        - uri_param_name: String. If given, will use this parameter to get cw_uri

                          for entities.

        """

        super(MassiveObjectStore, self).__init__(cnx)

        self.logger = logging.getLogger('dataio.relationmixin')

        self._cnx = cnx

        self.sql = cnx.system_sql

        self.commit_at_flush = commit_at_flush

        self._data_uri_relations = defaultdict(list)

        self._initialized = {'init_uri_eid': set(),

                             'uri_eid_inserted': set(),

                             'uri_rtypes': set(),

                             'entities': set(),

                             'rtypes': set(),

                            }

        self.sql = self._cnx.system_sql

        self.logger = logging.getLogger('dataio.massiveimport')

        self.autoflush_metadata = autoflush_metadata

        self.slave_mode = slave_mode

        self.size_constraints = get_size_constraints(cnx.vreg.schema)

        self.default_values = get_default_values(cnx.vreg.schema)

        pg_schema = cnx.repo.config.system_source_config.get('db-namespace', 'public')

        self._dbh = PGHelper(self._cnx, pg_schema)

        self._data_entities = defaultdict(list)

        self._data_relations = defaultdict(list)

        self._now = datetime.now()

        self._default_cwuri = make_uid('_auto_generated')

        self.uri_param_name = uri_param_name

        self._count_cwuri = 0

        self.commit_at_flush = commit_at_flush

        self.on_commit_callback = on_commit_callback

        self.on_rollback_callback = on_rollback_callback

        # Initialized the meta tables of dataio for warm restart

        self._init_dataio_metatables()

        # Internal markers of initialization

        if self.eids_seq_start is not None and not self.slave_mode:

            self._cnx.system_sql(self._cnx.repo.system_source.dbhelper.sql_restart_numrange(

                'entities_id_seq', initial_value=self.eids_seq_start + 1))

            cnx.commit()

        self.get_next_eid = lambda g=self._get_eid_gen(): next(g)

        # recreate then when self.finish() is called

        if not self.slave_mode:

            self._drop_metatables_constraints()

        if source is None:

            source = cnx.repo.system_source

        self.source = source

        self._etype_eid_idx = dict(cnx.execute('Any XN,X WHERE X is CWEType, X name XN'))

        cnx.read_security = False

        cnx.write_security = False

    ### INIT FUNCTIONS ########################################################

    def init_rtype_table(self, etype_from, rtype, etype_to):

        """ Build temporary table a for standard rtype """

        # Create an uri_eid table for each etype for a better

        # control of which etype is concerns for a particular

        # possibly multivalued relation.

        for etype in (etype_from, etype_to):

            if etype and etype not in self._initialized['init_uri_eid']:

                self._init_uri_eid_table(etype)

        if rtype not in self._initialized['uri_rtypes']:

            # Create the temporary tables

            if not self._cnx.repo.schema.rschema(rtype).inlined:

                try:

                    sql = 'CREATE TABLE %(r)s_relation_iid_tmp (uri_from character ' \

                          'varying(%(s)s), uri_to character varying(%(s)s))'

                    self.sql(sql % {'r': rtype, 's': self.iid_maxsize})

                except ProgrammingError:

                    # XXX Already exist (probably due to multiple import)

                    pass

            else:

                self.logger.warning("inlined relation %s: cannot insert it", rtype)

            #Add it to the initialized set

            self._initialized['uri_rtypes'].add(rtype)

    def _init_uri_eid_table(self, etype):

        """ Build a temporary table for id/eid convertion

        """

        try:

            sql = "CREATE TABLE uri_eid_%(e)s (uri character varying(%(size)s), eid integer)"

            self.sql(sql % {'e': etype.lower(), 'size': self.iid_maxsize,})

        except ProgrammingError:

            # XXX Already exist (probably due to multiple import)

            pass

        # Add it to the initialized set

        self._initialized['init_uri_eid'].add(etype)

    def _init_dataio_metatables(self):

        """ Initialized the meta tables of dataio for warm restart

        """

        # Check if dataio tables are not already created (i.e. a restart)

        self._initialized_table_created = self._dbh.table_exists('dataio_initialized')

        self._constraint_table_created = self._dbh.table_exists('dataio_constraints')

        self._metadata_table_created = self._dbh.table_exists('dataio_metadata')

    ### RELATE FUNCTION #######################################################

    def relate_by_iid(self, iid_from, rtype, iid_to):

        """Add new relation based on the internal id (iid)

        of the entities (not the eid)"""

        # Push data

        if isinstance(iid_from, unicode):

            iid_from = iid_from.encode('utf-8')

        if isinstance(iid_to, unicode):

            iid_to = iid_to.encode('utf-8')

        self._data_uri_relations[rtype].append({'uri_from': iid_from, 'uri_to': iid_to})

    ### FLUSH FUNCTIONS #######################################################

    def flush_relations(self):

        """ Flush the relations data

        """

        for rtype, data in self._data_uri_relations.items():

            if not data:

                self.logger.info('No data for rtype %s', rtype)

            buf = StringIO('\n'.join(['%(uri_from)s\t%(uri_to)s' % d for d in data]))

            if not buf:

                self.logger.info('Empty Buffer for rtype %s', rtype)

                continue

            cursor = self._cnx.cnxset.cu

            if not self._cnx.repo.schema.rschema(rtype).inlined:

                cursor.copy_from(buf, '%s_relation_iid_tmp' % rtype.lower(),

                                 null='NULL', columns=('uri_from', 'uri_to'))

            else:

                self.logger.warning("inlined relation %s: cannot insert it", rtype)

            buf.close()

            # Clear data cache

            self._data_uri_relations[rtype] = []

            # Commit if asked

            if self.commit_at_flush:

                self.commit()

    def fill_uri_eid_table(self, etype, uri_label):

        """ Fill the uri_eid table

        """

        self.logger.info('Fill uri_eid for etype %s', etype)

        sql = 'INSERT INTO uri_eid_%(e)s SELECT cw_%(l)s, cw_eid FROM cw_%(e)s'

        self.sql(sql % {'l': uri_label, 'e': etype.lower()})

        # Add indexes

        self.sql('CREATE INDEX uri_eid_%(e)s_idx ON uri_eid_%(e)s' '(uri)' % {'e': etype.lower()})

        # Set the etype as converted

        self._initialized['uri_eid_inserted'].add(etype)

        self.commit()

    def convert_relations(self, etype_from, rtype, etype_to,

                          uri_label_from='cwuri', uri_label_to='cwuri'):

        """ Flush the converted relations

        """

        # Always flush relations to be sure

        self.logger.info('Convert relations %s %s %s', etype_from, rtype, etype_to)

        self.flush_relations()

        if uri_label_from and etype_from not in self._initialized['uri_eid_inserted']:

            self.fill_uri_eid_table(etype_from, uri_label_from)

        if uri_label_to and etype_to not in self._initialized['uri_eid_inserted']:

            self.fill_uri_eid_table(etype_to, uri_label_to)

        if self._cnx.repo.schema.rschema(rtype).inlined:

            self.logger.warning("Can't insert inlined relation %s", rtype)

            return

        if uri_label_from and uri_label_to:

            sql = '''INSERT INTO %(r)s_relation (eid_from, eid_to) SELECT DISTINCT O1.eid, O2.eid

            FROM %(r)s_relation_iid_tmp AS T, uri_eid_%(ef)s as O1, uri_eid_%(et)s as O2

            WHERE O1.uri=T.uri_from AND O2.uri=T.uri_to AND NOT EXISTS (

            SELECT 1 FROM %(r)s_relation AS TT WHERE TT.eid_from=O1.eid AND TT.eid_to=O2.eid);

            '''

        elif uri_label_to:

            sql = '''INSERT INTO %(r)s_relation (eid_from, eid_to) SELECT DISTINCT

            CAST(T.uri_from AS INTEGER), O1.eid

            FROM %(r)s_relation_iid_tmp AS T, uri_eid_%(et)s as O1

            WHERE O1.uri=T.uri_to AND NOT EXISTS (

            SELECT 1 FROM %(r)s_relation AS TT WHERE

            TT.eid_from=CAST(T.uri_from AS INTEGER) AND TT.eid_to=O1.eid);

            '''

        elif uri_label_from:

            sql = '''INSERT INTO %(r)s_relation (eid_from, eid_to) SELECT DISTINCT O1.eid, T.uri_to

            O1.eid, CAST(T.uri_to AS INTEGER)

            FROM %(r)s_relation_iid_tmp AS T, uri_eid_%(ef)s as O1

            WHERE O1.uri=T.uri_from AND NOT EXISTS (

            SELECT 1 FROM %(r)s_relation AS TT WHERE

            TT.eid_from=O1.eid AND TT.eid_to=CAST(T.uri_to AS INTEGER));

            '''

        try:

            self.sql(sql % {'r': rtype.lower(),

                            'et': etype_to.lower() if etype_to else u'',

                            'ef': etype_from.lower() if etype_from else u''})

        except Exception as ex:

            self.logger.error("Can't insert relation %s: %s", rtype, ex)

        self.commit()

    ### SQL UTILITIES #########################################################

    def drop_and_store_indexes_constraints(self, tablename):

        # Drop indexes and constraints

        if not self._constraint_table_created:

            # Create a table to save the constraints

            # Allow reload even after crash

            sql = "CREATE TABLE dataio_constraints (origtable text, query text, type varchar(256))"

            self.sql(sql)

            self._constraint_table_created = True

        self._drop_table_constraints_indexes(tablename)

    def _drop_table_constraints_indexes(self, tablename):

        """ Drop and store table constraints and indexes """

        indexes, constraints = self._dbh.application_indexes_constraints(tablename)

        for name, query in constraints.items():

            sql = 'INSERT INTO dataio_constraints VALUES (%(e)s, %(c)s, %(t)s)'

            self.sql(sql, {'e': tablename, 'c': query, 't': 'constraint'})

            sql = 'ALTER TABLE %s DROP CONSTRAINT %s CASCADE' % (tablename, name)

            self.sql(sql)

        for name, query in indexes.items():

            sql = 'INSERT INTO dataio_constraints VALUES (%(e)s, %(c)s, %(t)s)'

            self.sql(sql, {'e': tablename, 'c': query, 't': 'index'})

            sql = 'DROP INDEX %s' % name

            self.sql(sql)

    def reapply_constraint_index(self, tablename):

        if not self._dbh.table_exists('dataio_constraints'):

            self.logger.info('The table dataio_constraints does not exist '

                             '(keep_index option should be True)')

            return

        sql = 'SELECT query FROM dataio_constraints WHERE origtable = %(e)s'

        crs = self.sql(sql, {'e': tablename})

        for query, in crs.fetchall():

            self.sql(query)

            self.sql('DELETE FROM dataio_constraints WHERE origtable = %(e)s '

                     'AND query = %(q)s', {'e': tablename, 'q': query})

    def _drop_metatables_constraints(self):

        """ Drop all the constraints for the meta data"""

        for tablename in ('created_by_relation', 'owned_by_relation',

                          'is_instance_of_relation', 'is_relation',

                          'entities'):

            self.drop_and_store_indexes_constraints(tablename)

    def _create_metatables_constraints(self):

        """ Create all the constraints for the meta data"""

        for tablename in ('entities',

                          'created_by_relation', 'owned_by_relation',

                          'is_instance_of_relation', 'is_relation'):

            # Indexes and constraints

            self.reapply_constraint_index(tablename)

    def init_relation_table(self, rtype):

        """ Get and remove all indexes for performance sake """

        # Create temporary table

        if not self.slave_mode and rtype not in self._initialized['rtypes']:

            sql = "CREATE TABLE %s_relation_tmp (eid_from integer, eid_to integer)" % rtype.lower()

            self.sql(sql)

            # Drop indexes and constraints

            tablename = '%s_relation' % rtype.lower()

            self.drop_and_store_indexes_constraints(tablename)

            # Push the etype in the initialized table for easier restart

            self.init_create_initialized_table()

            sql = 'INSERT INTO dataio_initialized VALUES (%(e)s, %(t)s)'

            self.sql(sql, {'e': rtype, 't': 'rtype'})

            # Mark rtype as "initialized" for faster check

            self._initialized['rtypes'].add(rtype)

    def init_create_initialized_table(self):

        """ Create the dataio initialized table

        """

        if not self._initialized_table_created:

            sql = "CREATE TABLE dataio_initialized (retype text, type varchar(128))"

            self.sql(sql)

            self._initialized_table_created = True

    def init_etype_table(self, etype):

        """ Add eid sequence to a particular etype table and

        remove all indexes for performance sake """

        if etype not in self._initialized['entities']:

            # Only for non-initialized etype and not slave mode store

            if not self.slave_mode:

                if self.eids_seq_range is None:

                    # Eids are directly set by the entities_id_seq.

                    # We attach this sequence to all the created etypes.

                    sql = ("ALTER TABLE cw_%s ALTER COLUMN cw_eid "

                           "SET DEFAULT nextval('entities_id_seq')" % etype.lower())

                    self.sql(sql)

                # Drop indexes and constraints

                tablename = 'cw_%s' % etype.lower()

                self.drop_and_store_indexes_constraints(tablename)

                # Push the etype in the initialized table for easier restart

                self.init_create_initialized_table()

                sql = 'INSERT INTO dataio_initialized VALUES (%(e)s, %(t)s)'

                self.sql(sql, {'e': etype, 't': 'etype'})

            # Mark etype as "initialized" for faster check

            self._initialized['entities'].add(etype)

    ### ENTITIES CREATION #####################################################

    def _get_eid_gen(self):

        """ 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, self.eids_seq_range)

            for eid in range(last_eid - self.eids_seq_range + 1, last_eid + 1):