"""Read .proto files as text and build a field table: {message: {num: [name, type, repeated]}}.

Types are either a scalar keyword ('string', 'bytes', 'double', 'int32', 'bool', ...),
'enum', or the fully qualified name of a message ('rv.data.UUID').
"""
import json
import os
import re
import sys

SCALARS = {"double", "float", "int32", "int64", "uint32", "uint64", "sint32", "sint64",
           "fixed32", "fixed64", "sfixed32", "sfixed64", "bool", "string", "bytes"}


def tokenize(src):
    src = re.sub(r"//[^\n]*", "", src)
    src = re.sub(r"/\*.*?\*/", "", src, flags=re.S)
    return re.findall(r'"(?:[^"\\]|\\.)*"|[A-Za-z_][\w.]*|\.[A-Za-z_][\w.]*|-?\d+(?:\.\d+)?|[{}\[\]();=<>,]', src)


def parse_file(path, messages, enums):
    toks = tokenize(open(path).read())
    pkg = ""
    i = 0
    scope = []  # stack of (kind, fqname)
    pending = []  # fields with unresolved types: (msg_fq, num, name, rawtype, repeated, scope_names)

    def cur_msg():
        for kind, fq in reversed(scope):
            if kind == "message":
                return fq
        return None

    while i < len(toks):
        t = toks[i]
        if t == "package":
            pkg = toks[i + 1]
            i += 3
            continue
        if t in ("syntax", "import", "option"):
            while toks[i] != ";":
                i += 1
            i += 1
            continue
        if t in ("message", "enum", "oneof", "service"):
            name = toks[i + 1]
            parent = scope[-1][1] if scope else pkg
            fq = parent + "." + name if parent else name
            if t == "oneof":
                fq = scope[-1][1]  # fields belong to the enclosing message
            scope.append((t, fq))
            if t == "message":
                messages.setdefault(fq, {})
            elif t == "enum":
                enums.add(fq)
            i += 3  # name {
            continue
        if t == "}":
            scope.pop()
            i += 1
            continue
        if t == "reserved" or t == "extensions":
            while toks[i] != ";":
                i += 1
            i += 1
            continue
        if scope and scope[-1][0] == "enum":
            while toks[i] != ";":
                i += 1
            i += 1
            continue
        if scope and scope[-1][0] in ("message", "oneof"):
            repeated = False
            if t in ("repeated", "optional", "required"):
                repeated = t == "repeated"
                i += 1
                t = toks[i]
            if t == "map":
                # map<k, v> name = n;
                j = toks.index(";", i)
                name = toks[j - 3]
                num = int(toks[j - 1]) if toks[j - 1] != "]" else None
                vt = toks[i + 4]
                pending.append((cur_msg(), int(toks[toks.index("=", i) + 1]), name, "map:" + vt, True,
                                [fq for _, fq in scope]))
                i = j + 1
                continue
            rawtype = t
            name = toks[i + 1]
            num = int(toks[i + 3])
            j = i + 4
            if toks[j] == "[":
                while toks[j] != "]":
                    j += 1
                j += 1
            pending.append((cur_msg(), num, name, rawtype, repeated, [fq for _, fq in scope]))
            i = j + 1
            continue
        i += 1
    return pending


def resolve(rawtype, scopes, messages, enums):
    if rawtype in SCALARS:
        return rawtype
    if rawtype.startswith("."):
        cand = rawtype[1:]
        if cand in messages:
            return cand
        if cand in enums:
            return "enum"
        raise KeyError(rawtype)
    # search enclosing scopes from innermost outwards
    for sc in reversed(scopes):
        parts = sc.split(".")
        for k in range(len(parts), -1, -1):
            cand = ".".join(parts[:k] + [rawtype]) if k else rawtype
            if cand in messages:
                return cand
            if cand in enums:
                return "enum"
    raise KeyError(rawtype)


def build(proto_dir):
    messages, enums, pending = {}, set(), []
    for fn in sorted(os.listdir(proto_dir)):
        if fn.endswith(".proto"):
            pending += parse_file(os.path.join(proto_dir, fn), messages, enums)
    for gdir in ("google/protobuf",):
        d = os.path.join(proto_dir, gdir)
        if os.path.isdir(d):
            for fn in sorted(os.listdir(d)):
                if fn.endswith(".proto"):
                    pending += parse_file(os.path.join(d, fn), messages, enums)
    for msg, num, name, rawtype, rep, scopes in pending:
        if rawtype.startswith("map:"):
            messages[msg][num] = [name, "map", True]
            continue
        messages[msg][num] = [name, resolve(rawtype, scopes, messages, enums), rep]
    return messages


def reachable(messages, root):
    out, todo = {}, [root]
    while todo:
        m = todo.pop()
        if m in out:
            continue
        out[m] = messages[m]
        for name, typ, rep in messages[m].values():
            if typ in messages and typ not in out:
                todo.append(typ)
    return out


if __name__ == "__main__":
    msgs = build(sys.argv[1])
    sub = reachable(msgs, "rv.data.Presentation")
    sub.update(reachable(msgs, "rv.data.Template.Document"))
    sub.update(reachable(msgs, "rv.data.PlaylistDocument"))
    json.dump({k: {str(n): v for n, v in f.items()} for k, f in sorted(sub.items())},
              open(sys.argv[2], "w"), separators=(",", ":"), sort_keys=True)
    print(len(msgs), "messages parsed;", len(sub), "reachable from Presentation + Template.Document")