#!/usr/bin/env python

# Maps one datasource to another, using a map spreadsheet if needed

# Exit status is the # of errors in the import, up to the maximum exit status

# For outputting an XML file to a PostgreSQL database, use the general format of

# http://vegbank.org/vegdocs/xml/vegbank_example_ver1.0.2.xml

import csv

import os.path

import sys

import xml.dom.minidom as minidom

sys.path.append(os.path.dirname(__file__)+"/../lib")

import exc

import opts

import Parser

import profiling

import sql

import strings

import term

import util

import xpath

import xml_dom

import xml_func

def metadata_value(name):

    if util.is_str(name) and name.startswith(':'): return name[1:]

    else: return None

def main_():

    env_names = []

    def usage_err():

        raise SystemExit('Usage: '+opts.env_usage(env_names, True)+' '

            +sys.argv[0]+' [map_path...] [<input] [>output]')

    # Get db config from env vars

    db_config_names = ['engine', 'host', 'user', 'password', 'database']

    def get_db_config(prefix):

        return opts.get_env_vars(db_config_names, prefix, env_names)

    in_db_config = get_db_config('in')

    out_db_config = get_db_config('out')

    in_is_db = 'engine' in in_db_config

    out_is_db = 'engine' in out_db_config

    # Get other config from env vars

    end = util.cast(int, opts.get_env_var('n', None, env_names))

    start = util.cast(int, opts.get_env_var('start', '0', env_names))

    if end != None: end += start

    test = opts.env_flag('test', False, env_names)

    commit = opts.env_flag('commit', False, env_names) and not test

        # never commit in test mode

    redo = opts.env_flag('redo', test, env_names) and not commit

        # never redo in commit mode (manually run `make empty_db` instead)

    debug = opts.env_flag('debug', False, env_names)

    sql.run_raw_query.debug = debug

    verbose = debug or opts.env_flag('verbose', False, env_names)

    opts.get_env_var('profile_to', None, env_names) # add to env_names

    # Logging

    def log(msg, on=verbose):

        if on: sys.stderr.write(msg)

    def log_start(action, on=verbose): log(action+'...\n', on)

    # Parse args

    map_paths = sys.argv[1:]

    if map_paths == []:

        if in_is_db or not out_is_db: usage_err()

        else: map_paths = [None]

    def connect_db(db_config):

        log_start('Connecting to '+sql.db_config_str(db_config))

        return sql.connect(db_config)

    ex_tracker = exc.ExPercentTracker(iter_text='row')

    profiler = profiling.ItersProfiler(start_now=True, iter_text='row')

    doc = xml_dom.create_doc()

    root = doc.documentElement

    out_is_xml_ref = [False]

    in_label_ref = [None]

    def update_in_label():

        if in_label_ref[0] != None:

            xpath.get(root, '/_ignore/inLabel="'+in_label_ref[0]+'"', True)

    def prep_root():

        root.clear()

        update_in_label()

    prep_root()

    def process_input(root, row_ready, map_path):

        '''Inputs datasource to XML tree, mapping if needed'''

        # Load map header

        in_is_xpaths = True

        out_is_xpaths = True

        out_label = None

        if map_path != None:

            metadata = []

            mappings = []

            stream = open(map_path, 'rb')

            reader = csv.reader(stream)

            in_label, out_label = reader.next()[:2]

            def split_col_name(name):

                name, sep, root = name.partition(':')

                return name, sep != '', root

            in_label, in_is_xpaths, in_root = split_col_name(in_label)

            in_label_ref[0] = in_label

            update_in_label()

            out_label, out_is_xpaths, out_root = split_col_name(out_label)

            has_types = out_root.startswith('/*s/') # outer elements are types

            for row in reader:

                in_, out = row[:2]

                if out != '':

                    if out_is_xpaths: out = xpath.parse(out_root+out)

                    mappings.append((in_, out))

            stream.close()

            root.ownerDocument.documentElement.tagName = out_label

        in_is_xml = in_is_xpaths and not in_is_db

        out_is_xml_ref[0] = out_is_xpaths and not out_is_db

        if in_is_xml:

            doc0 = minidom.parse(sys.stdin)

            doc0_root = doc0.documentElement

            if out_label == None: out_label = doc0_root.tagName

        def process_rows(process_row, rows):

            '''Processes input rows

            @param process_row(in_row, i)

            '''

            i = -1 # in case for loop does not execute

            for i, row in enumerate(rows):

                if i < start: continue

                if end != None and i >= end: break

                process_row(row, i)

                row_ready(i, row)

            row_ct = i-start+1

            return row_ct

        def map_rows(get_value, rows):

            '''Maps input rows

            @param get_value(in_, row):str

            '''

            def process_row(row, i):

                row_id = str(i)

                for in_, out in mappings:

                    value = metadata_value(in_)

                    if value == None:

                        log_start('Getting '+str(in_), debug)

                        value = get_value(in_, row)

                    if value != None: xpath.put_obj(root, out, row_id,

                        has_types, strings.cleanup(value))

            return process_rows(process_row, rows)

        if map_path == None:

            iter_ = xml_dom.NodeElemIter(doc0_root)

            util.skip(iter_, xml_dom.is_text) # skip metadata

            row_ct = process_rows(lambda row, i: root.appendChild(row), iter_)

        elif in_is_db:

            assert in_is_xpaths

            import db_xml

            in_root_xml = xpath.path2xml(in_root)

            for i, mapping in enumerate(mappings):

                in_, out = mapping

                if metadata_value(in_) == None:

                    mappings[i] = (xpath.path2xml(in_root+'/'+in_), out)

            in_db = connect_db(in_db_config)

            in_pkeys = {}

            def get_value(in_, row):

                pkey, = row

                in_ = in_.cloneNode(True) # don't modify orig value!

                xml_dom.set_id(xpath.get(in_, in_root)[0], pkey)

                value = sql.value_or_none(db_xml.get(in_db, in_, in_pkeys))

                if value != None: return str(value)

                else: return None

            row_ct = map_rows(get_value, sql.rows(db_xml.get(in_db, in_root_xml,

                in_pkeys, end, 0)))

            in_db.close()

        elif in_is_xml:

            def get_value(in_, row):

                nodes = xpath.get(row, in_, allow_rooted=False)

                if nodes != []: return xml_dom.value(nodes[0])

                else: return None

            rows = xpath.get(doc0_root, in_root, limit=end)

            if rows == []: raise SystemExit('Map error: Root "'+in_root

                +'" not found in input')

            row_ct = map_rows(get_value, rows)

        else: # input is CSV

            map_ = dict(mappings)

            reader = csv.reader(sys.stdin)

            cols = reader.next()

            col_idxs = dict([(value, idx) for idx, value in enumerate(cols)])

            for i, mapping in enumerate(mappings):

                in_, out = mapping

                if metadata_value(in_) == None:

                    try: mappings[i] = (col_idxs[in_], out)

                    except KeyError: pass

            def get_value(in_, row):

                value = row[in_]

                if value != '': return value

                else: return None

            row_ct = map_rows(get_value, reader)

        return row_ct

    def process_inputs(root, row_ready):

        row_ct = 0

        for map_path in map_paths:

            row_ct += process_input(root, row_ready, map_path)

        return row_ct

    if out_is_db:

        import db_xml

        out_db = connect_db(out_db_config)

        out_pkeys = {}

        try:

            if redo: sql.empty_db(out_db)

            row_ins_ct_ref = [0]

            def row_ready(row_num, input_row):

                def on_error(e):

                    exc.add_msg(e, term.emph('row #:')+' '+str(row_num))

                    exc.add_msg(e, term.emph('input row:')+'\n'+str(input_row))

                    exc.add_msg(e, term.emph('output row:')+'\n'+str(root))

                    ex_tracker.track(e)

                xml_func.process(root, on_error)

                if not xml_dom.is_empty(root):

                    assert xml_dom.has_one_child(root)

                    try:

                        sql.with_savepoint(out_db,

                            lambda: db_xml.put(out_db, root.firstChild,

                                out_pkeys, row_ins_ct_ref, on_error))

                        if commit: out_db.commit()

                    except sql.DatabaseErrors, e: on_error(e)

                prep_root()

            row_ct = process_inputs(root, row_ready)

            sys.stdout.write('Inserted '+str(row_ins_ct_ref[0])+

                ' new rows into database\n')

        finally:

            out_db.rollback()

            out_db.close()

    else:

        def on_error(e): ex_tracker.track(e)

        def row_ready(row_num, input_row): pass

        row_ct = process_inputs(root, row_ready)

        xml_func.process(root, on_error)

        if out_is_xml_ref[0]:

            doc.writexml(sys.stdout, **xml_dom.prettyxml_config)

        else: # output is CSV

            raise NotImplementedError('CSV output not supported yet')

    profiler.stop(row_ct)

    ex_tracker.add_iters(row_ct)

    if verbose:

        sys.stderr.write('Processed '+str(row_ct)+' input rows\n')

        sys.stderr.write(profiler.msg()+'\n')

        sys.stderr.write(ex_tracker.msg()+'\n')

    ex_tracker.exit()

def main():

    try: main_()

    except Parser.SyntaxException, e: raise SystemExit(str(e))

if __name__ == '__main__':

    profile_to = opts.get_env_var('profile_to', None)

    if profile_to != None:

        import cProfile

        cProfile.run(main.func_code, profile_to)

    else: main()