#!/usr/bin/env python

import os
import sys
import optparse
from sql2graph.export2 import SQL2GraphExporter
from musicbrainz_schema import mbschema, mbentities, make_link_entity_list

# ----------------------------------------------------------------------

DEFAULT_NODES_FILE = 'D:/musicbrainz__nodes__full.csv'
DEFAULT_RELS_FILE = 'D:/musicbrainz__rels__full.csv'
MULTIPLE_FILES = False

option_parser = optparse.OptionParser()

option_parser.add_option("--exclude",
                         dest="excluded_entities",
                         help="Entities to exclude",
                         default=None)
option_parser.add_option("--include",
                         dest="included_entities",
                         help="Entities to include",
                         default=None)

option_parser.add_option("--nodes",
                         dest="nodes_filename",
                         help="Nodes file",
                         default=DEFAULT_NODES_FILE)
option_parser.add_option("--relations",
                         dest="relations_filename",
                         help="Relationships file",
                         default=DEFAULT_RELS_FILE)
option_parser.add_option(
    "--multiple",
    action="store_true",
    dest="multiple_files",
    help="whether to output multiple nodes files and relationships files",
    default=MULTIPLE_FILES)
option_parser.add_option("--limit", type="int", dest="limit", default=None)
(options, args) = option_parser.parse_args()


# this is used to concatenate some columns as a comma-separated
# list of Labels for batch-import (branch Neo4j 2.0)
def concat_translate(*args):
    return """translate(
    array_to_string(
        %s,
        ','),
    ' _"/-', '')""" % ("ARRAY[\n            %s]" % ",\n            ".join(
        ["initcap(%s)" % col for col in args]))


class MusicBrainzExporter(SQL2GraphExporter):
    nodes_header_override = {
        "mbid":
        '"mbid:string:mb_exact"',
        "pk":
        '"pk:long:mb_exact"',
        "name":
        '"name:string:mb_fulltext"',
        "kind":
        None,
        "type":
        None,
        "format":
        None,
        "gender":
        None,
        "status":
        None,
        "packaging":
        None,
        (
            "kind",
            "type",
            "format",
            "gender",
            "status",
            "packaging",
        ): (
            concat_translate,
            '"l:label"',
        ),
        #"latitude": '"latitude:float"',
        #"longitude": '"longitude:float"',
        "latitude":
        None,
        "longitude":
        None,
        "code":
        None,
        "pk":
        None,
        "locale":
        None,
    }


if all([options.included_entities, options.excluded_entities]):
    raise RuntimeError(
        "You must choose between including or exluding entities")
if options.excluded_entities is not None:
    options.excluded_entities = options.excluded_entities.split(',')
    entities = [e for e in mbentities if e not in options.excluded_entities]

elif options.included_entities is not None:
    options.included_entities = options.included_entities.split(',')
    entities = options.included_entities
    entities += make_link_entity_list(entities)
else:
    entities = mbentities
exporter = MusicBrainzExporter(mbschema, entities, strict=False)

exporter.set_nodes_filename(options.nodes_filename)
exporter.set_rels_filename(options.relations_filename)

if options.limit:
    exporter.set_entity_export_limit(options.limit)

print(r"""
-- Change TABs to spaces in "name" column for "track" and "work" tables
-- somehow these TABs can make batch-import CSV parsing choked
UPDATE track SET name=translate(name, E'\t', ' ') WHERE name LIKE E'%\t%';
UPDATE work SET name=translate(name, E'\t', ' ') WHERE name LIKE E'%\t%';
""")
print(exporter.create_mapping_table_query(multiple=options.multiple_files))
print(exporter.create_nodes_query(multiple=options.multiple_files))
print(exporter.create_relationships_query(multiple=options.multiple_files))
