# coding: utf-8

# copyright 2015-2016 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 itertools import chain

from six.moves import range

import pytz

from yams.constraints import SizeConstraint

from cubicweb.schema import PURE_VIRTUAL_RTYPES

from cubicweb.server.schema2sql import rschema_has_table

from cubicweb.server.sqlutils import SQL_PREFIX

from cubicweb.dataimport import stores, pgstore

from cubicweb.utils import make_uid

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)

       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()

    """

    # 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,

                 eids_seq_range=10000):

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

        - cnx: CubicWeb cnx

        - eids_seq_range: size of eid range reserved by the store for each batch

        """

        super(MassiveObjectStore, self).__init__(cnx)

        self.on_commit_callback = on_commit_callback

        self.on_rollback_callback = on_rollback_callback

        self.slave_mode = slave_mode

        self.eids_seq_range = eids_seq_range

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

        self.sql = cnx.system_sql

        self.schema = self._cnx.vreg.schema

        self.default_values = get_default_values(self.schema)

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

        self._dbh = PGHelper(cnx)

        cnx.read_security = False

        cnx.write_security = False

        self._data_entities = defaultdict(list)

        self._data_relations = defaultdict(list)

        self._initialized = set()

        # uri handling

        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()

        self._now = datetime.now(pytz.utc)

        self._default_cwuri = make_uid('_auto_generated')

        if not self.slave_mode:

            # drop constraint and metadata table, they will be recreated when self.finish() is

            # called

            self._drop_all_constraints()

            self._drop_metatables_constraints()

        if source is None:

            source = cnx.repo.system_source

        self.source = source

    # URI related things #######################################################

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

                self.sql('CREATE TABLE IF NOT EXISTS uri_eid_%(e)s'

                         '(uri character varying(%(size)s), eid integer)'

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

        if rtype not in self._uri_rtypes:

            # Create the temporary table

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

                self.sql('CREATE TABLE IF NOT EXISTS %(r)s_relation_iid_tmp'

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

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

                self._uri_rtypes.add(rtype)

            else:

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

    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})

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

        """

        if etype not in self._uri_eid_inserted:

            self._uri_eid_inserted.add(etype)

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

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

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

            self.sql('CREATE INDEX uri_eid_%(e)s_idx ON uri_eid_%(e)s(uri)'

                     % {'e': etype.lower()})

    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:

            self.fill_uri_eid_table(etype_from, uri_label_from)

        if uri_label_to:

            self.fill_uri_eid_table(etype_to, uri_label_to)

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

        etypes_tables = ('cw_%s' % eschema.type.lower() for eschema in self.schema.entities()

                         if not eschema.final)

        rtypes_tables = ('%s_relation' % rschema.type.lower() for rschema in self.schema.relations()

                         if rschema_has_table(rschema, skip_relations=PURE_VIRTUAL_RTYPES))

        # Create a table to save the constraints, it allows reloading even after crash

        self.sql('CREATE TABLE IF NOT EXISTS cwmassive_constraints'

                 '(origtable text, query text, type varchar(256))')

        for tablename in chain(etypes_tables, rtypes_tables, ('entities',)):

            constraints = self._dbh.table_constraints(tablename)

            for name, query in constraints.items():

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

                         {'e': tablename, 'c': query, 't': 'constraint'})

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

    def _reapply_all_constraints(self):

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

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

            return

        cu = self.sql("SELECT query FROM cwmassive_constraints WHERE type='constraint'")

        for query, in cu.fetchall():

            self.sql(query)

            self.sql("DELETE FROM cwmassive_constraints WHERE type='constraint' AND query=%(q)s",

                     {'q': query})

    def drop_and_store_indexes(self, tablename):

        """Drop indexes and constraints"""

        # Create a table to save the constraints, it allows reloading even after crash

        self.sql('CREATE TABLE IF NOT EXISTS cwmassive_constraints'

                 '(origtable text, query text, type varchar(256))')

        indexes = self._dbh.table_indexes(tablename)

        for name, query in indexes.items():

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

                     {'e': tablename, 'c': query, 't': 'index'})

            self.sql('DROP INDEX %s' % name)

    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

        cu = self.sql('SELECT query FROM cwmassive_constraints WHERE origtable = %(e)s',

                      {'e': tablename})

        for query, in cu.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(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 restart_eid_sequence(self, start_eid):

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

            'entities_id_seq', initial_value=start_eid))

        self._cnx.commit()

    # 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):

                yield eid

    def _apply_default_values(self, etype, kwargs):

        """Apply the default values for a given etype, attribute and value."""

        default_values = self.default_values[etype]

        missing_keys = set(default_values) - set(kwargs)

        kwargs.update((key, default_values[key]) for key in missing_keys)

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

    def prepare_insert_entity(self, etype, **kwargs):

        """Given an entity type, attributes and inlined relations, returns the inserted entity's

        eid.

        """

        # Add meta data if not given

        if 'modification_date' not in kwargs:

            kwargs['modification_date'] = self._now

        if 'creation_date' not in kwargs:

            kwargs['creation_date'] = self._now

        if 'cwuri' not in kwargs:

            kwargs['cwuri'] = self._default_cwuri + str(self._count_cwuri)

            self._count_cwuri += 1

        if 'eid' not in kwargs:

            # If eid is not given and the eids sequence is set,

            # use the value from the sequence

            kwargs['eid'] = self.get_next_eid()

        self._apply_default_values(etype, kwargs)

        self._data_entities[etype].append(kwargs)

        return kwargs.get('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``.

        """

        self._data_relations[rtype].append({'eid_from': eid_from, 'eid_to': eid_to})

    def flush(self):

        """Flush the data"""

        self.flush_entities()

        self.flush_internal_relations()

        self.flush_relations()

    def commit(self):

        """Commit the database transaction."""

        self.on_commit()

        super(MassiveObjectStore, self).commit()