import argparse
import os
import sys
import pickle
import json
import shutil
import sqlparse

output_vocab = ['_UNK', '_EOS', '=', 'select', 'value', ')', '(', 'where', ',', 'count', 'group_by', 'order_by', 'distinct', 'and', 'limit_value', 'limit', 'desc', '>', 'avg', 'having', 'max', 'in', '<', 'sum', 'intersect', 'not', 'min', 'except', 'or', 'asc', 'like', '!=', 'union', 'between', '-', '+', '/']


def get_candidate_tables(format_sql, schema):
    candidate_tables = []

    tokens = format_sql.split()
    for i, token in enumerate(tokens):
        if '.' in token:
            table_name = token.split('.')[0]
            candidate_tables.append(table_name)

    candidate_tables = list(set(candidate_tables))
    candidate_tables.sort()
    table_names_original = [table_name.lower() for table_name in schema['table_names_original']]
    candidate_tables_id = [table_names_original.index(table_name) for table_name in candidate_tables]

    assert -1 not in candidate_tables_id
    table_names_original = schema['table_names_original']

    return candidate_tables_id, table_names_original
def read_database_schema(database_schema, schema_tokens, column_names, database_schemas_dict):
    def get_schema_tokens(table_schema):
        column_names_surface_form = []
        column_names = []
        column_names_original = table_schema['column_names_original']
        table_names = table_schema['table_names']
        table_names_original = table_schema['table_names_original']
        for i, (table_id, column_name) in enumerate(column_names_original):
            if table_id >= 0:
                table_name = table_names_original[table_id]
                column_name_surface_form = '{}.{}'.format(table_name, column_name)
            else:
                # this is just *
                column_name_surface_form = column_name
            column_names_surface_form.append(column_name_surface_form.lower())
            column_names.append(column_name.lower())

        # also add table_name.*
        for table_name in table_names_original:
            column_names_surface_form.append('{}.*'.format(table_name.lower()))

        return column_names_surface_form, column_names

    database_id = database_schema['db_id']
    database_schemas_dict[database_id] = database_schema
    schema_tokens[database_id], column_names[database_id] = get_schema_tokens(database_schema)

    return schema_tokens, column_names, database_schemas_dict


def remove_from_with_join(format_sql_2):
    used_tables_list = []
    format_sql_3 = []
    table_to_name = {}
    table_list = []
    old_table_to_name = {}
    old_table_list = []
    for sub_sql in format_sql_2.split('\n'):
        if 'select ' in sub_sql:
            # only replace alias: t1 -> table_name, t2 -> table_name, etc...
            if len(table_list) > 0:
                for i in range(len(format_sql_3)):
                    for table, name in table_to_name.items():
                        format_sql_3[i] = format_sql_3[i].replace(table, name)

            old_table_list = table_list
            old_table_to_name = table_to_name
            table_to_name = {}
            table_list = []
            format_sql_3.append(sub_sql)
        elif sub_sql.startswith('from'):
            new_sub_sql = None
            sub_sql_tokens = sub_sql.split()
            for t_i, t in enumerate(sub_sql_tokens):
                if t == 'as':
                    table_to_name[sub_sql_tokens[t_i + 1]] = sub_sql_tokens[t_i - 1]
                    table_list.append(sub_sql_tokens[t_i - 1])
                elif t == ')' and new_sub_sql is None:
                    # new_sub_sql keeps some trailing parts after ')'
                    new_sub_sql = ' '.join(sub_sql_tokens[t_i:])
            if len(table_list) > 0:
                # if it's a from clause with join
                if new_sub_sql is not None:
                    format_sql_3.append(new_sub_sql)

                used_tables_list.append(table_list)
            else:
                # if it's a from clause without join
                table_list = old_table_list
                table_to_name = old_table_to_name
                assert 'join' not in sub_sql
                if new_sub_sql is not None:
                    sub_sub_sql = sub_sql[:-len(new_sub_sql)].strip()
                    assert len(sub_sub_sql.split()) == 2
                    used_tables_list.append([sub_sub_sql.split()[1]])
                    format_sql_3.append(sub_sub_sql)
                    format_sql_3.append(new_sub_sql)
                elif 'join' not in sub_sql:
                    assert len(sub_sql.split()) == 2 or len(sub_sql.split()) == 1
                    if len(sub_sql.split()) == 2:
                        used_tables_list.append([sub_sql.split()[1]])

                    format_sql_3.append(sub_sql)
                else:
                    print('bad from clause in remove_from_with_join')
                    exit()
        else:
            format_sql_3.append(sub_sql)

    if len(table_list) > 0:
        for i in range(len(format_sql_3)):
            for table, name in table_to_name.items():
                format_sql_3[i] = format_sql_3[i].replace(table, name)

    used_tables = []
    for t in used_tables_list:
        for tt in t:
            used_tables.append(tt)
    used_tables = list(set(used_tables))

    return format_sql_3, used_tables, used_tables_list


