import sys
from logging import INFO, basicConfig, error, info
from os import environ
from time import sleep

import psycopg2
import simplejson as json
from tqdm import tqdm

basicConfig(format='%(levelname)s %(message)s', level=INFO)
try:
    table = sys.argv[1]
except:
    error('Usage: data_retrieval.py [table/view]')
    sys.exit(1)

USER = environ.get('DB_USER')
PASSWORD = environ.get('DB_PASSWORD')
if not USER or not PASSWORD:
    error('no USER or PASSWORD for connection found')
    sys.exit(1)

info('Connecting to database...')
conn = psycopg2.connect(database='twitter', user=USER, password=PASSWORD,
                        host='143.107.137.90')
cur = conn.cursor()
sql = 'SELECT json FROM {};'.format(table)
info('Executing SQL command...')
cur.execute(sql)
info('Fetching and writing data...')


def fetch_and_save_data_to_file(db_cursor, filenames):
    filenames['sampleresult'] = 'sample-' + filenames['result']
    with open(filenames['result'], 'w') as result, \
         open(filenames['sampleresult'], 'w') as sampleresult:
        info('File: %s', filenames['result'])
        i = 0
        samples = []
        for row in tqdm(db_cursor):
            if i < filenames['opt']:
                i += 1
                samples.append(row[0])
            result.write(json.dumps(row[0], result))
            result.write('\n')
        info('File: %s', filenames['sampleresult'])
        for sample in tqdm(samples):
            sampleresult.write(json.dumps(sample, sampleresult))
            sampleresult.write('\n')
    result.close()
    sampleresult.close()
    info('Closing conection...')
    cur.close()
    conn.close()

files = {
            'result':'tweets-json',
            'opt': 50,
        }
fetch_and_save_data_to_file(cur, files)