# -*- coding: utf-8 -*-

"""This module provides tools to import tabular data.

:organization: Logilab

:copyright: 2001-2010 LOGILAB S.A. (Paris, FRANCE), license is LGPL v2.

:contact: http://www.logilab.fr/ -- mailto:contact@logilab.fr

:license: GNU Lesser General Public License, v2.1 - http://www.gnu.org/licenses

Example of use (run this with `cubicweb-ctl shell instance import-script.py`):

.. sourcecode:: python

  from cubicweb.devtools.dataimport import *

  # define data generators

  GENERATORS = []

  USERS = [('Prenom', 'firstname', ()),

           ('Nom', 'surname', ()),

           ('Identifiant', 'login', ()),

           ]

  def gen_users(ctl):

      for row in ctl.get_data('utilisateurs'):

          entity = mk_entity(row, USERS)

          entity['upassword'] = u'motdepasse'

          ctl.check('login', entity['login'], None)

          ctl.store.add('CWUser', entity)

          email = {'address': row['email']}

          ctl.store.add('EmailAddress', email)

          ctl.store.relate(entity['eid'], 'use_email', email['eid'])

          ctl.store.rql('SET U in_group G WHERE G name "users", U eid %(x)s', {'x':entity['eid']})

  CHK = [('login', check_doubles, 'Utilisateurs Login',

          'Deux utilisateurs ne devraient pas avoir le même login.'),

         ]

  GENERATORS.append( (gen_users, CHK) )

  # create controller

  ctl = CWImportController(RQLObjectStore())

  ctl.askerror = 1

  ctl.generators = GENERATORS

  ctl.store._checkpoint = checkpoint

  ctl.store._rql = rql

  ctl.data['utilisateurs'] = lazytable(utf8csvreader(open('users.csv')))

  # run

  ctl.run()

  sys.exit(0)

.. BUG fichier à une colonne pose un problème de parsing

.. TODO rollback()

"""

__docformat__ = "restructuredtext en"

import sys

import csv

import traceback

import os.path as osp

from StringIO import StringIO

from copy import copy

from logilab.common import shellutils

from logilab.common.date import strptime

from logilab.common.decorators import cached

from logilab.common.deprecation import deprecated

def ucsvreader_pb(filepath, encoding='utf-8', separator=',', quote='"',

                  skipfirst=False, withpb=True):

    """same as ucsvreader but a progress bar is displayed as we iter on rows"""

    if not osp.exists(filepath):

        raise Exception("file doesn't exists: %s" % filepath)

    rowcount = int(shellutils.Execute('wc -l "%s"' % filepath).out.strip().split()[0])

    if skipfirst:

        rowcount -= 1

    if withpb:

        pb = shellutils.ProgressBar(rowcount, 50)

    for urow in ucsvreader(file(filepath), encoding, separator, quote, skipfirst):

        yield urow

        if withpb:

            pb.update()

    print ' %s rows imported' % rowcount

def ucsvreader(stream, encoding='utf-8', separator=',', quote='"',

               skipfirst=False):

    """A csv reader that accepts files with any encoding and outputs unicode

    strings

    """

    it = iter(csv.reader(stream, delimiter=separator, quotechar=quote))

    if skipfirst:

        it.next()

    for row in it:

        yield [item.decode(encoding) for item in row]

def commit_every(nbit, store, it):

    for i, x in enumerate(it):

        yield x

        if nbit is not None and i % nbit:

            store.checkpoint()

    if nbit is not None:

        store.checkpoint()

def lazytable(reader):

    """The first row is taken to be the header of the table and

    used to output a dict for each row of data.

    >>> data = lazytable(utf8csvreader(open(filename)))

    """

    header = reader.next()

    for row in reader:

        yield dict(zip(header, row))

