# 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.server.schema2sql import rschema_has_table

from cubicweb.schema import PURE_VIRTUAL_RTYPES

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.prepare_insert_entity('Person', subj_iid_attribute=person_iid, ...)

       store.prepare_insert_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.prepare_insert_entity('Person',

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

                                   name='Toto')

       store.prepare_insert_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.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.commit()

       store.finish()

    """

    # 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

    # max size of the iid, used to create the iid_eid conversion table

    iid_maxsize = 1024

    def __init__(self, cnx,

                 on_commit_callback=None, on_rollback_callback=None,

                 slave_mode=False,

                 source=None):

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

        - cnx: CubicWeb cnx

        """

        super(MassiveObjectStore, self).__init__(cnx)

        self.logger = logging.getLogger('dataimport.massive_store')

        self._cnx = cnx

        self.sql = cnx.system_sql

        self._data_uri_relations = defaultdict(list)

        # etypes for which we have a uri_eid_%(etype)s table

        self._init_uri_eid = set()

        # etypes for which we have a uri_eid_%(e)s_idx index

        self._uri_eid_inserted = set()

        # set of rtypes for which we have a %(rtype)s_relation_iid_tmp table

        self._uri_rtypes = set()

        # set of etypes whose tables are created

        self._entities = set()

        # set of rtypes for which we have a %(rtype)s_relation_tmp table

        self._rtypes = set()

        self.slave_mode = slave_mode

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

        pg_schema = cnx.repo.config.system_source_config.get('db-namespace') or '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._count_cwuri = 0

        self.on_commit_callback = on_commit_callback

        self.on_rollback_callback = on_rollback_callback

        # Do our meta tables already exist?

        self._init_massive_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_all_constraints()

            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 _drop_all_constraints(self):

        schema = self._cnx.vreg.schema

        tables = ['cw_%s' % etype.type.lower()

                  for etype in schema.entities() if not etype.final]

        for rschema in schema.relations():

            if rschema.inlined:

                continue

            elif rschema_has_table(rschema, skip_relations=PURE_VIRTUAL_RTYPES):

                tables.append('%s_relation' % rschema.type.lower())

        tables.append('entities')

        for tablename in tables:

            self._store_and_drop_constraints(tablename)

    def _store_and_drop_constraints(self, tablename):

        if not self._constraint_table_created:

            # Create a table to save the constraints

            # Allow reload even after crash

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

            self.sql(sql)

            self._constraint_table_created = True

        constraints = self._dbh.application_constraints(tablename)

        for name, query in constraints.items():

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

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

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

            self.sql(sql)

    def reapply_all_constraints(self):

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

            self.logger.info('The table cwmassive_constraints does not exist')

            return

        sql = 'SELECT query FROM cwmassive_constraints WHERE type = %(t)s'

        crs = self.sql(sql, {'t': 'constraint'})

        for query, in crs.fetchall():

            self.sql(query)

            self.sql('DELETE FROM cwmassive_constraints WHERE type = %(t)s '

                     'AND query = %(q)s', {'t': 'constraint', 'q': query})

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

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

        # Create an uri_eid table for each etype for a better

        # control of which etype is concerned by a particular

        # possibly multivalued relation.

        for etype in (etype_from, etype_to):

            if etype and etype not in self._init_uri_eid:

                self._init_uri_eid_table(etype)

        if rtype not in self._uri_rtypes:

            # Create the temporary table

            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._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._init_uri_eid.add(etype)

    def _init_massive_metatables(self):

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

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

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

        self._metadata_table_created = self._dbh.table_exists('cwmassive_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] = []

    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._uri_eid_inserted.add(etype)

    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._uri_eid_inserted:

            self.fill_uri_eid_table(etype_from, uri_label_from)

        if uri_label_to and etype_to not in self._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)

    ### 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 cwmassive_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 cwmassive_constraints VALUES (%(e)s, %(c)s, %(t)s)'

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

            self.sql(sql)

        for name, query in indexes.items():

            sql = 'INSERT INTO cwmassive_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('cwmassive_constraints'):

            self.logger.info('The table cwmassive_constraints does not exist')

            return

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

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

        for query, in crs.fetchall():

            self.sql(query)

            self.sql('DELETE FROM cwmassive_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._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 cwmassive_initialized VALUES (%(e)s, %(t)s)'

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

            # Mark rtype as "initialized" for faster check

            self._rtypes.add(rtype)

    def init_create_initialized_table(self):

        """ Create the cwmassive initialized table

        """

        if not self._initialized_table_created:

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

            self.sql(sql)