#
# Copyright 2010-2011,2015 Ettus Research LLC
# Copyright 2018 Ettus Research, a National Instruments Company
#
# SPDX-License-Identifier: GPL-3.0-or-later
#

import re
import sys
import math
from mako.template import Template

COMMON_TMPL = """<% import time %>\
/***********************************************************************
 * This file was generated by ${file} on ${time.strftime("%c")}
 **********************************************************************/

#ifndef INCLUDED_${name.upper()}_HPP
#define INCLUDED_${name.upper()}_HPP

#include <uhd/config.hpp>
#include <uhd/exception.hpp>
#include <set>
#include <vector>
#include <stdint.h>

class ${name}_t{
public:
    % for reg in regs:
        % if reg.get_enums():
    enum ${reg.get_type()}{
            % for i,enum in enumerate(reg.get_enums()):
        ${reg.get_name().upper()}_${enum[0].upper()} = ${enum[1]}<% comma = ',' if i != (len(reg.get_enums())-1) else '' %>${comma}
            % endfor
    };
        % endif
        % if reg.is_array:
    std::vector<${reg.get_type()}> ${reg.get_name()};
        % else:
    ${reg.get_type()} ${reg.get_name()};
        % endif
    % endfor

    ${name}_t(void){
        _state = NULL;
        % for reg in regs:
            % if reg.is_array:
        ${reg.get_name()}.resize(${reg.get_array_len()}, ${reg.get_default()});
            % else:
        ${reg.get_name()} = ${reg.get_default()};
            % endif
        % endfor
    }

    ~${name}_t(void){
        delete _state;
    }

    ${body}

    void save_state(void){
        if (_state == NULL) _state = new ${name}_t();
        % for reg in regs:
        _state->${reg.get_name()} = this->${reg.get_name()};
        % endfor
    }

    template<typename T> std::set<T> get_changed_addrs(void){
        if (_state == NULL) throw uhd::runtime_error("no saved state");
        //check each register for changes
        std::set<T> addrs;
        % for reg in regs:
            % if reg.is_array:
        for (size_t i = 0; i < ${reg.get_array_len()}; i++) {
            if(_state->${reg.get_name()}[i] != this->${reg.get_name()}[i]) {
                addrs.insert(${reg.get_addr()} + i * ${reg.get_addr_step_size()});
            }
        }
            % else:
        if(_state->${reg.get_name()} != this->${reg.get_name()}){
            addrs.insert(${reg.get_addr()});
        }
            % endif
        % endfor
        return addrs;
    }

    % for mreg in mregs:
    ${mreg.get_type()} get_${mreg.get_name()}(void){
        return <% shift = 0 %>
        % for reg in mreg.get_regs():
        (${mreg.get_type()}(${reg.get_name()} & ${reg.get_mask()}) << ${shift}) |<% shift = shift + reg.get_bit_width() %>
        % endfor
        0;
    }

    void set_${mreg.get_name()}(${mreg.get_type()} reg){<% shift = 0 %>
        % for reg in mreg.get_regs():
        ${reg.get_name()} = (reg >> ${shift}) & ${reg.get_mask()};<% shift = shift + reg.get_bit_width() %>
        % endfor
    }

    % endfor
private:
    ${name}_t *_state;
};

#endif /* INCLUDED_${name.upper()}_HPP */
"""

