import pymongo
import json
import sys

from utility_prediction import data_dict_load

"""
A simple program for loading JSON dicts into a MongoDB collection
"""

#TODO: Add support for GZIP json

def main(json_json_files_to_process, runtime_json, collection_name_override=None):
    with open(runtime_json, "r") as fj:
        runtime_config = json.load(fj)

    mongo_db_config = runtime_config["mongo_db_config"]
    connection_string = mongo_db_config["connection_string"]
    database_name = mongo_db_config["database_name"]

    if collection_name_override is None:
        collection_name = mongo_db_config["collection_name"]
    else:
        collection_name = collection_name_override

    refresh_collection = mongo_db_config["refresh_collection"]
    print("Connecting to MongoDB '%s'" % connection_string)
    client = pymongo.MongoClient(connection_string)

    database = client[database_name]
    collection = database[collection_name]

    if refresh_collection:
        collection.delete_many({})

    initial_collection_count = collection.count()

    with open(json_json_files_to_process, "r") as fj:
        json_files_to_to_process = json.load(fj)

        j = 0
        for json_file_dict in json_files_to_to_process:
            json_file_name = json_file_dict["data_json_file"]
            print("Loading '%s'" % json_file_name)

            data_dict = data_dict_load(json_file_name)
            i = 0
            for datum_key in data_dict:

                collection.insert(data_dict[datum_key], check_keys=False) # Have to makes sure keys do not have a "." or "$" in it

                if i > 0 and i % 500 == 0:
                    print("Inserted '%s' documents" % i)
                i += 1
            j += i

        print("Inserted '%s' total documents" % j)

    print("Added %s items in collection '%s' housed in database '%s'" % (collection.count() - initial_collection_count, collection_name, database_name))

if __name__ == "__main__":
    if len(sys.argv) == 3:
        main(sys.argv[1], sys.argv[2])
    elif len(sys.argv) == 4:
        main(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print("""Usage: python load_documents_from_doc_db.py batch_files.json runtime.json [collection_name]
  where batch_files.json is a json file generated by "build_document_mapping_from_db.py"
        """)


