#!/usr/bin/python
#
# Usage:
#    pkgdraw <header> <outfile> [<parser names>...] [-- <cpp options>...]
#
# Extract packet structure from <header>. Write generated PNG diagram
# to <outfile>.  If <parser names> is given, it is a list names of
# parsers to generate diagrams for. All other parsers from the same
# header file will be skipped.
#
# <cpp options> are parsed to the C preprocessor while processing the
# header file.
#
# pkgdraw will interpret most SENF_PARSER statements, it does however
# *not* understand GOTO statements. Special comments may be added
# directly before or after a field to pass parameters to pkgdraw
#
#    SENF_PARSER_FIELD( id, senf::UInt16Parser ); //<pkgdraw: <option>, <option>...
#
# or
#
#    //>pkgdraw: <option>, <option>...
#    SENF_PARSER_FIELD( id, senf::UInt16Parser );
#
# <option> is any valid option:
#
#   hide                Completely skip this field (Helps with GOTO)
#   name=<name>         Sets the field name to <name>
#   size=<min>[-<max>]  Sets the field size in bits.
#

import sys, re, signal, tempfile, os, os.path, shutil, atexit

basedir=os.path.abspath(os.path.split(sys.argv[0])[0])

charsPerBit = 1.4

TEX_HEADER = r"""\documentclass{scrartcl}
\usepackage[german]{babel}
\usepackage[latin1]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{ae,aecompl}

\usepackage{color}
\usepackage{bytefield}

\pagestyle{empty}

\begin{document}
\sffamily
"""

PACKET_HEADER=r"""\begin{bytefield}{32}
  \bitheader{0-31} \\
"""

PACKET_FOOTER=r"""\end{bytefield}
\bigskip

"""

TEX_FOOTER = r"""\end{document}
"""

def formatField(width, start, size):
    areas = []
    sz = width - start
    while size > 0:
        if sz > size:
            areas.append({'start': start,
                          'size': size})
            size = 0
        else:
            areas.append({'start': start,
                          'size': sz})
            size -= sz
        sz = width
        start = 0
    for i in range(len(areas)-1):
        if areas[i]['start'] < areas[i+1]['start']+areas[i+1]['size']:
            areas[i]['bottom'] = False
            areas[i+1]['top'] = False
    return areas

def formatSimpleField(width, start, field):
    areas = formatField(width, start, field['size'])
    nameix = 0
    namesz = 0
    for i in range(len(areas)):
        if areas[i]['size'] > namesz:
            namesz = areas[i]['size']
            nameix = i
    areas[nameix]['name'] = field['name'][:int(areas[nameix]['size'] * charsPerBit)]
    if field['name'] and len(areas) == 2 and areas[0].get('bottom',True):
        if areas[0].get('name','') : ix = 1
        else                       : ix = 0
        if 6 <= int(areas[ix]['size'] * charsPerBit):
            areas[ix]['name'] = '(cont)'
    return areas
    
def formatPacket(width, fields):
    rows = [ [] ]
    start = 0
    for field in fields:
        areas = []
        if field.get('repeat', False):
            if start > 0 and start < width:
                areas.append({ 'start': start, 'size': width-start, 'bottom': False,
                               'right': False})
            start = 0
        if field.get('size',None):
            areas.extend(formatSimpleField(width, start, field))
            start = areas[-1]['start'] + areas[-1]['size']
        elif field.get('minsize', None):
            f = dict(field)
            f['size'] = field['minsize']
            areas.extend(formatSimpleField(width, start, f))
            start = areas[-1]['start'] + areas[-1]['size']
            if start >= width : start = 0
            addareas = formatField(width, start, field['maxsize'] - field['minsize'])
            for area in addareas:
                area['filled'] = True
            areas += addareas
            start = areas[-1]['start'] + areas[-1]['size']
            if start > 0 and start < width:
                areas.append({ 'start': start, 'size': width-start, 'bottom': False,
                               'right': False})
            start = 0
        else:
            if start > 0 and start < width:
                areas.append({ 'start': start, 'size': width-start, 'bottom': False,
                               'right': False})
            areas.extend([ { 'start': 0, 'size': width, 'bottom': False,
                             'name': field['name'] },
                           { 'start': 0, 'size': width, 'skip': True },
                           { 'start': 0, 'size': width, 'top': False } ])
            start = 0
        if field.get('optional', False):
            for area in areas:
                area['optional'] = True
            if start > 0 and start < width:
                areas.append({ 'start': start, 'size': width-start, 'bottom': False,
                               'right': False})
            start = 0
        if field.get('repeat'):
            if start > 0 and start < width:
                areas.append({ 'start': start, 'size': width-start, 'bottom': False,
                               'right': False})
            start = 0
            areas.append({ 'start': 0, 'size': width, 'dots': True })
            da = areas[(areas[0].get('right', True) and (0,) or (1,))[0]:-1]
            for i in range(len(da)):
                if da[i].get('name','') :
                    da[i] = dict(da[i])
                    del da[i]['name']
            areas.extend(da)
        if start == width : start = 0
        
        while areas:
            while areas and (not(rows[-1]) or rows[-1][-1]['start'] + rows[-1][-1]['size'] < width):
                if areas[0].get('right', True) == False:
                    # This is a fillup field. Check, wether to draw top:
                    if len(rows) <= 1:
                        areas[0]['top'] = False
                    elif rows[-2][-1].get('bottom', True) or not rows[-2][-1].get('right', True):
                        areas[0]['top'] = False
                rows[-1].append(areas.pop(0))
            if areas:
                rows.append([])
    return rows