# This template matches the functionality of COMMON_TMPL in Python syntax
COMMON_PY_TMPL = """<% import time %>\
##########################################################################
# This file was generated by ${file} on ${time.strftime("%c")}
<%text>##########################################################################</%text>

from enum import Enum

## Create a class for the register map
class ${name}_t:
    ## Create an enum for each register which has defined values
    % for reg in regs:
    % if reg.get_enums():
    class ${reg.get_type()} (Enum):
        % for i, enum in enumerate(reg.get_enums()):
        ${reg.get_name().upper()}_${enum[0].upper()} = ${enum[1]}
        % endfor
    % endif
    %endfor

    def __init__(self):
        ## Assign each register to its default value
        self._state = None
        % for reg in regs:
            % if reg.get_enums():
                % if reg.is_array:
        self.${reg.get_name()} = [self.${reg.get_name()}_t.${reg.get_default()},] * ${reg.get_array_len()}
                % else:
        self.${reg.get_name()} = self.${reg.get_name()}_t.${reg.get_default()}
                % endif
            % else:
                % if reg.is_array:
        self.${reg.get_name()} = [${reg.get_default()},] * ${reg.get_array_len()}
                % else:
        self.${reg.get_name()} = ${reg.get_default()}
                % endif
            % endif
        self.${reg.get_name()}_addr = ${reg.get_addr()}
        self.${reg.get_name()}_mask = ${reg.get_mask()}
        self.${reg.get_name()}_shift = ${reg.get_shift()}
        % endfor

    ${body}

    def save_state(self):
        if self._state is None:
            self._state = ${name}_t()
        % for reg in regs:
            % if reg.is_array:
        self._state.${reg.get_name()} = self.${reg.get_name()}.copy()
            % else:
        self._state.${reg.get_name()} = self.${reg.get_name()}
            % endif
        % endfor

    def get_changed_addrs(self):
        if self._state is None:
            raise RuntimeError("No saved state")
        #check each register for changes
        addrs = set()
        % for reg in regs:
        % if reg.is_array:
        for index, value in enumerate(self.${reg.get_name()}):
            if self._state.${reg.get_name()}[index] != value:
                addrs.add(${reg.get_addr()} + index * ${reg.get_addr_step_size()})
        % else:
        if self._state.${reg.get_name()} != self.${reg.get_name()}:
            addrs.add(${reg.get_addr()})
        % endif
        % endfor
        return addrs

    % for mreg in mregs:
    def get_${mreg.get_name()}(self):
        return ( <% shift = 0 %>
            % for reg in mreg.get_regs():
            % if reg.get_enums():
            (self.${reg.get_name()}.value & ${reg.get_mask()} << ${shift}) |<% shift = shift + reg.get_bit_width() %>
            % else:
            (self.${reg.get_name()} & ${reg.get_mask()} << ${shift}) |<% shift = shift + reg.get_bit_width() %>
            % endif
            % endfor
            0
        )

    def set_${mreg.get_name()}(self, reg): <% shift = 0 %>
        % for reg in mreg.get_regs():
        self.${reg.get_name()} = (reg >> ${shift}) & ${reg.get_mask()}<% shift = shift + reg.get_bit_width() %>
        % endfor

    %endfor
"""

def parse_tmpl(_tmpl_text, **kwargs):
    return Template(_tmpl_text).render(**kwargs)

def to_num(arg):
    """
    Return the integer representation of arg, which is usually a string, but can
    be a Python representation
    """
    return int(eval(arg))

