#!/usr/bin/env python3
# Converts Google's protobuf python definitions of TREZOR wire messages
# to plain-python objects as used in TREZOR Core and python-trezor

import argparse
import importlib
import logging
import os
import sys


def import_pb2(name):
    return importlib.import_module("%s_pb2" % name)


def create_message_import(name):
    return "from .%s import %s" % (name, name)


def create_const(name, value, is_upy):
    if is_upy:
        return "%s = const(%s)" % (name, value)
    else:
        return "%s = %s" % (name, value)


def remove_from_start(s, prefix):
    if s.startswith(prefix):
        return s[len(prefix):]
    else:
        return s


def process_message_imports(descriptor):
    imports = set()

    for field in descriptor.fields:
        if field.type == field.TYPE_MESSAGE:
            imports.add(field.message_type.name)

    for name in sorted(imports):
        yield create_message_import(name)


def process_message(descriptor, protobuf_module, msg_id, indexfile, is_upy):
    logging.debug("Processing message %s", descriptor.name)

    if is_upy:
        yield "import protobuf as p"
    else:
        yield "from .. import protobuf as p"
    yield from process_message_imports(descriptor)

    yield ""
    yield ""
    yield "class %s(p.MessageType):" % descriptor.name

    if descriptor.fields_by_number:
        yield "    FIELDS = {"
    elif msg_id is None:
        yield "    pass"

    for number, field in descriptor.fields_by_number.items():
        field_name = field.name
        field_type = None
        repeated = (field.label == field.LABEL_REPEATED)
        required = (field.label == field.LABEL_REQUIRED)

        types = {
            field.TYPE_UINT64:  'p.UVarintType',
            field.TYPE_UINT32:  'p.UVarintType',
            field.TYPE_ENUM:    'p.UVarintType',
            field.TYPE_SINT32:  'p.Sint32Type',
            field.TYPE_SINT64:  'p.Sint64Type',
            field.TYPE_STRING:  'p.UnicodeType',
            field.TYPE_BOOL:    'p.BoolType',
            field.TYPE_BYTES:   'p.BytesType'
        }

        if field.type == field.TYPE_MESSAGE:
            field_type = field.message_type.name
        else:
            try:
                field_type = types[field.type]
            except KeyError:
                raise ValueError("Unknown field type %d for field %s" % (field.type, field_name))

        comments = []
        if required:
            comments.append('required')
        if field.has_default_value:
            comments.append("default=%s" % repr(field.default_value))

        if comments:
            comment = "  # %s" % ' '.join(comments)
        else:
            comment = ''

        if repeated:
            flags = 'p.FLAG_REPEATED'
        else:
            flags = '0'

        yield "        %d: ('%s', %s, %s),%s" % (number, field_name, field_type, flags, comment)

    if descriptor.fields_by_name:
        yield "    }"

    if msg_id is not None:
        yield "    MESSAGE_WIRE_TYPE = %d" % msg_id
        if indexfile is not None:
            indexfile.write(create_const(t, msg_id, is_upy))


def process_enum(descriptor, is_upy):
    logging.debug("Processing enum %s", descriptor.name)

    if is_upy:
        yield "from micropython import const"
        yield ""

    for name, value in descriptor.values_by_name.items():
        # Remove type name from the beginning of the constant
        # For example "PinMatrixRequestType_Current" -> "Current"
        enum_prefix = descriptor.name
        name = remove_from_start(name, "%s_" % enum_prefix)

        # If type ends with *Type, but constant use type name without *Type, remove it too :)
        # For example "ButtonRequestType & ButtonRequest_Other" => "Other"
        if enum_prefix.endswith("Type"):
            enum_prefix, _ = enum_prefix.rsplit("Type", 1)
            name = remove_from_start(name, "%s_" % enum_prefix)

        yield create_const(name, value.number, is_upy)


def process_file(descriptor, protobuf_module, genpath, indexfile, modlist, is_upy):
    logging.info("Processing module %s", descriptor.name)

    msg_types = import_pb2('messages').MessageType

    for name, message_descriptor in sorted(descriptor.message_types_by_name.items()):
        # Find message type for given class
        try:
            msg_id = msg_types.Value("MessageType_%s" % name)
        except ValueError:
            msg_id = None

        out = process_message(message_descriptor, protobuf_module, msg_id, indexfile, is_upy)
        write_to_file(genpath, name, out)
        if modlist:
            modlist.write(create_message_import(name) + "\n")

    for name, enum_descriptor in descriptor.enum_types_by_name.items():
        out = process_enum(enum_descriptor, is_upy)
        write_to_file(genpath, name, out)
        if modlist:
            modlist.write("from . import %s\n" % name)


def write_to_file(genpath, t, out):
    # Write generated sourcecode to given file
    with open(os.path.join(genpath, "%s.py" % t), 'w') as f:
        f.write("# Automatically generated by pb2py\n")
        for line in out:
            f.write(line + "\n")


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    parser = argparse.ArgumentParser()
    parser.add_argument('module', help="Name of module to generate")
    parser.add_argument('genpath', help="Directory for generated source code")
    parser.add_argument('-P', '--protobuf-module', default="protobuf", help="Name of protobuf module")
    parser.add_argument('-i', '--indexfile', type=argparse.FileType('a'), help="Generate index file of wire types")
    parser.add_argument('-l', '--modlist', type=argparse.FileType('a'), help="Generate list of modules")
    parser.add_argument('-p', '--protopath', type=str, help="Path to search for pregenerated Google's python sources")
    parser.add_argument('-m', '--micropython', action='store_true', help="Use micropython-favoured source code")
    args = parser.parse_args()

    if args.protopath:
        sys.path.append(args.protopath)

    # This must be done after sys.path.append
    module = import_pb2(args.module)

    process_file(module.DESCRIPTOR, args.protobuf_module, args.genpath, args.indexfile, args.modlist, args.micropython)