#!/usr/bin/python3

import sys, os

try:
    import dbl
except ImportError:
    sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'lib', 'wstools'))
    import dbl

from palaso.sldr.ldml_exemplars import Exemplars
import argparse
import multiprocessing, logging

def process(infname, lang, outdir=".", update=False):
    path = os.path.join(outdir, lang[0])
    if not os.path.exists(path):
        os.makedirs(path)
    outfile = os.path.join(path, lang.replace("-","_")+".tsv")
    if update and os.path.exists(outfile):
        logging.debug("Skipping: {}".format(infname))
        return False
    logging.debug("Processing: {}".format(infname))
    dblObj = dbl.DBL(infname)
    exemplars = Exemplars()
    exemplars.frequent = 0.0
    try:
        for t in dblObj.analyze_text():
            exemplars.process(t)
    except Exception as e:
        pass
    dblObj.close_project()
    exemplars.analyze()
    clusters = {k.base+k.trailers: v for k, v in exemplars.raw_clusters.items()}
    
    with open(outfile, "w") as outf:
        outf.write("cluster\tcount\n")
        for k, v in sorted(clusters.items(), key=lambda a: (-a[1], a[0])):
            outf.write("{}\t{}\n".format(k, v))
    return True

parser = argparse.ArgumentParser()
parser.add_argument('inputdir',help='Input directory containing project zips')
parser.add_argument('-o','--outdir',help='Output directory for generated files')
parser.add_argument('-l','--loglevel',help='Set logging level')
parser.add_argument('-L','--lang',help='Langtag to process')
parser.add_argument('-d','--dbl',action='store_true',help='Update DBL into inputdir')
parser.add_argument('-j','--jobs',type=int,help='number of parallel jobs, default 0 = num processors')
parser.add_argument('--listlangs',action='store_true',help='list all language codes available')
parser.add_argument('-u','--update',action='store_true',help='only pull new dbl files')
args = parser.parse_args()

if args.loglevel:
    logging.basicConfig(stream=sys.stdout, level=args.loglevel.upper(),
            format="%(levelname)s:%(module)s %(message)s")

if args.dbl:
    dreader = dbl.DBLReader()
    dreader.download(args.inputdir, lang=args.lang, update=args.update)

def doit(a):
    process(a[0], a[1], outdir=args.outdir, update=args.update)
    return True

if args.listlangs:
    dreader = dbl.DBLReader()
    entries = dreader.getEntries()
    for (entryId, entryInfo) in entries.items():
        (entryLangCode, entryAccessType) = entryInfo
        logging.debug("{} -> {}".format(entryLangCode, entryId))
else:
    filelist = [os.path.join(args.inputdir, f) for f in os.listdir(args.inputdir) if f.endswith(".zip")]
    allfiles = list(dbl.process_projects(filelist, args.lang))
    if args.jobs == 100:
        for a in allfiles:
            doit(a)
    else:
        p = multiprocessing.Pool(processes=args.jobs)
        list(p.map_async(doit, allfiles).get())
        