def texquote(s):
    s = s.replace('_', '\\_')
    return s

def makeTex(rows):
    lines = []
    for row in rows:
        line = []
        for area in row:
            sides=""
            if area.get('left',   True) : sides += "l"
            if area.get('right',  True) : sides += "r"
            if area.get('top',    True) : sides += "t"
            if area.get('bottom', True) : sides += "b"
            if sides == "lrtb" : sides = ""
            else               : sides = "[%s]" % sides
            if area.get('filled', False):
                line.append(r"\bitbox%s{%s}{\color[gray]{0.93}\rule{\width}{\height}}" % (sides, area['size']))
            elif area.get('skip', False):
                line.append(r"\skippedwords")
            elif area.get('dots', False):
                line.append(r"\wordbox[]{1}{$\vdots$\\[1ex]}")
            else:
                name = texquote(area.get('name',''))
                if name and area.get('optional', False):
                    name = "[%s]" % name
                line.append(r"\bitbox%s{%s}{\strut %s}" % (sides, area['size'], name))
        lines.append(" & ".join(line))
    return " \\\\\n".join(lines) + "\n"

COMMENT_RE = re.compile(r'//.*?$|/\*.*?\*/|"(?:\\.|[^\\"])*"', re.S | re.M)

def stripComments(text):
    def replacer(match):
        s = match.group(0)
        if s.startswith('//<pkgdraw:') or s.startswith('//>pkgdraw:'):
            return "@@" + s[2:]
        if s.startswith('/'):
            return ""
        return s
    return COMMENT_RE.sub(replacer, text)

SENF_INCLUDE_RE = re.compile(r"#\s*include\s*SENF_")

def quoteMacros(text):
    return SENF_INCLUDE_RE.sub("PKGDRAW_", text).replace("SENF_PARSER_","PKGDRAW_PARSER_")

def cppExpand(text, cppopts, dir):
    tmpf = tempfile.NamedTemporaryFile(dir=dir)
    tmpf.write(text)
    tmpf.flush()
    cmd = "gcc %s -E -o - -x c++-header %s" % (" ".join(cppopts), tmpf.name)
    return os.popen(cmd).read()
    
FIELD_TYPES = {
    'UInt8Parser' :  {'size': 8 },
    'UInt16Parser' : {'size': 16 },
    'UInt24Parser' : {'size': 24 },
    'UInt32Parser' : {'size': 32 },
    'UInt64Parser' : {'size': 64 },
    'Int8Parser' : {'size': 8 },
    'Int16Parser' : {'size': 16 },
    'Int24Parser' : {'size': 24 },
    'Int32Parser' : {'size': 32 },
    'Int64Parser' : {'size': 64 },
    'UInt16LSBParser' : {'size': 16 },
    'UInt24LSBParser' : {'size': 24 },
    'UInt32LSBParser' : {'size': 32 },
    'UInt64LSBParser' : {'size': 64 },
    'Int16LSBParser' : {'size': 16 },
    'Int24LSBParser' : {'size': 24 },
    'Int32LSBParser' : {'size': 32 },
    'Int64LSBParser' : {'size': 64 },
    'MACAddressParser': {'size': 48 },
    'INet4AddressParser' : {'size': 32 },
    'INet6AddressParser' : {'size': 128 },
    'VoidPacketParser' : {'size': 0 },
    }
    
def parse_FIELD(args, flags):
    args = [ arg.strip() for arg in args.split(',') ]
    if len(args) != 2:
        sys.stderr.write("Failed to parse FIELD: %s\n" % args)
        return None
    field = dict(FIELD_TYPES.get(args[1].split(':')[-1], {}))
    field['name'] = args[0]
    return field

