#!/usr/bin/env python3
"""
Tool for generating fuselage (interface between P4 and HLS C/C++).
"""

import argparse, os.path, re
from collections import OrderedDict

def main():
  """
  Main function
  """
  print('Fuselage Factory')
  print('================\n')

  parser = argparse.ArgumentParser()
  parser.add_argument('stubFilename', help = 'Verilog external function stub')
  parser.add_argument('fuselageFilename', help = 'Output fuselage file')
  parser.add_argument('headerFilename', help = 'Output HLS C/C++ header')
  parser.add_argument('hlsFilename', help = 'Output HLS C/C++ external function stub')
  args = parser.parse_args()

  parseInfo = parseStub(args.stubFilename)
  generateFuselage(parseInfo, args.fuselageFilename)
  generateHeader(parseInfo, args.headerFilename)
  print('Your fuselage has been completed.  Be aware that it is not airworthy yet!')

  if (args.hlsFilename != None):
    generateStub(parseInfo, args.hlsFilename, args.headerFilename)

def parseStub(filename):
  print('Checking pilot license...')
  with open(filename, 'rt') as stubFile:
    content = stubFile.read()
  result = re.search(r'module (\S+) \(\n(.*?)\n\);', content, re.MULTILINE | re.DOTALL)
  parseInfo = {'module': result.group(1), 'ports': result.group(2)}
  io = re.search(r'(^(input|output).*^(input|output).*?$)', content, re.MULTILINE | re.DOTALL).group(1)
  io = io.replace(' /* unused */ ', '')
  parseInfo['io'] = io.replace(' /* undriven */ ', '')
  wires = re.search(r'(^wire.*^wire.*?$)', content, re.MULTILINE | re.DOTALL).group(1)
  wires = wires.replace(' /* unused */ ', '')
  parseInfo['wires'] = wires.replace(' /* undriven */ ', '')
  extractTuples(parseInfo, content)
  parseInfo['function'] = re.match(r'(\S+)_.*_.*', parseInfo['module']).group(1)
  return parseInfo

def extractTuples(parseInfo, content):
  for direction in ['in', 'out']:
    tuples = OrderedDict()
    for line in parseInfo['io'].split('\n'):
      result = re.match(r'\S+ \[(\S+):0\] (tuple_' + direction + '_\S+)_DATA', line)
      if result != None:
        name = result.group(2)
        width = int(result.group(1)) + 1
        tuples[name] = OrderedDict([('width', width)])
    for tup in tuples:
      pattern = r'/\* Tuple format for ' + direction + 'put: ' + tup + '\n(.*?)\n\n\*/'
      result = re.search(pattern, content, re.MULTILINE | re.DOTALL)
      if result == None:
        continue
      lines = result.group(1)
      fields = OrderedDict()
      for line in lines.split('\n'):
        result = re.match('\s+\[(\S+):\S+\]\s+:\s+(\S+)', line)
        fields[result.group(2)] = int(result.group(1)) + 1
      offset = 0
      for name, width in reversed(list(fields.items())):
        fields[name] -= offset
        offset = width
      tuples[tup]['fields'] = fields
    parseInfo[direction + 'put_tuples'] = tuples

def generateFuselage(parseInfo, filename):
  print('Please wait while a fuselage is built...')

  with open(filename, "wt") as outputFile:
    outputFile.write('// Code generated by the Fuselage Factory\n\n')
    outputFile.write('// WE DISCOURAGE MODIFYING THIS FILE DIRECTLY!!!\n\n')

    outputFile.write('`timescale 1 ps / 1 ps\n\n')

    outputFile.write('module ' + parseInfo['module'] + ' (\n')
    outputFile.write(parseInfo['ports'] + ',\n')
    outputFile.write('\tbackpressure_in,\n')
    outputFile.write('\tbackpressure_out\n')
    outputFile.write(');\n\n')

    outputFile.write(parseInfo['io'] + '\n')
    outputFile.write('input backpressure_in;\n')
    outputFile.write('output backpressure_out;\n\n')

    outputFile.write(parseInfo['wires'] + '\n\n')

    writeWires(parseInfo, outputFile)
    writeFunction(parseInfo, outputFile)
    writeTupleFIFO(parseInfo, outputFile)
    writePacketFIFO(parseInfo, outputFile)
    writeLogic(parseInfo, outputFile)

    outputFile.write('endmodule\n')

