wmlunits

#!/usr/bin/env python3
# encoding: utf-8
"""
wmlunits -- tool to output information on all units in HTML

Run without arguments to see usage.
"""

# Makes things faster on 32-bit systems
try:
    import psyco
    psyco.full()
except ImportError:
    pass

import argparse
import os
import shutil
import sys
import traceback
import subprocess
import multiprocessing
import queue
import yaml

import wesnoth.wmlparser3 as wmlparser3
import unit_tree.helpers as helpers
import unit_tree.animations as animations
import unit_tree.html_output as html_output
import unit_tree.overview
import unit_tree.wiki_output as wiki_output

TIMEOUT = 20

def copy_images():
    print("Recolorizing pictures.")
    image_collector.copy_and_color_images(options.output)
    shutil.copy2(os.path.join(image_collector.datadir,
                              "data/tools/unit_tree/menu.js"),
                 options.output)

def shell(com):
    #print(com)
    p = subprocess.Popen(com,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE,
                         shell=True)
    out, err = p.communicate()
    #if out: sys.stdout.write(out)
    #if err: sys.stdout.write(err)

    return p.returncode

def shell_out(com):
    p = subprocess.Popen(com,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    out, err = p.communicate()
    return out

def bash(name):
    return "'" + name.replace("'", "'\\''") + "'"

def move(f, t, name):

    if os.path.exists(f + "/" + name + ".cfg"):
        com = "mv " + f + "/" + bash(name + ".cfg") + " " + t + "/"
        shell(com)

    com = "mv " + f + "/" + bash(name) + " " + t + "/"
    return shell(com)

_info = {}
def get_info(addon):
    global _info
    if addon in _info:
        return _info[addon]
    _info[addon] = None
    try:
        path = options.addons + "/" + addon + "/_info.cfg"
        if os.path.exists(path):
            parser = wmlparser3.Parser(options.wesnoth,
                                       options.config_dir,
                                       options.data_dir)
            parser.parse_file(path)
            _info[addon] = parser
        else:
            print("Cannot find " + path)
    except wmlparser3.WMLError as e:
        print(e)
    return _info[addon]

_deps = {}
global_addons = set()
def get_dependencies(addon):
    global _deps
    global global_addons
    if addon in _deps:
        return _deps[addon]
    _deps[addon] = []
    try:
        info = get_info(addon).get_all(tag="info")[0]
        row = info.get_text_val("dependencies")
        if row:
            deps1 = row.split(",")
        else:
            deps1 = []
        for d in deps1:
            if d in global_addons:
                _deps[addon].append(d)
            else:
                print("Missing dependency for " + addon + ": " + d)
    except Exception as e:
        print(e)
    return _deps[addon]

def set_dependencies(addon, depends_on):
    _deps[addon] = depends_on

def get_all_dependencies(addon):
    result = []
    check = get_dependencies(addon)[:]
    while check:
        d = check.pop()
        if d == addon or d in result:
            continue
        result.append(d)
        check += get_dependencies(d)
    return result

def sorted_by_dependencies(addons):
    sorted = []
    unsorted = addons[:]
    while unsorted:
        n = 0
        for addon in unsorted:
            for d in get_dependencies(addon):
                if d not in sorted:
                    break
            else:
                sorted.append(addon)
                unsorted.remove(addon)
                n += 1
                continue
        if n == 0:
            print("Cannot sort dependencies for these addons: " + str(unsorted))
            sorted += unsorted
            break
    return sorted

def search(batchlist, name):
    for info in batchlist:
        if info and info["name"] == name:
            return info
    batchlist.append({})
    batchlist[-1]["name"] = name
    return batchlist[-1]

def list_contents():
    class Empty:
        pass

    local = Empty()
    mainline_eras = set()
    filename = options.list

    def append(info, id, define, c=None, name=None, domain=None):
        info.append({})
        info[-1]["id"] = id
        info[-1]["define"] = define
        info[-1]["name"] = c.get_text_val("name") if c else name
        info[-1]["units"] = "?"
        info[-1]["translations"] = {}

        for isocode in languages:
            translation = html_output.Translation(options.transdir, isocode)
            def translate(string, domain):
                return translation.translate(string, domain)
            if c:
                t = c.get_text_val("name", translation=translate)
            else:
                t = translate(name, domain)
            if t != info[-1]["name"]:
                info[-1]["translations"][isocode] = t

    def get_dependency_eras(batchlist, addon):
        dependency_eras = list(mainline_eras)
        for d in get_all_dependencies(addon):
            dinfo = search(batchlist, d)
            for era in dinfo.get("eras", []):
                dependency_eras.append(era["id"])
        return dependency_eras

    def list_eras(batchlist, addon):
        eras = local.wesnoth.parser.get_all(tag="era")
        if addon != "mainline":
            dependency_eras = get_dependency_eras(batchlist, addon)
            eras = [x for x in eras if not x.get_text_val("id") in dependency_eras]
        info = []
        for era in eras:
            eid = era.get_text_val("id")
            if addon == "mainline":
                mainline_eras.add(eid)
            append(info, eid, "MULTIPLAYER", c=era)

        return info

    def list_campaigns(batchlist, addon):
        campaigns = local.wesnoth.parser.get_all(tag="campaign")
        info = []

        for campaign in campaigns:
            cid = campaign.get_text_val("id")
            d = campaign.get_text_val("define")
            d2 = campaign.get_text_val("extra_defines")
            if d2:
                d += "," + d2
            d3 = campaign.get_text_val("difficulties")
            if d3:
                d += "," + d3.split(",")[0]
            else:
                difficulty = campaign.get_all(tag="difficulty")
                if difficulty:
                    d3 = difficulty[0].get_text_val("define")
                    if d3:
                        d += "," + d3

            append(info, cid, d, c=campaign)

        return info

    def parse(wml, defines):
        def f(options, wml, defines, q):
            local.wesnoth = helpers.WesnothList(
                options.wesnoth,
                options.config_dir,
                options.data_dir,
                options.transdir)
            #print("remote", local.wesnoth)
            try:
                local.wesnoth.parser.parse_text(wml, defines)
                q.put(("ok", local.wesnoth))
            except Exception as e:
                q.put(("e", e))

        q = multiprocessing.Queue()
        p = multiprocessing.Process(target=f, args=(options, wml, defines, q))
        p.start()
        try:
            s, local.wesnoth = q.get(timeout=TIMEOUT)
        except queue.Empty:
            p.terminate()
            raise
        #print("local", s, local.wesnoth)
        p.join()
        if s == "e":
            remote_exception = local.wesnoth
            raise remote_exception

    def get_version(addon):
        parser = get_info(addon)
        if parser:
            for info in parser.get_all(tag="info"):
                return info.get_text_val("version") + "*" + info.get_text_val("uploads")

    try:
        os.makedirs(options.output + "/mainline")
    except OSError:
        pass

    try:
        batchlist = yaml.load(open(options.list))
    except IOError:
        batchlist = []

    print("mainline")

    info = search(batchlist, "mainline")
    info["version"] = "mainline"
    info["parsed"] = "false"

    parse("{core}{multiplayer/eras.cfg}", "SKIP_CORE")
    info["eras"] = list_eras(batchlist, "mainline")

    # Fake mainline campaign to have an overview of the mainline units
    info["campaigns"] = []
    append(info["campaigns"], "mainline", "", name="Units", domain="wesnoth-help")

    if not options.addons_only:
        parse("{core}{campaigns}", "SKIP_CORE")
        info["campaigns"] += list_campaigns(batchlist, "mainline")

    addons = []
    if options.addons:
        addons = os.listdir(options.addons)
    global global_addons
    global_addons = set(addons)

    # fill in the map for all dependencies
    for addon in addons:
        get_dependencies(addon)

    # this ensures that info about eras in dependant addons is available
    # already
    addons = sorted_by_dependencies(addons)

    for i, addon in enumerate(addons):
        if not os.path.isdir(options.addons + "/" + addon):
            continue
        mem = "? KB"
        try:
            mem = dict([x.split("\t", 1) for x in open("/proc/self/status").read().split("\n") if x])["VmRSS:"]
        except:
            pass
        sys.stdout.write("%4d/%4d " % (1 + i, len(addons)) + addon + " [" + mem + "] ... ")
        sys.stdout.flush()
        d = options.output + "/" + addon
        logname = d + "/error.log"
        try:
            os.makedirs(d)
        except OSError:
            pass
        version = get_version(addon)
        move(options.addons, options.config_dir + "/data/add-ons", addon)
        for d in get_dependencies(addon):
            move(options.addons, options.config_dir + "/data/add-ons", d)
        try:
            info = search(batchlist, addon)
            if info.get("version", "") == version and info.get("parsed", False) is True:
                sys.stdout.write("up to date\n")
                continue
            info["parsed"] = False
            info["dependencies"] = get_dependencies(addon)
            parse("{core}{multiplayer}{multiplayer/eras.cfg}{~add-ons}", "MULTIPLAYER,SKIP_CORE")
            info["eras"] = list_eras(batchlist, addon)
            info["campaigns"] = list_campaigns(batchlist, addon)
            info["version"] = version
            sys.stdout.write("ok\n")
        except wmlparser3.WMLError as e:
            ef = open(logname, "w")
            ef.write("<PARSE ERROR>\n")
            ef.write(str(e))
            ef.write("</PARSE ERROR>\n")
            ef.close()
            sys.stdout.write("failed\n")
        except queue.Empty as e:
            ef = open(logname, "w")
            ef.write("<TIMEOUT ERROR>\n")
            ef.write("Failed to parse the WML within " + str(TIMEOUT) + " seconds.")
            ef.write("</TIMEOUT ERROR>\n")
            ef.close()
            sys.stdout.write("failed\n")
        except Exception as e:
            ef = open(logname, "w")
            ef.write("<INTERNAL ERROR>\n")
            ef.write(repr(e))
            ef.write("</INTERNAL ERROR>\n")
            ef.close()
            sys.stdout.write("failed\n")
        finally:
            move(options.config_dir + "/data/add-ons", options.addons, addon)
            for d in get_dependencies(addon):
                move(options.config_dir + "/data/add-ons", options.addons, d)

    yaml.safe_dump(batchlist, open(filename, "w"),
                   encoding="utf-8", default_flow_style=False)

def process_campaign_or_era(addon, cid, define, batchlist):
    n = 0

    print(str(addon) + ": " + str(cid) + " " + str(define))

    if not define:
        define = ""

    wesnoth = helpers.WesnothList(
        options.wesnoth,
        options.config_dir,
        options.data_dir,
        options.transdir)
    wesnoth.batchlist = batchlist
    wesnoth.cid = cid

    wesnoth.parser.parse_text("{core/units.cfg}", "NORMAL")
    wesnoth.add_units("mainline")

    if define == "MULTIPLAYER":
        wesnoth.parser.parse_text("{core}{multiplayer}{multiplayer/eras.cfg}{~add-ons}", "MULTIPLAYER,SKIP_CORE")
        wesnoth.add_units(cid)
    else:
        if addon == "mainline":
            if cid != "mainline":
                wesnoth.parser.parse_text("{core}{campaigns}", "SKIP_CORE,NORMAL," + define)
                wesnoth.add_units(cid)
        else:
            wesnoth.parser.parse_text("{core}{~add-ons}", "SKIP_CORE," + define)
            wesnoth.add_units(cid)

    if addon == "mainline" and cid == "mainline":
        write_animation_statistics(wesnoth)

    wesnoth.add_binary_paths(addon, image_collector)

    if define == "MULTIPLAYER":
        eras = wesnoth.parser.get_all(tag="era")
        for era in eras:
            wesnoth.add_era(era)
        wesnoth.find_unit_factions()
    else:
        campaigns = wesnoth.parser.get_all(tag="campaign")
        for campaign in campaigns:
            wesnoth.add_campaign(campaign)

    wesnoth.add_languages(languages)
    wesnoth.add_terrains()
    wesnoth.check_units()

    for isocode in languages:
        if addon != "mainline" and isocode != "en_US":
            continue
        if define == "MULTIPLAYER":
            for era in list(wesnoth.era_lookup.values()):
                if era.get_text_val("id") == cid:
                    n = html_output.generate_era_report(addon, isocode, era, wesnoth)
                    break
        else:
            if cid == "mainline":
                n = html_output.generate_campaign_report(addon, isocode, None, wesnoth)

            for campaign in list(wesnoth.campaign_lookup.values()):
                if campaign.get_text_val("id") == cid:
                    n = html_output.generate_campaign_report(addon, isocode, campaign, wesnoth)
                    break
        html_output.generate_single_unit_reports(addon, isocode, wesnoth)

    return n

def batch_process():
    batchlist = yaml.load(open(options.batch))

    for addon in batchlist:
        name = addon["name"]
        set_dependencies(name, addon.get("dependencies", []))

    for addon in batchlist:
        name = addon["name"]

        if not options.reparse and addon.get("parsed", False) == True:
            continue

        if name == "mainline":
            worked = True
        else:
            worked = (move(options.addons, options.config_dir + "/data/add-ons", name) == 0)
            for d in addon.get("dependencies", []):
                move(options.addons, options.config_dir + "/data/add-ons", d)
        d = options.output + "/" + name
        try:
            os.makedirs(d)
        except OSError:
            pass
        logname = d + "/error.log"

        def err(mess):
            ef = open(logname, "a")
            ef.write(str(mess))
            ef.close()

        html_output.write_error = err

        try:
            if not worked:
                print(name + " not found")
                continue

            for era in addon.get("eras", []):
                eid = era["id"]
                n = process_campaign_or_era(name, eid, era["define"], batchlist)
                era["units"] = n

            for campaign in addon.get("campaigns", []):
                cid = campaign["id"]
                if cid is None:
                    cid = campaign["define"]
                if cid is None:
                    cid = name
                n = process_campaign_or_era(name, cid, campaign["define"], batchlist)
                campaign["units"] = n

        except wmlparser3.WMLError as e:
            ef = open(logname, "a")
            ef.write("<WML ERROR>\n")
            ef.write(str(e))
            ef.write("</WML ERROR>\n")
            ef.close()
            print("    " + name + " failed")
        except Exception as e:
            traceback.print_exc()
            print("    " + name + " failed")
            ef = open(logname, "a")
            ef.write("<INTERNAL ERROR>\n")
            ef.write("please report as bug\n")
            ef.write(str(e))
            ef.write("</INTERNAL ERROR>\n")
            ef.close()
        finally:
            if name != "mainline":
                move(options.config_dir + "/data/add-ons", options.addons, name)
                for d in addon.get("dependencies", []):
                    move(options.config_dir + "/data/add-ons", options.addons, d)

        addon["parsed"] = True

        yaml.safe_dump(batchlist, open(options.batch, "w"),
                       encoding="utf-8", default_flow_style=False)

        try:
            unit_tree.overview.write_addon_overview(
                os.path.join(options.output, name), addon)
        except Exception as e:
            pass

    html_output.html_postprocess_all(batchlist)

def write_unit_ids_UNUSED():
    # Write a list with all unit ids, just for fun.
    uids = list(wesnoth.unit_lookup.keys())
    def by_race(u1, u2):
        r = cmp(wesnoth.unit_lookup[u1].rid,
                wesnoth.unit_lookup[u2].rid)
        if r == 0:
            r = cmp(u1, u2)
        return r
    uids.sort(by_race)
    race = None
    f = MyFile(os.path.join(options.output, "uids.html"), "w")
    f.write("<html><body>")
    for uid in uids:
        u = wesnoth.unit_lookup[uid]
        if u.rid != race:
            if race is not None:
                f.write("</ul>")
            f.write("<p>%s</p>\n" % (u.rid,))
            f.write("<ul>")
            race = u.rid
        f.write("<li>%s</li>\n" % (uid, ))
    f.write("</ul>")
    f.write("</body></html>")
    f.close()

def write_animation_statistics(wesnoth):
    # Write animation statistics
    f = html_output.MyFile(os.path.join(options.output, "animations.html"), "w")
    animations.write_table(f, wesnoth)
    f.close()

if __name__ == '__main__':
    # We change the process name to "wmlunits"
    try:
        import ctypes
        libc = ctypes.CDLL("libc.so.6")
        libc.prctl(15, b"wmlunits", 0, 0, 0)
    except: # oh well...
        pass

    global options
    global image_collector

    ap = argparse.ArgumentParser()
    ap.add_argument("-C", "--config-dir",
                    help="Specify the user configuration dir (wesnoth --config-path).")
    ap.add_argument("-D", "--data-dir",
                    help="Specify the wesnoth data dir (wesnoth --path).")
    ap.add_argument("-l", "--language", default="all",
                    help="Specify a language to use. Else output is produced for all languages.")
    ap.add_argument("-o", "--output",
                    help="Specify the output directory.")
    ap.add_argument("-n", "--nocopy", action="store_true",
                    help="No copying of files. By default all images are copied to the output dir.")
    ap.add_argument("-w", "--wesnoth",
                    help="Specify the wesnoth executable to use. Whatever data " +
                    "and config paths that executable is configured for will be " +
                    "used to find game files and addons.")
    ap.add_argument("-t", "--transdir",
                    help="Specify the directory with gettext message catalogues. " +
                    "Defaults to ./translations.", default="translations")
    ap.add_argument("-T", "--timeout",
                    help="Use the timeout iven in seconds when parsing WML, default is 20 " +
                    "seconds.")
    ap.add_argument("-r", "--reparse", action="store_true",
                    help="Reparse everything.")
    ap.add_argument("-a", "--addons",
                    help="Specify path to a folder with all addons. This should be " +
                    "outside the user config folder.")
    ap.add_argument("-L", "--list",
                    help="List available eras and campaigns.")
    ap.add_argument("-B", "--batch",
                    help="Batch process the given list.")
    ap.add_argument("-A", "--addons-only", action="store_true",
                    help="Do only process addons (for debugging).")
    ap.add_argument("-v", "--verbose", action="store_true")
    ap.add_argument("-W", "--wiki", action="store_true",
                    help="write wikified units list to stdout")
    options = ap.parse_args()

    html_output.options = options
    helpers.options = options
    unit_tree.overview.options = options
    wiki_output.options = options

    if not options.output and not options.wiki:
        sys.stderr.write("Need --output (or --wiki).\n")
        ap.print_help()
        sys.exit(-1)

    if options.output:
        options.output = os.path.expanduser(options.output)

    if options.timeout:
        TIMEOUT = int(options.timeout)

    if not options.wesnoth:
        options.wesnoth = "wesnoth"

    if not options.data_dir:
        options.data_dir = shell_out([options.wesnoth, "--path"]).strip().decode("utf8")
        print("Using " + options.data_dir + " as data dir.")

    if not options.config_dir:
        options.config_dir = shell_out([options.wesnoth, "--config-path"]).strip().decode("utf8")
        print("Using " + options.config_dir + " as config dir.")

    if not options.transdir:
        options.transdir = os.getcwd()

    if options.wiki:
        wiki_output.main()
        sys.exit(0)

    image_collector = helpers.ImageCollector(options.wesnoth,
                                             options.config_dir,
                                             options.data_dir)
    html_output.image_collector = image_collector

    if options.language == "all":
        languages = []
        parser = wmlparser3.Parser(options.wesnoth,
                                   options.config_dir,
                                   options.data_dir)
        parser.parse_text("{languages}")

        for locale in parser.get_all(tag="locale"):
            isocode = locale.get_text_val("locale")
            name = locale.get_text_val("name")
            if isocode == "ang_GB":
                continue
            languages.append(isocode)
        languages.sort()
    else:
        languages = options.language.split(",")

    if not options.list and not options.batch:
        sys.stderr.write("Need --list or --batch (or both).\n")
        sys.exit(-1)

    if options.output:
        # Generate output dir.
        if not os.path.isdir(options.output):
            os.mkdir(options.output)

    if options.list:
        list_contents()

    if options.batch:
        batch_process()

        unit_tree.overview.main(options.output)

        if not options.nocopy:
            copy_images()

    html_output.write_index(options.output)