#!/usr/bin/env python3

# Copyright (c) 2019-2021 Antmicro <www.antmicro.com>
# Copyright (c) 2021 Henk Vergonet <henk.vergonet@gmail.com>
#
# Zephyr DTS & config overlay generator for LiteX SoC.
#
# This script parses LiteX 'csr.json' file and generates DTS and config
# files overlay for Zephyr.

# Changelog:
# - 2021-07-05 Henk Vergonet <henk.vergonet@gmail.com>
#    removed dependency on intermediate interpretation layers
#    switch to JSON csr
#    fix uart size parameter
# - 2021-07-15 Henk Vergonet <henk.vergonet@gmail.com>
#    added identifier_mem handler as dna0
#    added spiflash as spi0
#

import argparse
import json


# DTS formatting
def dts_open(name, parm): return "&{} {{\n".format(parm.get('alias', name))
def dts_close():          return "};\n"
def dts_intr(name, csr):  return "    interrupts = <{} 0>;\n".format(
                                    hex(csr['constants'][name + '_interrupt']))
def dts_reg(regs):        return "    reg = <{}>;\n".format(regs)


# DTS handlers
def disabled_handler(name, parm, csr):
    return "    status = \"disabled\";\n"


def ram_handler(name, parm, csr):
    return dts_reg(" ".join([
                hex(csr['memories'][name]['base']),
                hex(csr['memories'][name]['size'])]))


def ethmac_handler(name, parm, csr):
    dtsi  = dts_reg(" ".join([
                hex(csr['csr_bases'][name]),
                hex(parm['size']),
                hex(csr['memories'][name]['base']),
                hex(csr['memories'][name]['size'])]))
    dtsi += dts_intr(name, csr)
    return dtsi


def i2c_handler(name, parm, csr):
    dtsi  = dts_reg(" ".join([
                hex(csr['csr_bases'][name]),
                hex(parm['size']),
                hex(csr['csr_bases'][name] + parm['size']),
                hex(parm['size'])]))
    dtsi += dts_intr(name, csr)
    return dtsi


def peripheral_handler(name, parm, csr):
    dtsi  = dts_reg(" ".join([
                hex(csr['csr_bases'][name]),
                hex(parm['size'])]))
    try:
        dtsi += dts_intr(name, csr)
    except KeyError as e:
        print('  dtsi key', e, 'not found, no interrupt override')
    return dtsi


overlay_handlers = {
    'uart': {
        'handler': peripheral_handler,
        'alias': 'uart0',
        'size': 0x20,
        'config_entry': 'UART_LITEUART'
    },
    'timer0': {
        'handler': peripheral_handler,
        'size': 0x40,
        'config_entry': 'LITEX_TIMER'
    },
    'ethmac': {
        'handler': ethmac_handler,
        'alias': 'eth0',
        'size': 0x80,
        'config_entry': 'ETH_LITEETH'
    },
    'spiflash': {
        'handler': peripheral_handler,
        'alias': 'spi0',
        'size': 12,
        'config_entry': 'SPI_LITESPI'
    },
    'i2c0' : {
        'handler': i2c_handler,
        'size': 0x4,
        'config_entry': 'I2C_LITEX'
    },
    'main_ram': {
        'handler': ram_handler,
        'alias': 'ram0',
    },
    'identifier_mem': {
        'handler': peripheral_handler,
        'alias': 'dna0',
        'size': 0x100,
    }
}


def generate_dts_config(csr):
    dts = cnf = ''

    for name, parm in overlay_handlers.items():
        print('Generating overlay for:',name)
        enable = 'y'
        dtsi = dts_open(name, parm)

        try:
            dtsi += parm['handler'](name, parm, csr)
        except KeyError as e:
            print('  dtsi key', e, 'not found, disable', name)
            enable = 'n'
            dtsi += disabled_handler(name, parm, csr)

        dtsi += dts_close()
        dts += dtsi
        if 'config_entry' in parm:
            cnf += ' -DCONFIG_' + parm['config_entry'] + '=' + enable 

    for name, value in csr['csr_bases'].items():
        if name not in overlay_handlers.keys():
            print('No overlay handler for:', name, 'at', hex(value))

    return dts, cnf


# helpers
def print_or_save(filepath, lines):
    """ Prints given string on standard output or to the file.

    Args:
        filepath (string): path to the file lines should be written to
                           or '-' to write to a standard output
        lines (string): content to be printed/written
    """
    if filepath == '-':
        print(lines)
    else:
        with open(filepath, 'w') as f:
            f.write(lines)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('conf_file',
                        help='JSON configuration generated by LiteX')
    parser.add_argument('--dts', action='store', required=True,
                        help='Output DTS overlay file')
    parser.add_argument('--config', action='store', required=True,
                        help='Output config overlay file')
    return parser.parse_args()


def main():
    args = parse_args()

    with open(args.conf_file) as f:
        csr = json.load(f)
    dts, config = generate_dts_config(csr)

    print_or_save(args.dts, dts)
    print_or_save(args.config, config)


if __name__ == '__main__':
    main()
