#!/usr/bin/env python

from __future__ import print_function

import argparse
from collections import OrderedDict
import errno
from io import StringIO, BytesIO
import json
import os
from sqlalchemy import *
from sqlalchemy import types
from sqlalchemy.exc import ArgumentError, OperationalError, InvalidRequestError, SAWarning
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import create_session, mapper
from sqlalchemy.sql import expression, functions
import sys
import warnings

from catsql import cmdline
from catsql.nullify import Nullify

if sys.version_info[0] == 2:
    import unicodecsv as csv
else:
    import csv

warnings.simplefilter("ignore", category=SAWarning)

# Get approximate length of header
class CsvRowWriter(object):
    def __init__(self):
        self.writer = None

    def writerow(self, row):
        row = [str(s) for s in row]
        try:
            queue = StringIO()
            self.writer = csv.writer(queue, lineterminator='')
            self.writer.writerow(row)
        except TypeError:
            queue = BytesIO()
            self.writer = csv.writer(queue, lineterminator='')
            self.writer.writerow(row)
        return queue.getvalue()

class Viewer(object):
    def __init__(self, args):
        self.args = args
        self.failure = False
        self.setup_filters(args)
        self.connect_database()

    def setup_filters(self, args):
        self.url = args.url
        self.tables = args.table

        self.context_columns = set()
        self.context_filters = dict()
        if args.load_bookmark:
            with open(self.url, 'r') as fin:
                nargs = json.loads(fin.read())
            self.url = nargs['url']
            self.tables = nargs['table']
            self.context_filters = nargs['context']
            self.context_columns = set(nargs['hidden_columns'])
            if args.value is None:
                self.args.value = []

        if self.tables is not None:
            self.tables = set(self.tables)
        self.row_filter = args.sql
        self.output_in_csv = args.csv

        if args.value is not None:
            for context in args.value:
                if '=' in context:
                    key, value = context.split('=', 1)
                    self.context_filters[key] = value
                    self.context_columns.add(key)
                else:
                    self.context_columns.add(context)

    def connect_database(self):
        self.Base = declarative_base()
        try:
            self.engine = create_engine(self.url, echo=self.args.verbose)
        except ImportError as e:
            print("Support library for this database not installed - {}".format(e))
            exit(1)
        except ArgumentError:
            try:
                # maybe this is a local sqlite database?
                sqlite_url = 'sqlite:///{}'.format(self.url)
                self.engine = create_engine(sqlite_url, echo=self.args.verbose)
                self.args.url = sqlite_url
            except ArgumentError:
                # no joy, recreate the original problem and die.
                self.engine = create_engine(self.url, echo=self.args.verbose)

        self.Base.metadata.reflect(self.engine)

        self.session = create_session(bind=self.engine)

    def ok_column(self, name):
        if self.args.terse:
            if name in context_columns:
                return False
        return True

    def start_table(self, table_name, columns):
        self.header_shown = False
        self.header_considered = False
        self.table_name = table_name
        self.columns = columns

    def show_header_on_need(self):
        if self.header_shown or self.header_considered:
            return
        self.header_considered = True
        if len(self.tables_so_far) > 0:
            if self.output_in_csv:
                self.failure = True
                self.tables_so_far.append(self.table_name)
                return
            print("", file=self.output_file)
        if not self.output_in_csv:
            print('== {} =='.format(self.table_name), file=self.output_file)
        header_writer = CsvRowWriter()

        header = header_writer.writerow(list(column for column in self.columns 
                                             if self.ok_column(column)))
        print(header, file=self.output_file)
        if not self.output_in_csv:
            print('-' * len(header), file=self.output_file)
        self.header_shown = True
        self.tables_so_far.append(self.table_name)


    def show(self):

        self.tables_so_far = []

        work = None
        output_filename = None
        self.output_file = sys.stdout
        work_file = None
        self.start_table(None, None)

        try:
            if self.args.edit:
                import tempfile
                work = tempfile.mkdtemp()
                output_filename = os.path.join(work, 'reference.csv')
                work_file = open(output_filename, 'wt')
                self.output_file = work_file
                self.output_in_csv = True
                self.args.safe_null = True
                self.args.save_bookmark = [ os.path.join(work, 'bookmark.json') ]

            table_items = self.Base.metadata.tables.items()
            table_items.sort(key=lambda p: p[0])

            viable_tables = []
            for table_name, table in table_items:
                if self.tables is not None:
                    if table_name not in self.tables:
                        continue
                rows = self.session.query(table)
                if self.row_filter is not None:
                    for filter in self.row_filter:
                        rows = rows.filter(text(filter))
                    try:
                        count = rows.count()
                    except OperationalError as e:
                        # should cache these and show if no results at all found
                        continue

                if self.args.value is not None:
                    try:
                        rows = rows.filter_by(**self.context_filters)
                    except InvalidRequestError as e:
                        continue

                if self.args.grep is not None:
                    # functions.concat would be neater, but doesn't seem to translate
                    # correctly on sqlite
                    parts = ''
                    for idx, column in enumerate(table.columns):
                        if idx > 0:
                            parts = parts + ' // '
                        parts = parts + functions.coalesce(expression.cast(column,
                                                                           types.Unicode),
                                                          '')
                    rows = rows.filter(parts.contains(self.args.grep[0]))

                primary_key = table.primary_key
                if len(primary_key) >= 1:
                    rows = rows.order_by(*primary_key)
                elif len(table.c) >= 1:
                    rows = rows.order_by(*table.c)

                if self.args.limit:
                    rows = rows.limit(int(self.args.limit[0]))

                self.header_shown = False
                self.start_table(table_name, table.columns.keys())
                viable_tables.append(table_name)

                if not self.args.count:
                    # csv spec is that eol is \r\n; we ignore this for our purposes
                    # for good reasons that unfortunately there isn't space to describe
                    # here on the back of this envelope
                    csv_writer = csv.writer(self.output_file, lineterminator='\n')
                    if self.args.safe_null:
                        nullify = Nullify()
                        for row in rows:
                            self.show_header_on_need()
                            csv_writer.writerow(list(nullify.encode_null(cell)
                                                     for c, cell in enumerate(row)
                                                     if self.ok_column(self.columns[c])))
                    else:
                        for row in rows:
                            self.show_header_on_need()
                            csv_writer.writerow(list(cell for c, cell in enumerate(row)
                                                     if self.ok_column(self.columns[c])))
                    del csv_writer
                else:
                    self.show_header_on_need()
                    ct = rows.count()
                    print("({} row{})".format(ct, '' if ct == 1 else 's'), file=self.output_file)

            if len(self.tables_so_far) == 0 and len(viable_tables) == 1:
                self.show_header_on_need()

            if self.args.save_bookmark:
                with open(self.args.save_bookmark[0], 'w') as fout:
                    link = OrderedDict()
                    link['url'] = self.args.url
                    link['table'] = self.args.table
                    link['context'] = self.context_filters
                    link['hidden_columns'] = sorted(self.context_columns)
                    link['sql'] = self.args.sql
                    fout.write(json.dumps(link, indent=2))

            if self.args.edit and not self.failure:
                work_file.close()
                work_file = None
                from shutil import copyfile
                edit_filename = os.path.join(work, 'variant.csv')
                copyfile(output_filename, edit_filename)
                from subprocess import call
                EDITOR = os.environ.get('EDITOR', 'nano')
                call([EDITOR, edit_filename])
                call(['patchsql', self.args.url] +
                     ['--table'] + self.tables_so_far + 
                     ['--follow', output_filename, edit_filename] + 
                     ['--safe-null'])

        finally:
            if self.failure:
                print("ERROR: "
                      "More than one table in CSV output, consider adding:",
                      file=sys.stderr)
                for name in self.tables_so_far:
                    print("  --table {}".format(name),
                          file=sys.stderr)

            if work:
                if work_file:
                    try:
                        work_file.close()
                    except:
                        pass
                    work_file = None
                import shutil
                shutil.rmtree(work)
                work = None


def catsql():

    parser = argparse.ArgumentParser(description='Quickly display and edit a slice of a database.')

    cmdline.add_options(parser)
    args = parser.parse_args()

    viewer = Viewer(args)
    viewer.show()


def main():
    try:
        catsql()
    except IOError as error:
        if error.errno == errno.EPIPE:
            # totally benign e.g. pipe through head/tail
            pass
        else:
            raise

if __name__ == "__main__":
    main()
