from astvisitors import *

MODEL = """
------------------------------------------------------------------------
-- Generated using Kripkomat
------------------------------------------------------------------------
MODULE {name}
VAR
    state : {{ {states} }};
{variables}
ASSIGN
    init(state) := {initial};
    next(state) :=
        case
{transitions}
            TRUE : state;
        esac;
{varassigns}
DEFINE
{properties}
"""

VARIABLE = "    {n}:{t};"

TRAN = "            (state = {n}) & ({cond}) : {m};"
CASE = "            {cond} : {value};"
VARASSIGN = """
    init({name}) := {initial};
    next({name}) :=
        case
{cases}
            TRUE : {name};
        esac;
"""

PROP = """    -- Property "{prop}"
    {alias} := {logic};
"""
PROP_LOGIC = "state = {state}"

LTL_FORM = """LTLSPEC
    {ltl};
"""

class ModelBuilder:
    def __init__(self, states, initial, transitions, assigns, ltls=[], name="main"):
        self._name = name
        self._states = states
        self._initial = initial
        self._tran = transitions
        self._assigns = assigns
        self._ltls = ltls

    def generate(self):
        cg = CGenerator()
        cond_tf = NuSmvConditionTransformer()

        # build model
        states_decl = ",".join(self._states)
        transitions = []
        
        # find variables in the condition
        varextract = NuSmvVariableExtractor()
        
        for tran in self._tran:
            cond = cond_tf.transform(tran.condition)
            if cond is not None:
                expr = cg.visit(cond)
            else:
                expr = 'TRUE'
            transition = TRAN.format(n=tran.state_from, m=tran.state_to, cond=expr)
            transitions.append(transition)
            if tran.condition is not None:
                varextract.visit(tran.condition)
        varassigns = []
        for var,assigns in self._assigns.items():
            cases = []
            var_label = var
            vartype = None
            for a in assigns:
                assign_varextract = NuSmvVariableExtractor()
                #todo convert rvalue to type
                assign_varextract.visit(a.lvalue)
                varname, vartype = list(assign_varextract.variables.items())[0]
                var_label = varname
                cond = cond_tf.transform(a.condition)
                cond_expr = cg.visit(cond)
                varextract.visit(a.condition)
                varextract.visit(a.lvalue)
                varextract.visit(a.rvalue)
                lvalue = var
                rvalue = cg.visit(cond_tf.transform(a.rvalue))
                if vartype == 'boolean':
                    if rvalue == '0':
                        rvalue = 'FALSE'
                    elif rvalue == '1':
                        rvalue = 'TRUE'
                    else:
                        pass
                cases.append(CASE.format(cond=cond_expr, value=rvalue))
                initial = '0'
                if vartype == 'boolean':
                    initial = 'FALSE'
            varassigns.append(VARASSIGN.format(name=var_label, cases="\n".join(cases),
                                               initial=initial))
            #print(cases)

        variables = []
        for v,t in varextract.variables.items():
            variables.append(VARIABLE.format(n=v, t=t))
        
        properties = []
        #for alias,(states,prop) in self._props.items():
        #    logic = " | ".join([PROP_LOGIC.format(state=x) for x in states])
        #    prop_str = PROP.format(prop=prop, alias=alias, logic=logic)
        #    properties.append(prop_str)

        out = MODEL.format(name=self._name,
                            states=states_decl,
                            variables="\n".join(variables),
                            initial=self._initial,
                            transitions="\n".join(transitions),
                            varassigns="\n".join(varassigns),
                            properties="\n".join(properties))

        # add LTL formulas
        for ltl in self._ltls:
            out += LTL_FORM.format(ltl=ltl)
        return out