import os
import yaml
import sys

def worldprefix(world,worldname):
    if worldname.startswith(world + "_") or worldname == world:
        return True
    else:
        return False
if len(sys.argv) == 1:
    print("Please specify world(s) to extract")
else:
    for world in sys.argv[1:]:
        print("Extracting Essentials userdata for world: "+world)

        if "userdata" in os.listdir():
            output_extraction = 'userdata_'+world
            output_remainder = 'userdata'
            try:
                os.stat(output_extraction)
            except:
                os.mkdir(output_extraction)
            try:
                os.stat(output_remainder)
            except:
                os.mkdir(output_remainder)
            for filename in os.listdir('userdata'):
                if filename.endswith('.yml'):
                    with open(os.path.join('userdata', filename), 'r') as stream:
                        try:
                            ess = yaml.load(stream)
                        except yaml.YAMLError as exc:
                            print(exc)
                            ess = yaml.load("")
                    if "lastAccountName" in ess.keys():
                        print("Processing user: "+ess["lastAccountName"])
                    else:
                        print("Processing unknown user in "+filename)
                    new_ess = dict.copy(ess)
                    if "homes" in ess.keys():
                        homes = ess["homes"]
                        new_homes = dict.copy(ess["homes"])
                        for home in dict.copy(homes):
                            if worldprefix(world,homes[home]["world"]):
                                del(homes[home])
                            else:
                                del(new_homes[home])
                        new_ess["homes"]=new_homes
                    if "lastlocation" in ess.keys():
                        if worldprefix(world,ess["lastlocation"]["world"]):
                            del(ess["lastlocation"])
                        else:
                            del(new_ess["lastlocation"])
                    if "logoutlocation" in ess.keys():
                        if worldprefix(world,ess["logoutlocation"]["world"]):
                            del(ess["logoutlocation"])
                        else:
                            del(new_ess["logoutlocation"])
                with open (os.path.join(output_extraction,filename),'w') as output:
                    output.write(yaml.dump(new_ess,default_flow_style=False))
                    output.close()
                with open (os.path.join(output_remainder,filename),'w') as output:
                    output.write(yaml.dump(ess,default_flow_style=False))
                    output.close()
        else:
            print("We need to be in the plugins/Essentials folder")