def writeWires(parseInfo, outputFile):
  function = parseInfo['function']

  input_width = 0
  for tup, info in parseInfo['input_tuples'].items():
    input_width += info['width']

  output_width = 0
  for tup, info in parseInfo['output_tuples'].items():
    output_width += info['width']

  outputFile.write('wire ' + function + '_ap_clk;\n')
  outputFile.write('wire ' + function + '_ap_rst;\n')
  outputFile.write('wire ' + function + '_ap_start;\n')
  outputFile.write('wire ' + function + '_ap_done;\n')
  outputFile.write('wire ' + function + '_ap_idle;\n')
  outputFile.write('wire ' + function + '_ap_ready;\n')
  outputFile.write('wire [' + str(input_width - 1) + ':0] ' + function + '_tuple_input_v_dout;\n')
  outputFile.write('wire ' + function + '_tuple_input_v_empty_n;\n')
  outputFile.write('wire ' + function + '_tuple_input_v_read;\n')
  outputFile.write('wire [' + str(output_width - 1) + ':0] ' + function + '_tuple_output_v;\n')
  outputFile.write('wire ' + function + '_tuple_output_v_ap_vld;\n')
  outputFile.write('wire ' + function + '_tuple_output_v_ap_ack;\n')
  outputFile.write('wire [70:0] ' + function + '_packet_input_v_dout;\n')
  outputFile.write('wire ' + function + '_packet_input_v_empty_n;\n')
  outputFile.write('wire ' + function + '_packet_input_v_read;\n')
  outputFile.write('wire [70:0] ' + function + '_packet_output_v;\n')
  outputFile.write('wire ' + function + '_packet_output_v_ap_vld;\n')
  outputFile.write('wire ' + function + '_packet_output_v_ap_ack;\n\n')

  outputFile.write('wire tuple_fifo_wr_en;\n')
  outputFile.write('wire tuple_fifo_rd_en;\n')
  outputFile.write('wire [' + str(input_width - 1) + ':0] tuple_fifo_din;\n')
  outputFile.write('wire [' + str(input_width - 1) + ':0] tuple_fifo_dout;\n')
  outputFile.write('wire tuple_fifo_empty;\n')
  outputFile.write('wire tuple_fifo_almost_full;\n')
  outputFile.write('wire tuple_fifo_full;\n\n')

  outputFile.write('wire packet_fifo_wr_en;\n')
  outputFile.write('wire packet_fifo_rd_en;\n')
  outputFile.write('wire [70:0] packet_fifo_din;\n')
  outputFile.write('wire [70:0] packet_fifo_dout;\n')
  outputFile.write('wire packet_fifo_empty;\n')
  outputFile.write('wire packet_fifo_almost_full;\n')
  outputFile.write('wire packet_fifo_full;\n\n')

def writeFunction(parseInfo, outputFile):
  function = parseInfo['function']
  outputFile.write(function + ' ' + function + '_inst\n')
  outputFile.write('(\n')
  outputFile.write('  .ap_clk(' + function + '_ap_clk),\n')
  outputFile.write('  .ap_rst(' + function + '_ap_rst),\n')
  outputFile.write('  .ap_start(' + function + '_ap_start),\n')
  outputFile.write('  .ap_done(' + function + '_ap_done),\n')
  outputFile.write('  .ap_idle(' + function + '_ap_idle),\n')
  outputFile.write('  .ap_ready(' + function + '_ap_ready),\n')
  outputFile.write('  .Tuple_input_V_dout(' + function + '_tuple_input_v_dout),\n')
  outputFile.write('  .Tuple_input_V_empty_n(' + function + '_tuple_input_v_empty_n),\n')
  outputFile.write('  .Tuple_input_V_read(' + function + '_tuple_input_v_read),\n')
  outputFile.write('  .Tuple_output_V(' + function + '_tuple_output_v),\n')
  outputFile.write('  .Tuple_output_V_ap_vld(' + function + '_tuple_output_v_ap_vld),\n')
  outputFile.write('  .Tuple_output_V_ap_ack(' + function + '_tuple_output_v_ap_ack),\n')
  outputFile.write('  .Packet_input_V_dout(' + function + '_packet_input_v_dout),\n')
  outputFile.write('  .Packet_input_V_empty_n(' + function + '_packet_input_v_empty_n),\n')
  outputFile.write('  .Packet_input_V_read(' + function + '_packet_input_v_read),\n')
  outputFile.write('  .Packet_output_V(' + function + '_packet_output_v),\n')
  outputFile.write('  .Packet_output_V_ap_vld(' + function + '_packet_output_v_ap_vld),\n')
  outputFile.write('  .Packet_output_V_ap_ack(' + function + '_packet_output_v_ap_ack)\n')
  outputFile.write(');\n\n')

  outputFile.write('assign ' + function + '_ap_clk = clk_line;\n')
  outputFile.write('assign ' + function + '_ap_rst = rst;\n')
  outputFile.write('assign ' + function + '_ap_start = 1;\n')
  outputFile.write('assign ' + function + '_tuple_input_v_dout = tuple_fifo_dout;\n')
  outputFile.write('assign ' + function + '_tuple_input_v_empty_n = ~tuple_fifo_empty;\n')
  outputFile.write('assign ' + function + '_tuple_output_v_ap_ack = 1;\n')
  outputFile.write('assign ' + function + '_packet_input_v_dout = packet_fifo_dout;\n')
  outputFile.write('assign ' + function + '_packet_input_v_empty_n = ~packet_fifo_empty;\n')
  outputFile.write('assign ' + function + '_packet_output_v_ap_ack = ~backpressure_in;\n\n')

