#!/usr/bin/env python2

import datetime as dt
from itertools import dropwhile, islice, takewhile
import os
import re
import shlex
from subprocess import Popen, PIPE
import sys
import textwrap

COPY_DEFAULT = 'pbcopy'

USAGE = textwrap.dedent("""
    Usage: {me} <file.sql> [--var <expr>] [--copy] [--expand-only] [--log [file]] [--explain]

    The provided SQL file should contain a top-level comment specifying
    a database -

        -- DATABASE: dbname

    Template variables can be specified as follows -

        -- VAR: foo='bar'
        -- VAR: baz=12

    If these variable names are found in the provided SQL file and are
    surrounded by word boundaries (e.g. they are not *contained* in another
    word, so, given the VARs above, `afooble` will not become `a'bar'ble`)
    they are replaced with the specified value literal.  For example -

        select a from b where c = foo and d = baz;

    Will become -

        select a from b where c = 'bar' and d = 12;

    Instead of specifying VARs from within the SQL file itself, you can
    also do so via the --var flag.

        {me} --var foo='bar' --var baz=12
    You can use the --copy argument to copy to the output to the clipboard.
    By default, this uses the command '{COPY_DEFAULT}'; however, this can be
    specified via the COPY environment variable.

    The --log flag logs the query, its VARs, and the query output to a file,
    which defaults to '<file.sql>.log'.  Logs are appended to the file starting with
    two newlines to separate log entries.

    The --explain flag will pipe the query output to the `explain` command line tool,
    which posts the query plan to http://explain.depesz.com
    Note that you will need to provide the explain syntax in your own query, this
    flag does not prepend it with an explain.
""".format(
    me=os.path.basename(__file__),
    COPY_DEFAULT=COPY_DEFAULT
)).strip()

DATABASE_COMMENT_REGEX = re.compile(r'^-- DATABASE: (\w+)$')
VAR_COMMENT_REGEX = re.compile(r'^-- VAR: (.*)$')
VAR_REGEX = re.compile(r'^(\w+)=(.*)$')
DOUBLE_NEWLINE_REGEX = re.compile('\n\n+', re.MULTILINE)


def main():
    if len(sys.argv) < 2:
        error(USAGE)

    if {'-h', '--help'}.intersection(sys.argv):
        print USAGE
        exit()

    filename = sys.argv[1]

    copy = '--copy' in sys.argv
    expand_only = '--expand-only' in sys.argv
    log_args = list(dropwhile(lambda x: x != '--log', sys.argv))[:2]
    log = len(log_args) > 0
    log_filename = filename + '.log'
    if len(log_args) > 1 and not log_args[1].startswith('-'):
        log_filename = log_args[1]

    (explain, explain_title) = get_explain_args()

    state = {
        'database': None,
        'varmap': {},
        'lines': []
    }

    for i, arg in enumerate(sys.argv):
        arg_len = len(sys.argv)
        if arg == '--var':
            if i + 1 >= arg_len:
                error('--var requires an expression')
            put_var(state, sys.argv[i + 1])

    with open(filename) as f:
        for line in f:
            process(state, line)

    if state['database'] is None:
        error("No '-- DATABASE: <dbname>' line found in " + filename)

    raw_query = ''.join(state['lines'])
    rendered_query = raw_query
    for k, v in state['varmap'].iteritems():
        rendered_query = v['regex'].sub(v['value'], rendered_query)
    rendered_query = cleanup_sql(rendered_query)

    if expand_only:
        print rendered_query
        if copy:
            copy_to_clipboard(rendered_query)
        return

    p = Popen(
        ['psql', state['database']],
        stdout=PIPE, stdin=PIPE, stderr=PIPE
    )
    (query_out, query_err) = p.communicate(input=rendered_query)
    sys.stdout.write(query_out)
    sys.stderr.write(query_err)

    if copy:
        copy_to_clipboard(query_out)

    explain_url = None
    if explain:
        (explain_out, explain_err) = pipe_to_command('explain', query_out)
        if explain_err:
            explain_url = "ERROR POSTING EXPLAIN"
        else:
            explain_url = explain_out.strip()
        print explain_url

    if log:
        with open(log_filename, 'a') as log_file:
            log_file.write(dt.datetime.now().strftime('%Y-%m-%d %H:%M:%S\n'))
            if state['varmap']:
                log_file.write('-- DEFINED VARS:\n')
                for (k, v) in state['varmap'].iteritems():
                    log_file.write('--     {0}: {1}\n'.format(k, v['value']))
            log_file.write('\n')
            log_file.write(cleanup_sql(raw_query) + '\n')
            log_file.write(query_err + '\n')
            log_file.write(query_out + '\n\n')
            if explain_url:
                log_file.write('-- EXPLAIN URL: ' + explain_url)


def error(msg):
    sys.stderr.write(msg + '\n')
    exit(1)


def cleanup_sql(query):
    def should_not_strip(line):
        return not line.startswith('--')
    return DOUBLE_NEWLINE_REGEX.sub(
        '\n\n',
        '\n'.join(filter(should_not_strip, query.splitlines()))
    ).strip()


def pipe_to_command(cmd, inp):
    p = Popen(shlex.split(cmd), stdout=PIPE, stdin=PIPE)
    return p.communicate(inp)


def copy_to_clipboard(text):
    pipe_to_command(os.environ.get('COPY', COPY_DEFAULT), text)


def process(state, line):
    m = DATABASE_COMMENT_REGEX.match(line)
    if m:
        state['database'] = m.group(1)
        return

    m = VAR_COMMENT_REGEX.match(line)
    if m:
        put_var(state, m.group(1))
        return

    state['lines'].append(line)


def put_var(state, expr):
    m = VAR_REGEX.match(expr)
    if not m:
        error('Invalid VAR expression: ' + expr)
    k = m.group(1)
    if k in state['varmap']:
        error('VAR "{0}" already defined in varmap'.format(k))
    v = m.group(2)
    # Ensure var value is quoted.
    if not v.startswith("'"):
        v = "'{0}'".format(v)
    state['varmap'][k] = {
        'regex': re.compile(r'\b' + k + r'\b'),
        'value': v
    }


def get_explain_args():
    args = list(islice(dropwhile(lambda arg: arg != '--explain', sys.argv), 2))
    if not args:
        return (False, '')
    maybe_title = args[1:]
    if maybe_title:
        title = maybe_title[0]
    else:
        title = ''
    return (True, title)


if __name__ == '__main__':
    main()