class reg:
    def __init__(self, reg_des):
        self.is_array = False
        self._array_len = 0
        self._addr_step = 0
        try:
            self.parse(reg_des)
        except Exception as e:
            raise Exception(
                'Error parsing register description: "{}"\nWhat: {}'
                .format(reg_des, e))

    def parse(self, reg_des):
        """
        Parse a regmap line.

        Unpacks lines like this:

        REG_NAME[ARRAY_LEN]  BASE_ADDRESS[BITS] DEFAULT LIST_OF_ENUMS

        Notes:
        - Array length and enums are optional
        - Base address == address if not an array

        Example:

        super_reg       0x1000[0:1]  0
        duper_reg[128]  0x2000[0:31] 0

        super_reg is a 2-bit field in register at address 0x1000. It will be
        initialized to 0. duper_reg is an array of length 128, of 32-bit registers.
        """
        x = re.match(
            r'^(\w*)(\[([0-9:]*)\])?\s*(\w*)\[(.*)\]\s*(\w*)\s*(.*)$',
            reg_des)
        name, _, addr_range, addr, bit_range, default, enums = x.groups()
        #store variables
        self._name = name
        self._addr = to_num(addr)
        if ':' in bit_range: self._addr_spec = sorted(map(int, bit_range.split(':')))
        else: self._addr_spec = int(bit_range), int(bit_range)
        self._default = to_num(default)
        #extract enum
        self._enums = list()
        if enums:
            enum_val = 0
            for enum_str_unstripped in enums.split(','):
                enum_str = enum_str_unstripped.strip()
                if '=' in enum_str:
                    enum_name, enum_val = enum_str.split('=')
                    enum_val = to_num(enum_val)
                else: enum_name = enum_str
                self._enums.append((enum_name, enum_val))
                enum_val += 1
        # Extract address range for arrays
        if addr_range is not None:
            self.is_array = True
            addr_range = addr_range.split(':', 3)
            if len(addr_range) == 1:
                self._array_len = to_num(addr_range[0])
                self._addr_step = 4
            else:
                addr_start = to_num(addr_range[0])
                addr_end = to_num(addr_range[1])
                self._array_len = addr_end - addr_start + 1
                self._addr_step = 4 if len(addr_range) == 2 else to_num(addr_range[2])
                if addr_start > 0:
                    self._addr += addr_start

    def get_addr(self):
        """
        Return register address. If it's an array, then return the base address.
        """
        return self._addr

    def get_enums(self): return self._enums
    def get_name(self): return self._name
    def get_default(self):
        for key, val in self.get_enums():
            if val == self._default: return ('%s_%s'%(self.get_name(), key)).upper()
        return self._default
    def get_type(self):
        if self.get_enums(): return '%s_t'%self.get_name()
        return 'uint%d_t'%max(2**math.ceil(math.log(self.get_bit_width(), 2)), 8)
    def get_shift(self): return self._addr_spec[0]
    def get_mask(self): return hex(int('1'*self.get_bit_width(), 2))
    def get_bit_width(self): return self._addr_spec[1] - self._addr_spec[0] + 1

    def get_array_len(self):
        """
        Return 1 for regular registers, or the array length for arrays
        """
        return self._array_len

    def get_addr_step_size(self):
        """
        Return the step size for register arrays
        """
        return self._addr_step

class mreg:
    def __init__(self, mreg_des, regs):
        try: self.parse(mreg_des, regs)
        except Exception as e:
            raise Exception('Error parsing meta register description: "%s"\nWhat: %s'%(mreg_des, e))

    def parse(self, mreg_des, regs):
        x = re.match('^~(\w*)\s+(.*)\s*$', mreg_des)
        self._name, reg_names = x.groups()
        regs_dict = dict([(reg.get_name(), reg) for reg in regs])
        self._regs = [regs_dict[reg_name.strip()] for reg_name in reg_names.split(',')]

    def get_name(self): return self._name
    def get_regs(self): return self._regs
    def get_bit_width(self): return sum(map(reg.get_bit_width, self._regs))
    def get_type(self):
        return 'uint%d_t'%max(2**math.ceil(math.log(self.get_bit_width(), 2)), 8)

def generate(name, regs_tmpl, body_tmpl='', py_body_tmpl='', file=__file__, append=False):
    # determine if the destination file is a Python or C++ header file
    out_file = sys.argv[1]
    if out_file.endswith('.py'): # Write a Python file
        template = COMMON_PY_TMPL
        body_template = py_body_tmpl
    else: # default to C++ Header
        template = COMMON_TMPL
        body_template = body_tmpl

    #evaluate the regs template and parse each line into a register
    regs = list()
    mregs = list()
    for entry in parse_tmpl(regs_tmpl).splitlines():
        if entry.startswith('~'):
            mregs.append(mreg(entry, regs))
        else:
            regs.append(reg(entry))

    #evaluate the body template with the list of registers
    body = '\n    '.join(parse_tmpl(body_template, regs=regs).splitlines())

    #evaluate the code template with the parsed registers and arguments
    code = parse_tmpl(template,
        name=name,
        regs=regs,
        mregs=mregs,
        body=body,
        file=file,
    )

    #write the generated code to file specified by argv1
    with open(out_file, 'a' if append else 'w') as f:
        f.write(code)