def remove_from_without_join(format_sql_3, column_names, schema_tokens):
    format_sql_4 = []
    table_name = None
    for sub_sql in format_sql_3.split('\n'):
        if 'select ' in sub_sql:
            if table_name:
                for i in range(len(format_sql_4)):
                    tokens = format_sql_4[i].split()
                    for ii, token in enumerate(tokens):
                        if token in column_names and tokens[ii - 1] != '.':
                            if (ii + 1 < len(tokens) and tokens[ii + 1] != '.' and tokens[
                                ii + 1] != '(') or ii + 1 == len(tokens):
                                if '{}.{}'.format(table_name, token) in schema_tokens:
                                    tokens[ii] = '{} . {}'.format(table_name, token)
                    format_sql_4[i] = ' '.join(tokens)

            format_sql_4.append(sub_sql)
        elif sub_sql.startswith('from'):
            sub_sql_tokens = sub_sql.split()
            if len(sub_sql_tokens) == 1:
                table_name = None
            elif len(sub_sql_tokens) == 2:
                table_name = sub_sql_tokens[1]
            else:
                print('bad from clause in remove_from_without_join')
                print(format_sql_3)
                exit()
        else:
            format_sql_4.append(sub_sql)

    if table_name:
        for i in range(len(format_sql_4)):
            tokens = format_sql_4[i].split()
            for ii, token in enumerate(tokens):
                if token in column_names and tokens[ii - 1] != '.':
                    if (ii + 1 < len(tokens) and tokens[ii + 1] != '.' and tokens[ii + 1] != '(') or ii + 1 == len(
                            tokens):
                        if '{}.{}'.format(table_name, token) in schema_tokens:
                            tokens[ii] = '{} . {}'.format(table_name, token)
            format_sql_4[i] = ' '.join(tokens)

    return format_sql_4


def add_table_name(format_sql_3, used_tables, column_names, schema_tokens):
    # If just one table used, easy case, replace all column_name -> table_name.column_name
    if len(used_tables) == 1:
        table_name = used_tables[0]
        format_sql_4 = []
        for sub_sql in format_sql_3.split('\n'):
            if sub_sql.startswith('from'):
                format_sql_4.append(sub_sql)
                continue

            tokens = sub_sql.split()
            for ii, token in enumerate(tokens):
                if token in column_names and tokens[ii - 1] != '.':
                    if (ii + 1 < len(tokens) and tokens[ii + 1] != '.' and tokens[ii + 1] != '(') or ii + 1 == len(
                            tokens):
                        if '{}.{}'.format(table_name, token) in schema_tokens:
                            tokens[ii] = '{} . {}'.format(table_name, token)
            format_sql_4.append(' '.join(tokens))
        return format_sql_4

    def get_table_name_for(token):
        table_names = []
        for table_name in used_tables:
            if '{}.{}'.format(table_name, token) in schema_tokens:
                table_names.append(table_name)
        if len(table_names) == 0:
            return 'table'
        if len(table_names) > 1:
            return None
        else:
            return table_names[0]

    format_sql_4 = []
    for sub_sql in format_sql_3.split('\n'):
        if sub_sql.startswith('from'):
            format_sql_4.append(sub_sql)
            continue

        tokens = sub_sql.split()
        for ii, token in enumerate(tokens):
            # skip *
            if token == '*':
                continue
            if token in column_names and tokens[ii - 1] != '.':
                if (ii + 1 < len(tokens) and tokens[ii + 1] != '.' and tokens[ii + 1] != '(') or ii + 1 == len(tokens):
                    table_name = get_table_name_for(token)
                    if table_name:
                        tokens[ii] = '{} . {}'.format(table_name, token)
        format_sql_4.append(' '.join(tokens))

    return format_sql_4