def writeTupleFIFO(parseInfo, outputFile):
  width = 0
  for tup, info in parseInfo['input_tuples'].items():
    width += info['width']

  outputFile.write('defparam tuple_fifo.WRITE_DATA_WIDTH = ' + str(width) + ';\n');
  outputFile.write('defparam tuple_fifo.FIFO_WRITE_DEPTH = 512;\n');
  outputFile.write('defparam tuple_fifo.PROG_FULL_THRESH = 287;\n');
  outputFile.write('defparam tuple_fifo.PROG_EMPTY_THRESH = 287;\n');
  outputFile.write('defparam tuple_fifo.READ_MODE = "fwft";\n');
  outputFile.write('defparam tuple_fifo.WR_DATA_COUNT_WIDTH = 9;\n');
  outputFile.write('defparam tuple_fifo.RD_DATA_COUNT_WIDTH = 9;\n');
  outputFile.write('defparam tuple_fifo.DOUT_RESET_VALUE = "0";\n');
  outputFile.write('defparam tuple_fifo.FIFO_MEMORY_TYPE = "bram";\n\n');

  outputFile.write('xpm_fifo_sync tuple_fifo (\n')
  outputFile.write('\t.wr_en(tuple_fifo_wr_en),\n')
  outputFile.write('\t.din(tuple_fifo_din),\n')
  outputFile.write('\t.rd_en(tuple_fifo_rd_en),\n')
  outputFile.write('\t.sleep(1\'b0),\n')
  outputFile.write('\t.injectsbiterr(),\n')
  outputFile.write('\t.injectdbiterr(),\n')
  outputFile.write('\t.prog_empty(),\n')
  outputFile.write('\t.dout(tuple_fifo_dout), \n')
  outputFile.write('\t.empty(tuple_fifo_empty), \n')
  outputFile.write('\t.prog_full(tuple_fifo_almost_full),\n')
  outputFile.write('\t.full(tuple_fifo_full),\n')
  outputFile.write('\t.rd_data_count(),\n')
  outputFile.write('\t.wr_data_count(),\n') 
  outputFile.write('\t.wr_rst_busy(),\n')
  outputFile.write('\t.rd_rst_busy(),\n') 
  outputFile.write('\t.overflow(),\n')
  outputFile.write('\t.underflow(),\n')
  outputFile.write('\t.sbiterr(),\n')
  outputFile.write('\t.dbiterr(),\n') 
  outputFile.write('\t.wr_clk(clk_line),\n')
  outputFile.write('\t.rst(rst)\n')
  outputFile.write(');\n\n')

  first_tuple = next(iter(parseInfo['input_tuples']))
  tuples = reversed([tup + '_DATA' for tup in parseInfo['input_tuples'].keys()])

  outputFile.write('assign tuple_fifo_wr_en = ' + first_tuple + '_VALID;\n')
  outputFile.write('assign tuple_fifo_din = {' + ', '.join(tuples) + '};\n')
  outputFile.write('assign tuple_fifo_rd_en = ' + parseInfo['function'] + '_tuple_input_v_read;\n\n')