def parse_PRIVATE_FIELD(args, flags):
    return parse_FIELD(args, flags)

def parse_FIELD_RO(args, flags):
    return parse_FIELD(args, flags)

def parse_BITFIELD(args, flags):
    args = [ arg.strip() for arg in args.split(',') ]
    if len(args) != 3:
        sys.stderr.write("Failed to parse BITFIELD: %s\n" % args)
        return None
    try:
        size = int(args[1])
    except ValueError:
        sys.stderr.write("Failed to parse BITFIELD: %s\n" % args)
        return None
    return { 'size' : size, 'name' : args[0] }

def parse_PRIVATE_BITFIELD(args, flags):
    return parse_BITFIELD(args, flags)

def parse_BITFIELD_RO(args, flags):
    return parse_BITFIELD(args, flags)

def parse_SKIP(args, flags):
    args = args.split(',')[0]
    try:
        bytes = int(args.strip())
    except ValueError:
        sys.stderr.write("Failed to parse SKIP: %s\n" % args)
        return None
    return { 'size': 8*bytes, 'name': '' }

def parse_SKIP_BITS(args, flags):
    try:
        bits = int(args.strip())
    except ValueError:
        sys.stderr.write("Failed to parse SKIP_BITS: %s\n" % args)
        return None
    return { 'size': bits, 'name': '' }

def parse_VECTOR(args, flags):
    args = [ arg.strip() for arg in args.split(',') ]
    if len(args) < 3:
        sys.stderr.write("Failed to aprse VECTOR: %s\n" % args)
        return None
    field = dict(FIELD_TYPES.get(args[-1].split(':')[-1], {}))
    field['name'] = args[0]
    field['repeat'] = True
    return field

def parse_LIST(args, flags):
    return parse_VECTOR(args, flags)

VARIANT_FIELD_RE_STR = r"""
    \(\s*(?:
        ([a-zA-Z0-9_:]+) |
        id\(\s*
            [a-zA-Z0-9_]+\s*,\s*
            (?:
                ([a-zA-Z0-9_:]+) |
                key\(\s*[^,]*,\s*([a-zA-Z0-9_:]+)\s*\)
            )\s*\) |
        ids\(\s*
            [a-zA-Z0-9_]+\s*,\s*
            [a-zA-Z0-9_]+\s*,\s*
            [a-zA-Z0-9_]+\s*,\s*
            (?:
                ([a-zA-Z0-9_:]+) |
                key\(\s*[^,]*,\s*([a-zA-Z0-9_:]+)\s*\)
            )\s*\) |
        novalue\(\s*
            [a-zA-Z0-9_]+\s*,\s*
            (?:
                ([a-zA-Z0-9_:]+) |
                key\(\s*[^,]*,\s*([a-zA-Z0-9_:]+)\s*\)
            )\s*\)
    )\s*\)
"""

VARIANT_FIELD_RE = re.compile(VARIANT_FIELD_RE_STR, re.X)
VARIANT_FIELDS_RE = re.compile(",\s*((?:%s\s*)+)$" % VARIANT_FIELD_RE_STR, re.X)

def parse_VARIANT(args, flags):
    name = args.split(',',1)[0].strip()
    fields_match = VARIANT_FIELDS_RE.search(args)
    if not fields_match:
        return { 'name': name }
    fields_str = fields_match.group(1)
    optional = False
    minsize = None
    maxsize = None
    for field_match in VARIANT_FIELD_RE.finditer(fields_str):
        parser = ([ group for group in field_match.groups() if group ] + [ None ])[0]
        field = dict(FIELD_TYPES.get(parser.split(':')[-1], {}))
        if field.has_key('minsize'):
            if minsize is None or field['minsize'] < minsize:
                minsize = field['minsize']
            if maxsize is None or field['maxsize'] > maxsize:
                maxsize = field['maxsize']
        elif field.has_key('size'):
            if field['size'] == 0:
                optional = True
            else:
                if minsize is None or field['size'] < minsize:
                    minsize = field['size']
                if maxsize is None or field['size'] > maxsize:
                    maxsize = field['size']
    if minsize is not None and minsize == maxsize:
        return { 'name': name, 'size': minsize, 'optional': optional }
    elif minsize is not None:
        return { 'name': name, 'minsize': minsize, 'maxsize': maxsize, 'optional': optional }
    else:
        return { 'name': name, 'optional': optional }

def parse_PRIVATE_VARIANT(args, flags):
    return parse_VARIANT(args, flags)

def parse_INIT(args, flags):
    return None

