#!/usr/bin/python3
"""
Generate lvbitx from template and bitfile
"""

from __future__ import print_function
from xml.etree import ElementTree
import optparse
import base64
import os
from hashlib import md5

def to_native_str(str_or_bstr):
    """
    Returns a native string, regardless of the input string type (binary or
    UTF-8), and the Python version (2 or 3).
    Note that the native string type is actually not the same in Python 2 and
    3: In the former, it's a binary string, in the latter, it's Unicode.
    >>> to_native_str(b'foo')
    'foo'
    >>> to_native_str(u'foo')
    'foo'
    """
    if isinstance(str_or_bstr, str):
        return str_or_bstr
    try:
        # This will either fail because we're running Python 2 (which doesn't)
        # have the encoding argument) or because we're not passing in a bytes-
        # like object (e.g., an integer)
        return str(str_or_bstr, encoding='ascii')
    except TypeError:
        return str(str_or_bstr)

def get_parser():
    """Parse args."""
    parser = optparse.OptionParser()
    parser.add_option(
        "--device",
        type="string",
        dest="device_type",
        help="Device Type. (Has to match the LVFPGA target plugin)",
        default=None
    )
    parser.add_option(
        "--input-bin",
        type="string",
        dest="input_bin",
        help="Path to bin file that needs to be merged with the LVBITX before exporting",
        default=None
    )
    parser.add_option(
        "--output-bin",
        type="string",
        dest="output_bin",
        help="Create a binary configuration bitstream",
        default=None
    )
    parser.add_option(
        "--output-lvbitx",
        type="string",
        dest="output_lvbitx_path",
        help="Output path for autogenerated LVBITX file",
        default=None
    )
    parser.add_option(
        "--output-src-path",
        type="string",
        dest="output_src_path",
        help="Output path for autogenerated src file",
        default=None
    )
    return parser

def main():
    """Go, go, go!"""
    parser = get_parser()
    options, args = parser.parse_args()
    # Args
    if len(args) < 1:
        print('ERROR: Please specify the input LVBITX file name')
        parser.print_help()
        exit(1)

    lvbitx_filename = args[0]
    input_filename = os.path.abspath(lvbitx_filename)

    if not os.path.isfile(input_filename):
        print("ERROR: LVBITX File `{}' could not be accessed or is not a file."
              .format(input_filename))
        parser.print_help()
        exit(1)

    if options.input_bin is not None and \
            not os.path.isfile(os.path.abspath(options.input_bin)):
        print("ERROR: FPGA Bin File `{}' could not be accessed or is not a file."
              .format(options.input_bin))
        parser.print_help()
        exit(1)

    if options.output_lvbitx_path is not None and \
            input_filename == options.output_lvbitx_path:
        print('ERROR: Input and output LVBITX files were the same. '
              'Choose a difference input or output file.')
        parser.print_help()
        exit(1)

    # Get XML Tree Node
    tree = ElementTree.parse(input_filename)
    root = tree.getroot()

    # Update device type
    if options.device_type is not None:
        root.find('Project').find('TargetClass').text += '; ' + options.device_type

    # Merge bitstream into LVBITX
    if options.input_bin is not None:
        with open(os.path.abspath(options.input_bin), 'rb') as bin_file:
            bitstream = bin_file.read()
            try:
                # allow the md5 call to work on FIPS systems
                # try python3.9+ keyword, which is also supported on some earlier versions
                bitstream_md5 = md5(bitstream, usedforsecurity=False).hexdigest()
            except TypeError:
                # unknown keywords trigger TypeError so default back to basic call
                bitstream_md5 = md5(bitstream).hexdigest()
            bitstream_b64 = base64.b64encode(bitstream)
            bitstream_b64_lb = b'\n'.join([
                bitstream_b64[i:i+76]
                for i in range(0, len(bitstream_b64), 76)
            ]) + b'\n'
            root.find('Bitstream').text = to_native_str(bitstream_b64_lb)
            root.find('BitstreamMD5').text = bitstream_md5

    # Write BIN file
    bitstream = base64.b64decode(root.find('Bitstream').text)

    try:
        # allow the md5 call to work on FIPS systems
        # try python3.9+ keyword, which is also supported on some earlier versions
        bitstream_md5 = md5(bitstream, usedforsecurity=False).hexdigest()
    except TypeError:
        # unknown keywords trigger TypeError so default back to basic call
        bitstream_md5 = md5(bitstream).hexdigest()
    
    if options.output_lvbitx_path is not None \
            and bitstream_md5 != root.find('BitstreamMD5').text:
        print('ERROR: The MD5 sum for the output LVBITX was incorrect. '
              'Make sure that the bitstream in the input LVBITX or BIN file is valid.')
        exit(1)

    if options.output_bin is not None:
        fpga_bin_file = open(options.output_bin, 'wb')
        fpga_bin_file.write(bitstream)
        fpga_bin_file.close()

    # Save LVBITX
    if options.output_lvbitx_path is not None:
        with open(options.output_lvbitx_path, 'wb') as lvbitx_file:
            tree.write(
                lvbitx_file,
                encoding="utf-8",
                xml_declaration=True,
                default_namespace=None,
                method="xml"
            )

if __name__ == "__main__":
    main()