def writePacketFIFO(parseInfo, outputFile):
  outputFile.write('defparam packet_fifo.WRITE_DATA_WIDTH = 71;\n');
  outputFile.write('defparam packet_fifo.FIFO_WRITE_DEPTH = 512;\n');
  outputFile.write('defparam packet_fifo.PROG_FULL_THRESH = 287;\n');
  outputFile.write('defparam packet_fifo.PROG_EMPTY_THRESH = 287;\n');
  outputFile.write('defparam packet_fifo.READ_MODE = "fwft";\n');
  outputFile.write('defparam packet_fifo.WR_DATA_COUNT_WIDTH = 9;\n');
  outputFile.write('defparam packet_fifo.RD_DATA_COUNT_WIDTH = 9;\n');
  outputFile.write('defparam packet_fifo.DOUT_RESET_VALUE = "0";\n');
  outputFile.write('defparam packet_fifo.FIFO_MEMORY_TYPE = "bram";\n\n');
  outputFile.write('xpm_fifo_sync packet_fifo (\n')
  outputFile.write('\t.wr_en(packet_fifo_wr_en),\n')
  outputFile.write('\t.din(packet_fifo_din),\n')
  outputFile.write('\t.rd_en(packet_fifo_rd_en),\n')
  outputFile.write('\t.sleep(1\'b0),\n')
  outputFile.write('\t.injectsbiterr(),\n')
  outputFile.write('\t.injectdbiterr(),\n')
  outputFile.write('\t.prog_empty(),\n')
  outputFile.write('\t.dout(packet_fifo_dout), \n')
  outputFile.write('\t.empty(packet_fifo_empty), \n')
  outputFile.write('\t.prog_full(packet_fifo_almost_full),\n')
  outputFile.write('\t.full(packet_fifo_full),\n')
  outputFile.write('\t.rd_data_count(),\n')
  outputFile.write('\t.wr_data_count(),\n') 
  outputFile.write('\t.wr_rst_busy(),\n')
  outputFile.write('\t.rd_rst_busy(),\n') 
  outputFile.write('\t.overflow(),\n')
  outputFile.write('\t.underflow(),\n')
  outputFile.write('\t.sbiterr(),\n')
  outputFile.write('\t.dbiterr(),\n') 
  outputFile.write('\t.wr_clk(clk_line),\n')
  outputFile.write('\t.rst(rst)\n')
  outputFile.write(');\n\n')

  outputFile.write('assign packet_fifo_wr_en = packet_in_packet_in_VAL;\n')
  outputFile.write('assign packet_fifo_din = {packet_in_packet_in_SOF, packet_in_packet_in_EOF, packet_in_packet_in_DAT,\n')
  outputFile.write('                          packet_in_packet_in_CNT, packet_in_packet_in_ERR};\n')
  outputFile.write('assign packet_fifo_rd_en = ' + parseInfo['function'] + '_packet_input_v_read;\n')

def writeLogic(parseInfo, outputFile):

  function = parseInfo['function']

  outputFile.write('assign packet_in_packet_in_RDY = 1;\n\n')

  outputFile.write('assign {packet_out_packet_out_SOF, packet_out_packet_out_EOF, packet_out_packet_out_DAT,\n')
  outputFile.write('        packet_out_packet_out_CNT, packet_out_packet_out_ERR} = ' + function + '_packet_output_v;\n')
  outputFile.write('assign packet_out_packet_out_VAL = ' + function + '_packet_output_v_ap_vld & ~backpressure_in;\n\n')

  start = 0
  end = 0
  for tup, info in parseInfo['output_tuples'].items():
    end += info['width']
    outputFile.write('assign ' + tup + '_DATA  = ' + function + '_tuple_output_v[' + str(end - 1) + ':' + str(start) + '];\n')
    start = end
  outputFile.write('\n')

  for tup in parseInfo['output_tuples']:
    outputFile.write('assign ' + tup + '_VALID  = ' + function + '_tuple_output_v_ap_vld;\n')
  outputFile.write('\n')

  outputFile.write('assign backpressure_out = packet_fifo_almost_full;\n\n')