PARSER_START_RE = re.compile(r"PKGDRAW_(FIXED_)?PARSER\s*\(\s*\)")
PARSER_END_RE = re.compile(r"PKGDRAW_PARSER_FINALIZE\s*\(([^)]*)\)\s*;")
PARSER_FIELD_RE = re.compile(r"(?:@@>pkgdraw:(.*)$\s*)?PKGDRAW_PARSER_([A-Z_]+)\s*\(([^;]*)\)\s*;(?:\s*@@<pkgdraw:(.*)$)?", re.M)

def scanPackets(data):
    global FIELD_TYPES
    
    packets = {}
    packetOrder = []
    end = 0
    while True:
        start =  PARSER_START_RE.search(data, end)
        if not start: return (packets, packetOrder)
        start = start.end(0)
        end = PARSER_END_RE.search(data, start)
        if not end: return (packets, packetOrder)
        name=end.group(1).strip()
        end = end.start(0)
        packets[name] = scanFields(data[start:end])
        packetOrder.append(name)
        minsize = maxsize = 0
        for field in packets[name]:
            if maxsize is not None:
                if field.get('repeat', False):
                    maxsize = None
                elif field.get('size', None) is not None:
                    maxsize += field['size']
                elif field.get('minsize', None) is not None:
                    maxsize += field['maxsize']
                else:
                    maxsize = None
            if not field.get('optional', False):
                if field.get('size', None) is not None:
                    minsize += field['size']
                elif field.get('minsize', None) is not None:
                    minsize += field['minsize']
        if minsize is not None and maxsize is not None:
            if minsize == maxsize:
                FIELD_TYPES[name] = { 'size' : minsize }
            else:
                FIELD_TYPES[name] = { 'minsize' : minsize, 'maxsize' : maxsize }

def scanFields(data):
    fields = []
    for match in PARSER_FIELD_RE.finditer(data):
        tp = match.group(2)
        flags = dict([ ([ arg.strip() for arg in flag.strip().split('=',1) ]+[True])[:2]
                       for flag in ((match.group(1) or '')+(match.group(4) or '')).split(',') ])
        if flags.has_key('hide') : continue
        parser = globals().get("parse_%s" % tp, None)
        if parser:
            field = parser(match.group(3).strip(), flags)
            if field:
                if flags.has_key('name') : field['name'] = flags['name']
                field['name'] = field['name'].strip('_')
                if flags.has_key('size'):
                    if '-' in flags['size']:
                        field['minsize'], field['maxsize'] = map(int, flags['size'].split('-',1))
                        del field['size']
                    else:
                        field['size'] = int(flags['size'])
                if not field['name'] and fields and not fields[-1]['name'] \
                       and field.has_key('size') and fields[-1].has_key('size'):
                    fields[-1]['size'] += field['size']
                else:
                    fields.append(field)
        else:
            sys.stderr.write("Unknown parser type: %s\n" % tp)
    return fields

tmpdir = tempfile.mkdtemp(prefix="pkgdraw_")

def cleanup():
    global tmpdir
    shutil.rmtree(tmpdir)

signal.signal(signal.SIGINT, cleanup)
signal.signal(signal.SIGTERM, cleanup)
signal.signal(signal.SIGHUP, cleanup)
atexit.register(cleanup)

args = sys.argv[1:]
names = []
gccopts = []

if len(args)<2 or args[0] == '--' or args[1] == '--':
    sys.stderr.write("Usage: %s <header> <outfile> [<parser names>...] [-- <cpp options>...]\n"
                     % sys.argv[0])
    sys.exit(1)

source = args.pop(0)
target = args.pop(0)

while args and args[0] != '--' : names.append(args.pop(0))
if args : gccopts = args[1:]

data, order = scanPackets(cppExpand(quoteMacros(stripComments(file(source).read())),
                                    gccopts, os.path.dirname(source)))

texf = file(os.path.join(tmpdir, "fields.tex"),"w")
texf.write(TEX_HEADER)

if not names:
    order.reverse()
    names = order

for name in names:
    texf.write("\\textbf{%s}\n\\bigskip\\par\n" % texquote(name))
    texf.write(PACKET_HEADER)
    texf.write(makeTex(formatPacket(32, data[name])))
    texf.write(PACKET_FOOTER)
    
texf.write(TEX_FOOTER)
texf.close()

if os.system("cd %s; %s/textogif -png -dpi 80 -res 0.25 fields >pkgdraw.log 2>&1"
             % (tmpdir, basedir)) != 0:
    sys.stderr.write("Conversion failed. See %s\n" % tmpdir)
    os._exit(1)

file(target,"w").write(file(os.path.join(tmpdir, "fields.png")).read())