#!/usr/bin/env python

from __future__ import print_function
from __future__ import unicode_literals

import collections
import csv
import logging
import logging.config
import re
import sys, os

import docopt
import psycopg2
import psycopg2.extras
import six

options = docopt.docopt("""
usage: codex-list-concepts [options]
                           [--ancestors <concept_id>]...
                           [--descendants <concept_id>]...
                           [--matching <regex>]...

options:
  -s <string>, --schema <string>    Database schema [default: public]
  -d <string>, --database <string>  Database name or URI [default: postgres]
  --ancestors <concept_id>          List all ancestors of this concept
  --descendants <concept_id>        List all descendants of this concept
  --matching <regex>                List all concepts whose name match a
                                    regular expression <regex>
  --as-table                        Print identifiers as a table (default)
  --as-csv                          Print identifiers as a comma-delimited table
  --as-tsv                          Print identifiers as a tab-delimited table
  --as-list                         Print identifiers as a comma-delimited list
  --no-header                       Exclude header from the output
  -v, --verbose                     Display additional debugging information
  -h, --help                        Display this help then exit
""")

logging.config.dictConfig({
    "version": 1,
    "disable_existing_loggers": True,
    "formatters": {"default": {
        "format": "[%(asctime)s] %(levelname)s: %(message)s"
        }},
    "handlers": {"default": {
        "class": "logging.StreamHandler",
        "formatter": "default",
        }},
    "loggers": {"": {
        "handlers": ["default"],
        "level": logging.DEBUG if (options["--verbose"]) else logging.INFO,
        "propagate": True
        }}
    })

logger = logging.getLogger(os.path.basename(__file__))

def error (msg, is_exception = False):
    if (is_exception) and (options["--verbose"]):
        logger.exception(msg)
    else:
        logger.error(msg)
    sys.exit(1)

output_format_options = \
    options["--as-table"] + \
    options["--as-csv"] + \
    options["--as-tsv"] + \
    options["--as-list"]

if (output_format_options == 0):
    options["--as-table"] = True
elif (output_format_options > 1):
    error("only one of --as-* options is allowed")

ancestors_of = {}
for entry in options["--ancestors"]:
    for concept_id in entry.split(','):
        try:
            concept_id = int(concept_id)
        except:
            error("invalid concept identifier: %s" % concept_id)
        ancestors_of[concept_id] = True

descendants_of = {}
for entry in options["--descendants"]:
    for concept_id in entry.split(','):
        try:
            concept_id = int(concept_id)
        except:
            error("invalid concept identifier: %s" % concept_id)
        descendants_of[concept_id] = True

matching_regex = []
for entry in options["--matching"]:
    try:
        re.compile(entry, re.IGNORECASE)
    except Exception as e:
        error("invalid regular expression: %s (error was: %s)" % (entry, e))
    matching_regex.append(entry)

if (len(ancestors_of) == 0) and \
   (len(descendants_of) == 0) and \
   (len(matching_regex) == 0):
    error("--ancestors, --descendants or --matching is required")

include_header = (not options["--no-header"])

#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

postgres_uri = options["--database"]
postgres_schema = options["--schema"]

try:
    # if the value provided is a URI, pass it as such
    if (postgres_uri.startswith("postgresql://")):
        connection = psycopg2.connect(postgres_uri)

    # if not, consider it to be a database name
    else:
        connection = psycopg2.connect(database = postgres_uri)

    with connection.cursor() as cursor:
        cursor.execute("SET search_path TO %s", (postgres_schema,))

except psycopg2.Error as e:
    error("error from database server:\n%s" % e, True)

#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

concept_getters, concepts = [], collections.OrderedDict()

ancestor_concepts = lambda concept_id: (
    "SELECT * FROM codex_concept_ancestors(%s)", concept_id)

for concept_id in ancestors_of:
    concept_getters.append(ancestor_concepts(concept_id))

descendant_concepts  = lambda concept_id: (
    "SELECT * FROM codex_concept_descendants(%s)", concept_id)

for concept_id in descendants_of:
    concept_getters.append(descendant_concepts(concept_id))

matching_concepts = lambda pattern: (
    "SELECT * FROM concept WHERE (concept_name ~ %s)", pattern)

for regex in matching_regex:
    concept_getters.append(matching_concepts(regex))

header = None
for query in concept_getters:
    cursor = connection.cursor(
        cursor_factory = psycopg2.extras.RealDictCursor)

    cursor.execute(query[0], query[1:])
    for entry in cursor:
        concepts[entry["concept_id"]] = entry

    if (header is None):
        header = [column.name for column in cursor.description]

if (options["--as-list"]):
    print(','.join(map(str, six.iterkeys(concepts))))

elif (options["--as-csv"] or options["--as-tsv"]):
    o = csv.DictWriter(
        sys.stdout, fieldnames = header,
        dialect = "excel" if (options["--as-csv"]) else "excel-tab")

    if (include_header):
        o.writeheader()

    for entry in six.itervalues(concepts):
        o.writerow(entry)

else:
    width = {column: 0 for column in header}
    for entry in six.itervalues(concepts):
        width = {column: max(width[column], len(str(entry[column]))) \
            for column in header}

    if (include_header):
        width = {column: max(width[column], len(column)) for column in header}
        print(' '.join([column.ljust(width[column]) for column in header]))

    for entry in six.itervalues(concepts):
        print(' '.join([str(entry[column]).ljust(width[column]) \
            for column in header]))