def generateStub(parseInfo, hlsFilename, headerFilename):
  with open(hlsFilename, "wt") as outputFile:
    outputFile.write('// Code generated by the Fuselage Factory\n\n')
    outputFile.write('#include <' + os.path.basename(headerFilename) + '>\n\n')
    outputFile.write('#include <hls_stream.h>\n\n')
    outputFile.write('void ' + parseInfo['function'] + '(hls::stream<input_tuples> & Tuple_input, hls::stream<output_tuples> & Tuple_output,\n')
    outputFile.write('    hls::stream<packet_interface> & Packet_input, hls::stream<packet_interface> & Packet_output)\n')
    outputFile.write('{\n')
    outputFile.write('#pragma HLS DATA_PACK variable=Tuple_input\n')
    outputFile.write('#pragma HLS DATA_PACK variable=Tuple_output\n')
    outputFile.write('#pragma HLS DATA_PACK variable=Packet_input\n')
    outputFile.write('#pragma HLS DATA_PACK variable=Packet_output\n')
    outputFile.write('#pragma HLS INTERFACE ap_fifo port=Tuple_input\n')
    outputFile.write('#pragma HLS INTERFACE ap_hs port=Tuple_output\n')
    outputFile.write('#pragma HLS INTERFACE ap_fifo port=Packet_input\n')
    outputFile.write('#pragma HLS INTERFACE ap_hs port=Packet_output\n\n')
    outputFile.write('  /* Fill in your code here. */\n')
    outputFile.write('}\n')

def generateHeader(parseInfo, filename):
  with open(filename, "wt") as outputFile:
    outputFile.write('// Code generated by the Fuselage Factory\n\n')

    outputFile.write('// WE DISCOURAGE MODIFYING THIS FILE DIRECTLY!!!\n')
    outputFile.write('// Changes to P4 code that do not affect the external function signature in P4\n')
    outputFile.write('// may still result in different variable widths.\n\n')

    outputFile.write('#ifndef FUSELAGE_HEADER_INCLUDED\n')
    outputFile.write('#define FUSELAGE_HEADER_INCLUDED\n\n')

    outputFile.write('#include <cstddef>\n')
    outputFile.write('#include <ap_int.h>\n')
    outputFile.write('#include <hls_stream.h>\n\n')

    tuples = set()
    for direction in ['in', 'out']:
      for tup, info in parseInfo[direction + 'put_tuples'].items():
        tuple_name = re.match(r'tuple_.*?_(.*)', tup).group(1)
        if tuple_name not in tuples:
          outputFile.write('struct tuple_' + tuple_name + '\n')
          outputFile.write('{\n')
          fields = info.get('fields')
          if fields != None:
            offset = 0
            for name, width in reversed(fields.items()):
              outputFile.write('  ap_uint<' + str(width) + '> ' + name.capitalize() + ';\n')
          else:
            outputFile.write('  ap_uint<' + str(info['width']) + '> ' + tuple_name.capitalize() + ';\n')
          outputFile.write('};\n\n')
          tuples.add(tuple_name)

    for direction in ['in', 'out']:
      outputFile.write('struct ' + direction + 'put_tuples\n')
      outputFile.write('{\n')
      for tup in parseInfo[direction + 'put_tuples']:
        tuple_name = re.match(r'tuple_.*?_(.*)', tup).group(1)
        outputFile.write('  tuple_' + tuple_name + ' ' + tuple_name.capitalize() + ';\n')
      outputFile.write('};\n\n')

    outputFile.write('struct packet_interface\n')
    outputFile.write('{\n')
    outputFile.write('  ap_uint<1> Error;\n')
    outputFile.write('  ap_uint<4> Count;\n')
    outputFile.write('  ap_uint<64> Data;\n')
    outputFile.write('  ap_uint<1> End_of_frame;\n')
    outputFile.write('  ap_uint<1> Start_of_frame;\n')
    outputFile.write('};\n\n')

    outputFile.write('void ' + parseInfo['function'] + '(hls::stream<input_tuples> & Tuple_input, hls::stream<output_tuples> & Tuple_output,\n')
    outputFile.write('    hls::stream<packet_interface> & Packet_input, hls::stream<packet_interface> & Packet_output);\n\n')

    outputFile.write('#endif\n')

if __name__ == '__main__':
  main()

