#
# Copyright 2010-2011,2015 Ettus Research LLC
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

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 <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
    ${reg.get_type()} ${reg.get_name()};
    % endfor

    ${name}_t(void){
        _state = NULL;
        % for reg in regs:
        ${reg.get_name()} = ${reg.get_default()};
        % 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(_state->${reg.get_name()} != this->${reg.get_name()}){
            addrs.insert(${reg.get_addr()});
        }
        % 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 */
"""

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

def to_num(arg): return int(eval(arg))

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

    def parse(self, reg_des):
        x = re.match('^(\w*)\s*(\w*)\[(.*)\]\s*(\w*)\s*(.*)$', reg_des)
        name, 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

    def get_addr(self): 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

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='', file=__file__, append=False):
    #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_tmpl, regs=regs).splitlines())

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

    #write the generated code to file specified by argv1
    open(sys.argv[1], 'a' if append else 'w').write(code)