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

# copyright 2003-2012 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/>.

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

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

.. sourcecode:: python

  from cubicweb.dataimport import *

  # define data generators

  GENERATORS = []

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

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

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

           ]

  def gen_users(ctl):

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

          entity = mk_entity(row, USERS)

          entity['upassword'] = 'motdepasse'

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

          entity = ctl.store.create_entity('CWUser', **entity)

          email = ctl.store.create_entity('EmailAddress', address=row['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

  if 'cnx' in globals():

      ctl = CWImportController(RQLObjectStore(cnx))

  else:

      print 'debug mode (not connected)'

      print 'run through cubicweb-ctl shell to access an instance'

      ctl = CWImportController(ObjectStore())

  ctl.askerror = 1

  ctl.generators = GENERATORS

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

  # run

  ctl.run()

.. BUG file with one column are not parsable

.. TODO rollback() invocation is not possible yet

"""

__docformat__ = "restructuredtext en"

import csv

import sys

import threading

import traceback

import cPickle

import os.path as osp

from collections import defaultdict

from contextlib import contextmanager

from copy import copy

from datetime import date, datetime

from time import asctime

from StringIO import StringIO

from logilab.common import shellutils, attrdict

from logilab.common.date import strptime

from logilab.common.decorators import cached

from logilab.common.deprecation import deprecated

from cubicweb import QueryError

from cubicweb.utils import make_uid

from cubicweb.schema import META_RTYPES, VIRTUAL_RTYPES

from cubicweb.server.edition import EditedEntity

from cubicweb.server.sqlutils import SQL_PREFIX

from cubicweb.server.utils import eschema_eid

def count_lines(stream_or_filename):

    if isinstance(stream_or_filename, basestring):

        f = open(stream_or_filename)

    else:

        f = stream_or_filename

        f.seek(0)

    for i, line in enumerate(f):

        pass

    f.seek(0)

    return i+1

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

                  skipfirst=False, withpb=True):

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

    if isinstance(stream_or_path, basestring):

        if not osp.exists(stream_or_path):

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

        stream = open(stream_or_path)

    else:

        stream = stream_or_path

    rowcount = count_lines(stream)

    if skipfirst:

        rowcount -= 1

    if withpb:

        pb = shellutils.ProgressBar(rowcount, 50)

    for urow in ucsvreader(stream, 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 callfunc_every(func, number, iterable):

    """yield items of `iterable` one by one and call function `func`

    every `number` iterations. Always call function `func` at the end.

    """

    for idx, item in enumerate(iterable):

        yield item

        if not idx % number:

            func()

    func()

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(ucsvreader(open(filename)))

    """

    header = reader.next()

    for row in reader:

        yield dict(zip(header, row))

def lazydbtable(cu, table, headers, orderby=None):

    """return an iterator on rows of a sql table. On each row, fetch columns

    defined in headers and return values as a dictionary.

    >>> data = lazydbtable(cu, 'experimentation', ('id', 'nickname', 'gps'))

    """

    sql = 'SELECT %s FROM %s' % (','.join(headers), table,)

    if orderby:

        sql += ' ORDER BY %s' % ','.join(orderby)

    cu.execute(sql)

    while True:

        row = cu.fetchone()

        if row is None:

            break

        yield dict(zip(headers, row))

def mk_entity(row, map):

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

    ValueError can be raised on unexpected values found in checkers

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

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

    >>> mk_entity(row, map)

    {'name': u'Dupont'}

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

    >>> map = [('myname', u'name', (call_transform_method('title'),)),

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

    >>> mk_entity(row, map)

    {'name': u'Dupont', 'optname': None}

    """

    res = {}

    assert isinstance(row, dict)

    assert isinstance(map, list)

    for src, dest, funcs in map:

        try:

            res[dest] = row[src]

        except KeyError:

            continue

        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)), None, sys.exc_info()[-1]

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

    """checker to filter optional field

    If value is undefined (ex: empty string), return None that will

    break the checkers validation chain

    General use is to add 'optional' check in first condition to avoid

    ValueError by further checkers

    >>> MAPPER = [(u'value', 'value', (optional, int))]

    >>> row = {'value': u'XXX'}

    >>> mk_entity(row, MAPPER)

    {'value': None}

    >>> row = {'value': u'100'}

    >>> mk_entity(row, MAPPER)

    {'value': 100}

    """

    if value:

        return value

    return None

def required(value):

    """raise ValueError if 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]

# sql generator utility functions #############################################

def _import_statements(sql_connect, statements, nb_threads=3,

                       dump_output_dir=None,

                       support_copy_from=True, encoding='utf-8'):

    """

    Import a bunch of sql statements, using different threads.

    """

    try:

        chunksize = (len(statements) / nb_threads) + 1

        threads = []

        for i in xrange(nb_threads):

            chunks = statements[i*chunksize:(i+1)*chunksize]

            thread = threading.Thread(target=_execmany_thread,

                                      args=(sql_connect, chunks,

                                            dump_output_dir,

                                            support_copy_from,

                                            encoding))

            thread.start()

            threads.append(thread)

        for t in threads:

            t.join()

    except Exception:

        print 'Error in import statements'

def _execmany_thread_not_copy_from(cu, statement, data, table=None,

                                   columns=None, encoding='utf-8'):

    """ Execute thread without copy from

    """

    cu.executemany(statement, data)

def _execmany_thread_copy_from(cu, statement, data, table,

                               columns, encoding='utf-8'):

    """ Execute thread with copy from

    """

    buf = _create_copyfrom_buffer(data, columns, encoding)

    if buf is None:

        _execmany_thread_not_copy_from(cu, statement, data)

    else:

        if columns is None:

            cu.copy_from(buf, table, null='NULL')

        else:

            cu.copy_from(buf, table, null='NULL', columns=columns)

def _execmany_thread(sql_connect, statements, dump_output_dir=None,

                     support_copy_from=True, encoding='utf-8'):

    """

    Execute sql statement. If 'INSERT INTO', try to use 'COPY FROM' command,

    or fallback to execute_many.

    """

    if support_copy_from:

        execmany_func = _execmany_thread_copy_from

    else:

        execmany_func = _execmany_thread_not_copy_from

    cnx = sql_connect()

    cu = cnx.cursor()

    try:

        for statement, data in statements:

            table = None

            columns = None

            try:

                if not statement.startswith('INSERT INTO'):

                    cu.executemany(statement, data)

                    continue

                table = statement.split()[2]

                if isinstance(data[0], (tuple, list)):

                    columns = None

                else:

                    columns = data[0].keys()

                execmany_func(cu, statement, data, table, columns, encoding)

            except Exception:

                print 'unable to copy data into table %s', table

                # Error in import statement, save data in dump_output_dir

                if dump_output_dir is not None:

                    pdata = {'data': data, 'statement': statement,

                             'time': asctime(), 'columns': columns}

                    filename = make_uid()

                    try:

                        with open(osp.join(dump_output_dir,

                                           '%s.pickle' % filename), 'w') as fobj:

                            fobj.write(cPickle.dumps(pdata))

                    except IOError:

                        print 'ERROR while pickling in', dump_output_dir, filename+'.pickle'

                        pass

                cnx.rollback()

                raise

    finally:

        cnx.commit()

        cu.close()