#program to convert xref metadata files into compatible files for mongo database

import pymongo, gzip, argparse, json, os
from pymongo import MongoClient
from alive_progress import alive_bar

import logging

logger = None
numFile = 0
ACCEPTED_TYPES = ["journal-article", "proceedings-article", "book-chapter"]
counters = [0] * len(ACCEPTED_TYPES)

#convert std xref metadata file to json string format
def convertFile(path):
    # dictionary for all items in the current file
    logging.info(counters)
    try:
        jsonItems = []
        logging.info("Unzipping and opening " + path)
        with gzip.open(path, "r") as gzippedFile:
            #load the json in the file as a dictionary in python
            jsonDictionary = json.load(gzippedFile)
            #for each item in the file:
            for item in jsonDictionary["items"]:
                #create a dictionary containing relevant contents of that item and add it to another dictionary to dump in future
                if "type" not in item.keys():
                    continue
                else:
                    for i in range(len(ACCEPTED_TYPES)):
                        if item["type"] == ACCEPTED_TYPES[i]:
                            counters[i] += 1
                            break
                type = item["type"]
                if type not in ACCEPTED_TYPES:
                    continue
                #dictionary to store a single item
                dictToDump = {}#{"_id":numFile,}
                if "DOI" in item.keys():
                    dictToDump.update({"DOI":item["DOI"]})
                if "title" in item.keys():
                    dictToDump.update({"title":item["title"][0]})
                if "subtitle" in item.keys():
                    dictToDump.update({"subtitle":item["subtitle"][0]})
                if "author" in item.keys():
                    dictToDump.update({"authors":item["author"]})
                dictToDump.update({"dateIndexed":item["indexed"]["date-time"]})
                dictToDump.update({"dateCreated":item["created"]["date-time"]})
                dictToDump.update({"dateDeposited":item["deposited"]["date-time"]})
                dictToDump.update({"typeOfWork":type})
                jsonItems.append(dictToDump)
                #logging.info("Processed Item " + dictToDump["DOI"])

        if len(jsonItems) == 0: return None 

        #now we have an array of things to dump, lest return it as a string
        stringToReturn = ""

        #format for mongodb
        for item in jsonItems:
            stringToReturn += json.dumps(item)
            stringToReturn += ","

        #should get format {item},{item},{item}
        stringToReturn = stringToReturn[:-1]

        logging.info("Returning processed Items from " + path)
        return stringToReturn
    except Exception as inst:
        #handle error
        logging.critical("Failed to process " + path + " with exception " + str(inst))
        return None

#process a directory of gzipped json metadata files into json strings and put them into files
def processDir(verbose, path, count, database, collection):
    #set processing function depending on where we want to pipe output
    processFunction = processDirFilesToJSON
    mongoCollection = None
    #if there is a database selected then get the mongo db information
    if database != None: 
        logging.info("Got Database!")
        processFunction = processDirFilesToMongo
        mongoClient = MongoClient("localhost", 27017)
        mongoDatabase = mongoClient[database]
        mongoCollection = mongoDatabase[collection]

    #lest get a list of the things in the directory
    listDirFiles = os.listdir(path)
    #using a progress bar library for easy ETA
    with alive_bar(len(listDirFiles)) as bar:
        for dirFile in listDirFiles:   
            logging.info("processing " + dirFile)
            processFunction(dirFile, path, mongoCollection)
            bar()

    logging.info("Successfully processed directory")

    return

#process the directory files output to json
def processDirFilesToJSON(dirFile, path, mongoCollection):
    with open("mongo" + dirFile + ".json", "w", encoding="utf-8") as outputFile:
        convertedString = convertFile(path + "/" + dirFile)
        outputFile.write(convertedString)
    
#process the directory files output to mongodb
def processDirFilesToMongo(dirFile, path, mongoCollection):
    try:
        convertedString = convertFile(path + "/" + dirFile)
        convertedString = "[" + convertedString + "]"
        jsonConvertedString = json.loads(convertedString)
        mongoCollection.insert_many(jsonConvertedString)
        logging.info("Inserted " + dirFile)
    except Exception as inst:
        logging.critical("Failed to insert " + dirFile + " with exception " + str(inst))

# main entry point and argument parser    
def main():
    parser = argparse.ArgumentParser(
        prog="xRefToMongo",
        description="convert xRef Metadata to compatible file for Mongo"
    )
    parser.add_argument("filePath", help="path to file or directory")
    parser.add_argument("-c", "--count", dest="count", help="number of file to convert if working with full metadata archive")
    parser.add_argument("-v", "--verbose", dest="verbose", default=0, action="count", help="enable verbose output")
    parser.add_argument("-db", "--database", dest="database", help="export directly to mongodb database (requires collection)")
    parser.add_argument("-cl", "--collection", dest="collection", help="collection in specified database")

    # when no input for count, equals None
    parsedArgs = parser.parse_args()

    # set up logger so that verbosity works
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger("verbosityLogger")
    logger.setLevel(parsedArgs.verbose * 10)

    # if filepath is a directory
    if "." not in parsedArgs.filePath:
        processDir(parsedArgs.verbose, parsedArgs.filePath, parsedArgs.count, parsedArgs.database, parsedArgs.collection)
    else:
        mongoString = convertFile(parsedArgs.filePath)
        logging.info(mongoString)
        mongoFile = open("mongo" + str(numFile) + ".json", "w", encoding="utf-8")
        mongoFile.write(mongoString)
        mongoFile.close()

    #print(parsedArgs.count)

if __name__ == "__main__": 
    main()