def mk_entity(row, map):

    """Return a dict made from sanitized mapped values.

    ValidationError can be raised on unexpected values found in checkers

    >>> row = {'myname': u'dupont'}

    >>> map = [('myname', u'name', (capitalize_if_unicase,))]

    >>> mk_entity(row, map)

    {'name': u'Dupont'}

    >>> row = {'myname': u'dupont', 'optname': u''}

    >>> map = [('myname', u'name', (capitalize_if_unicase,)),

    ...        ('optname', u'MARKER', (optional,))]

    >>> mk_entity(row, map)

    {'name': u'Dupont'}

    """

    res = {}

    assert isinstance(row, dict)

    assert isinstance(map, list)

    for src, dest, funcs in map:

        assert not (required in funcs and optional in funcs), \

               "optional and required checks are exclusive"

        res[dest] = row[src]

        try:

            for func in funcs:

                res[dest] = func(res[dest])

                if res[dest] is None:

                    break

        except ValueError, err:

            raise ValueError('error with %r field: %s' % (src, err))

    return res

# user interactions ############################################################

def tell(msg):

    print msg

def confirm(question):

    """A confirm function that asks for yes/no/abort and exits on abort."""

    answer = shellutils.ASK.ask(question, ('Y', 'n', 'abort'), 'Y')

    if answer == 'abort':

        sys.exit(1)

    return answer == 'Y'

class catch_error(object):

    """Helper for @contextmanager decorator."""

    def __init__(self, ctl, key='unexpected error', msg=None):

        self.ctl = ctl

        self.key = key

        self.msg = msg

    def __enter__(self):

        return self

    def __exit__(self, type, value, traceback):

        if type is not None:

            if issubclass(type, (KeyboardInterrupt, SystemExit)):

                return # re-raise

            if self.ctl.catcherrors:

                self.ctl.record_error(self.key, None, type, value, traceback)

                return True # silent

# base sanitizing/coercing functions ###########################################

def optional(value):

    """validation error will not been raised if you add this checker in chain"""

    if value:

        return value

    return None

def required(value):

    """raise ValueError is value is empty

    This check should be often found in last position in the chain.

    """

    if value:

        return value

    raise ValueError("required")

def todatetime(format='%d/%m/%Y'):

    """return a transformation function to turn string input value into a

    `datetime.datetime` instance, using given format.

    Follow it by `todate` or `totime` functions from `logilab.common.date` if

    you want a `date`/`time` instance instead of `datetime`.

    """

    def coerce(value):

        return strptime(value, format)

    return coerce

def call_transform_method(methodname, *args, **kwargs):

    """return value returned by calling the given method on input"""

    def coerce(value):

        return getattr(value, methodname)(*args, **kwargs)

    return coerce

def call_check_method(methodname, *args, **kwargs):

    """check value returned by calling the given method on input is true,

    else raise ValueError

    """

    def check(value):

        if getattr(value, methodname)(*args, **kwargs):

            return value

        raise ValueError('%s not verified on %r' % (methodname, value))

    return check

# base integrity checking functions ############################################

def check_doubles(buckets):

    """Extract the keys that have more than one item in their bucket."""

    return [(k, len(v)) for k, v in buckets.items() if len(v) > 1]

def check_doubles_not_none(buckets):

    """Extract the keys that have more than one item in their bucket."""

    return [(k, len(v)) for k, v in buckets.items()

            if k is not None and len(v) > 1]

# object stores #################################################################

