import sys
from logging import INFO, basicConfig, error, info  # ordem alfabética
from os import environ
from time import sleep  # imported but not used

import psycopg2
import simplejson as json
from tqdm import tqdm
# 2 blank lines

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

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)
# 2 blank lines

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...')


# filenames não me parece um bom nome, parece que você deveria passar uma
# lista com nomes de arquivos, e na verdade é só um arquivo só
# opção: file_options
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 = []

        # achei confuso esse for, acho que tem um jeito melhor de fazer isso
        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')
    # blank line
    # ao usar `with` pra abrir arquivos, o __exit__ dos arquivos é chamado no
    # momento que se sai do `with`, e o __exit__ gera um close(), daí não
    # precisa chamar o .close()
    result.close()
    sampleresult.close()
    info('Closing conection...')
    cur.close()
    conn.close()

# 2 blank lines
# identação desse dict passou né?
files = {
            'result':'tweets-json',  # result ou filename? :P
            'opt': 50,
        }
fetch_and_save_data_to_file(cur, files)