import sys
import json
import requests

try: 
    from prettytable import PrettyTable
except: 
    print('\nPlease install python pretty table (pip3 install ptable)\n')
    sys.exit(1)

search1_url = 'http://xps8920:8250/entities?attrs={}&featureMode=NONE&forceMinimal=false&withRaw=false'
dynamic_url = 'http://xps8920:8252/resolve'


# ==============================
class colors: 
    code = {}
    #--styles
    code['reset'] = '\033[0m'
    code['bold'] ='\033[01m'
    code['dim'] = '\033[02m'
    code['italics'] = '\033[03m'
    code['underline'] = '\033[04m'
    code['blink'] = '\033[05m'
    code['reverse'] = '\033[07m'
    code['strikethrough'] = '\033[09m'
    code['invisible'] = '\033[08m'
    code['fg.black'] = '\033[30m'
    code['fg.red'] = '\033[31m'
    code['fg.green'] = '\033[32m'
    code['fg.yellow'] = '\033[33m'
    code['fg.blue'] = '\033[34m'
    code['fg.magenta'] = '\033[35m'
    code['fg.cyan'] = '\033[36m'
    code['fg.lightgrey'] = '\033[37m'
    code['fg.darkgrey'] = '\033[90m'
    code['fg.lightred'] = '\033[91m'
    code['fg.lightgreen'] = '\033[92m'
    code['fg.lightyellow'] = '\033[93m'
    code['fg.lightblue'] = '\033[94m'
    code['fg.lightmagenta'] = '\033[95m'
    code['fg.lightcyan'] = '\033[96m'
    code['fg.white'] = '\033[97m'
    code['bg.black'] = '\033[40m'
    code['bg.red'] = '\033[41m'
    code['bg.green'] = '\033[42m'
    code['bg.orange'] = '\033[43m'
    code['bg.blue'] = '\033[44m'
    code['bg.magenta'] = '\033[45m'
    code['bg.cyan'] = '\033[46m'
    code['bg.lightgrey'] = '\033[47m'
    code['bg.darkgrey'] = '\033[100m'
    code['bg.lightred'] = '\033[101m'
    code['bg.lightgreen'] = '\033[102m'
    code['bg.yellow'] = '\033[103m'
    code['bg.lightblue'] = '\033[104m'
    code['bg.lightmagenta'] = '\033[105m'
    code['bg.lightcyan'] = '\033[106m'
    code['bg.white'] = '\033[107m'

def colorize(string, colorList=None):
    if colorList: 
        prefix = ''.join([colors.code[i.strip().lower()] for i in colorList.split(',')])
        suffix = colors.code['reset']
        return '{}{}{}'.format(prefix, string, suffix) 
    return string

def renderTable(columnList, recordList):
    headerList = [colorize(item, 'fg.magenta') for item in columnList]
    thisTable = PrettyTable(field_names = headerList, hrules = 1, align = 'l')
    for record in recordList:
        record[0] = colorize(record[0], 'fg.magenta')
        for i in range(len(columnList))[1:]:
            record[i] = '\n'.join(colorize(item.split(':')[0] + ':', 'fg.yellow' if i == 1 else 'fg.cyan') + item.split(':')[1] if ':' in item else item for item in record[i].split('\n'))
#                record[i] = colorize(record[i].split(':')[0], 'fg.yellow' + record[i].split(':')[1])
        thisTable.add_row(record)
    print(thisTable)

if __name__ == '__main__':

    option = '?'
    while option.upper() != 'Q':
        options = '\nOptions:\n\t1=Regular Search\n\t2=Dynamic Search\n\n\tQ=Quit\n\nSelect option: '
        option = input(options)

        if option in ('1', '2'):

            search_str = input('\nSearch for: ')
            if not search_str:
                search_str = 'Kohls'

            try: search_json = json.loads(search_str)
            except: 
                search_json = {"NAME_ORG": search_str}

            print()
            print('Initial Search Result ...')

            get_response = requests.get(search1_url.format(json.dumps(search_json)))
            if get_response.status_code == 404:
                print('404 error')
                sys.exit(1)
            elif get_response.status_code != 200:
                print(get_response.status_code, get_response.json())
                sys.exit(1)
            else:

                tableColumns = ['#', 'RECORD', 'NAME', 'ADDRESS']
                tableData = []
                recordCache = {}
                jsonRecords = ''
                entityCount = 0
                for entityData in get_response.json()['data']['searchResults']:         
                    entityCount += 1
                    record_keys = []
                    record_names = []
                    record_addresses = []
                    for recordData in entityData['records']:
                        record_key = recordData['dataSource'] + ': ' + recordData['recordId']
                        recordCache[record_key] = recordData
                        record_keys.append(record_key)
                        record_names.extend(recordData['nameData'])
                        record_addresses.extend(recordData['addressData'])
                        jsonRecords = jsonRecords + json.dumps(recordData['originalSourceData']) + '\n'
                    tableData.append([entityCount, '\n'.join(record_keys), '\n'.join(set(record_names)), '\n'.join(set(record_addresses))])
                renderTable(tableColumns, tableData)

     
        if option == '2':
            print()
            print('Dynamic Re-resolve ...')

            post_response = requests.post(dynamic_url, data=jsonRecords)
            if post_response.status_code != 200:
                print('\n--ERROR--')
                print(post_response.status_code, post_response.json())
                print()     
            else:

                tableColumns = ['#', 'RECORD', 'NAME', 'ADDRESS']
                tableData = []
                entityCount = 0
                for entityData in post_response.json():
                    entityCount += 1
                    record_keys = []
                    record_names = []
                    record_addresses = []
                    for recordPointer in entityData['RESOLVED_ENTITY']['RECORDS']:
                        record_key = recordPointer['DATA_SOURCE'] + ': ' + recordPointer['RECORD_ID']
                        recordData = recordCache[record_key]
                        record_keys.append(record_key)
                        record_names.extend(recordData['nameData'])
                        record_addresses.extend(recordData['addressData'])
                        jsonRecords = jsonRecords + json.dumps(recordData['originalSourceData']) + '\n'
                    tableData.append([entityCount, '\n'.join(record_keys), '\n'.join(set(record_names)), '\n'.join(set(record_addresses))])
                renderTable(tableColumns, tableData)

    print()
    sys.exit(0)