class ObjectStore(object):

    """Store objects in memory for *faster* validation (development mode)

    But it will not enforce the constraints of the schema and hence will miss some problems

    >>> store = ObjectStore()

    >>> user = {'login': 'johndoe'}

    >>> store.add('CWUser', user)

    >>> group = {'name': 'unknown'}

    >>> store.add('CWUser', group)

    >>> store.relate(user['eid'], 'in_group', group['eid'])

    """

    def __init__(self):

        self.items = []

        self.eids = {}

        self.types = {}

        self.relations = set()

        self.indexes = {}

        self._rql = None

        self._checkpoint = None

    def _put(self, type, item):

        self.items.append(item)

        return len(self.items) - 1

    def add(self, type, item):

        assert isinstance(item, dict), 'item is not a dict but a %s' % type(item)

        eid = item['eid'] = self._put(type, item)

        self.eids[eid] = item

        self.types.setdefault(type, []).append(eid)

    def relate(self, eid_from, rtype, eid_to, inlined=False):

        """Add new relation (reverse type support is available)

        >>> 1,2 = eid_from, eid_to

        >>> self.relate(eid_from, 'in_group', eid_to)

        1, 'in_group', 2

        >>> self.relate(eid_from, 'reverse_in_group', eid_to)

        2, 'in_group', 1

        """

        if rtype.startswith('reverse_'):

            eid_from, eid_to = eid_to, eid_from

            rtype = rtype[8:]

        relation = eid_from, rtype, eid_to

        self.relations.add(relation)

        return relation

    def build_index(self, name, type, func=None):

        index = {}

        if func is None or not callable(func):

            func = lambda x: x['eid']

        for eid in self.types[type]:

            index.setdefault(func(self.eids[eid]), []).append(eid)

        assert index, "new index '%s' cannot be empty" % name

        self.indexes[name] = index

    def build_rqlindex(self, name, type, key, rql, rql_params=False, func=None):

        """build an index by rql query

        rql should return eid in first column

        ctl.store.build_index('index_name', 'users', 'login', 'Any U WHERE U is CWUser')

        """

        rset = self.rql(rql, rql_params or {})

        for entity in rset.entities():

            getattr(entity, key) # autopopulate entity with key attribute

            self.eids[entity.eid] = dict(entity)

            if entity.eid not in self.types.setdefault(type, []):

                self.types[type].append(entity.eid)

        assert self.types[type], "new index type '%s' cannot be empty (0 record found)" % type

        # Build index with specified key

        func = lambda x: x[key]

        self.build_index(name, type, func)

    def fetch(self, name, key, unique=False, decorator=None):

        """

            decorator is a callable method or an iterator of callable methods (usually a lambda function)

            decorator=lambda x: x[:1] (first value is returned)

            We can use validation check function available in _entity

        """

        eids = self.indexes[name].get(key, [])

        if decorator is not None:

            if not hasattr(decorator, '__iter__'):

                decorator = (decorator,)

            for f in decorator:

                eids = f(eids)

        if unique:

            assert len(eids) == 1, u'expected a single one value for key "%s" in index "%s". Got %i' % (key, name, len(eids))

            eids = eids[0] # FIXME maybe it's better to keep an iterator here ?

        return eids

    def find(self, type, key, value):

        for idx in self.types[type]:

            item = self.items[idx]

            if item[key] == value:

                yield item

    def rql(self, *args):

        if self._rql is not None:

            return self._rql(*args)

    def checkpoint(self):

        pass

    @property

    def nb_inserted_entities(self):

        return len(self.eids)

    @property

    def nb_inserted_types(self):

        return len(self.types)

    @property

    def nb_inserted_relations(self):

        return len(self.relations)

    @deprecated('[3.6] get_many() deprecated. Use fetch() instead')

    def get_many(self, name, key):

        return self.fetch(name, key, unique=False)

    @deprecated('[3.6] get_one() deprecated. Use fetch(..., unique=True) instead')

    def get_one(self, name, key):

        return self.fetch(name, key, unique=True)

class RQLObjectStore(ObjectStore):

    """ObjectStore that works with an actual RQL repository (production mode)"""

    _rql = None # bw compat

    def __init__(self, session=None, checkpoint=None):

        ObjectStore.__init__(self)

        if session is not None:

            if not hasattr(session, 'set_pool'):

                # connection

                cnx = session

                session = session.request()

                session.set_pool = lambda : None

                checkpoint = checkpoint or cnx.commit

            else:

                session.set_pool()

            self.session = session

            self._checkpoint = checkpoint or session.commit

        elif checkpoint is not None:

            self._checkpoint = checkpoint

            # XXX .session

    def checkpoint(self):

        self._checkpoint()

        self.session.set_pool()

    def rql(self, *args):

        if self._rql is not None:

            return self._rql(*args)

        return self.session.execute(*args)

    def create_entity(self, *args, **kwargs):