def check_oov(format_sql_final, output_vocab, schema_tokens):
    for sql_tok in format_sql_final.split():
        if not (sql_tok in schema_tokens or sql_tok in output_vocab):
            print('OOV!', sql_tok)
            raise Exception('OOV')


def normalize_space(format_sql):
    format_sql_1 = [' '.join(
        sub_sql.strip().replace(',', ' , ').replace('.', ' . ').replace('(', ' ( ').replace(')', ' ) ').split()) for
                    sub_sql in format_sql.split('\n')]
    format_sql_1 = '\n'.join(format_sql_1)

    format_sql_2 = format_sql_1.replace('\njoin', ' join').replace(',\n', ', ').replace(' where', '\nwhere').replace(
        ' intersect', '\nintersect').replace('\nand', ' and').replace('order by t2 .\nstart desc',
                                                                      'order by t2 . start desc')

    format_sql_2 = format_sql_2.replace('select\noperator', 'select operator').replace('select\nconstructor',
                                                                                       'select constructor').replace(
        'select\nstart', 'select start').replace('select\ndrop', 'select drop').replace('select\nwork',
                                                                                        'select work').replace(
        'select\ngroup', 'select group').replace('select\nwhere_built', 'select where_built').replace('select\norder',
                                                                                                      'select order').replace(
        'from\noperator', 'from operator').replace('from\nforward', 'from forward').replace('from\nfor',
                                                                                            'from for').replace(
        'from\ndrop', 'from drop').replace('from\norder', 'from order').replace('.\nstart', '. start').replace(
        '.\norder', '. order').replace('.\noperator', '. operator').replace('.\nsets', '. sets').replace(
        '.\nwhere_built', '. where_built').replace('.\nwork', '. work').replace('.\nconstructor',
                                                                                '. constructor').replace('.\ngroup',
                                                                                                         '. group').replace(
        '.\nfor', '. for').replace('.\ndrop', '. drop').replace('.\nwhere', '. where')

    format_sql_2 = format_sql_2.replace('group by', 'group_by').replace('order by', 'order_by').replace('! =',
                                                                                                        '!=').replace(
        'limit value', 'limit_value')
    return format_sql_2


def normalize_final_sql(format_sql_5):
    format_sql_final = format_sql_5.replace('\n', ' ').replace(' . ', '.').replace('group by', 'group_by').replace(
        'order by', 'order_by').replace('! =', '!=').replace('limit value', 'limit_value')

    # normalize two bad sqls
    if 't1' in format_sql_final or 't2' in format_sql_final or 't3' in format_sql_final or 't4' in format_sql_final:
        format_sql_final = format_sql_final.replace('t2.dormid', 'dorm.dormid')

    # This is the failure case of remove_from_without_join()
    format_sql_final = format_sql_final.replace(
        'select city.city_name where city.state_name in ( select state.state_name where state.state_name in ( select river.traverse where river.river_name = value ) and state.area = ( select min ( state.area ) where state.state_name in ( select river.traverse where river.river_name = value ) ) ) order_by population desc limit_value',
        'select city.city_name where city.state_name in ( select state.state_name where state.state_name in ( select river.traverse where river.river_name = value ) and state.area = ( select min ( state.area ) where state.state_name in ( select river.traverse where river.river_name = value ) ) ) order_by city.population desc limit_value')

    return format_sql_final


def parse_sql(sql_string, column_names, output_vocab, schema_tokens, schema):
    format_sql = sqlparse.format(sql_string, reindent=True)
    format_sql_2 = normalize_space(format_sql)

    num_from = sum([1 for sub_sql in format_sql_2.split('\n') if sub_sql.startswith('from')])
    num_select = format_sql_2.count('select ') + format_sql_2.count('select\n')

    format_sql_3, used_tables, used_tables_list = remove_from_with_join(format_sql_2)

    format_sql_3 = '\n'.join(format_sql_3)
    format_sql_4 = add_table_name(format_sql_3, used_tables, column_names, schema_tokens)

    format_sql_4 = '\n'.join(format_sql_4)
    format_sql_5 = remove_from_without_join(format_sql_4, column_names, schema_tokens)

    format_sql_5 = '\n'.join(format_sql_5)
    format_sql_final = normalize_final_sql(format_sql_5)

    candidate_tables_id, table_names_original = get_candidate_tables(format_sql_final, schema)

    failure = False
    if len(candidate_tables_id) != len(used_tables):
        failure = True

    check_oov(format_sql_final, output_vocab, schema_tokens)

    return format_sql_final