from __future__ import print_function
import argparse
import sys
from modelbuilder import *

from pycparser import parse_file, c_ast

func_table = {}
enum_table = {}
assign_table = []
state_enums = []

class FuncDefVisitor(c_ast.NodeVisitor):
    def __init__(self):
        self.func_table = {}
    def visit_FuncDef(self, node):
        self.func_table[node.decl.name] = node

class StateAssignmentVisitor(c_ast.NodeVisitor):
    def __init__(self, state):
        super().__init__()
        self.state = state
        self.assignments = []

    def visit(self, node, path = None):
        """ Visit a node.
        """
        if path is None:
            path = []
        if self._method_cache is None:
            self._method_cache = {}

        visitor = self._method_cache.get(node.__class__.__name__, None)
        if visitor is None:
            method = 'visit_' + node.__class__.__name__
            visitor = getattr(self, method, self.generic_visit)
            self._method_cache[node.__class__.__name__] = visitor

        return visitor(node, path)

    def generic_visit(self, node, path):
        """ Called if no explicit visitor function exists for a
            node. Implements preorder visiting of the node.
        """
        path = path.copy()
        path.append(node)
        for c in node:
            self.visit(c, path)

    def visit_Assignment(self, n, path):
        rval_str = n.rvalue.name
        if rval_str == self.state:
            self.assignments.append((n,path))

class EnumDefVisitor(c_ast.NodeVisitor):
    def __init__(self, name):
        super().__init__()
        self._name = name

    def visit_Enum(self, node):
        enum_table[self._name] = node

class TypedefVisitor(c_ast.NodeVisitor):
    def visit_Typedef(self, node):
        ev = EnumDefVisitor(node.name)
        ev.visit(node)


class EnumVisitor(c_ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.enum_names = []

    def visit_Enumerator(self, node):
        self.enum_names.append(node.name)

def find_longest_path(paths):
    ms = 0
    mp = None
    for path in paths:
        s = len(path)
        if s > ms:
            ms = s
            mp = path
    return mp

def find_shortest_path(paths):
    ms = float('inf')
    mp = None
    for path in paths:
        s = len(path)
        if s < ms:
            ms = s
            mp = path
    return mp

def find_common_ancestor(paths):
    shortest_path = find_shortest_path(paths)
    
    last_common = None
    for gen_i,_ in enumerate(shortest_path):
        common = True
        last_node = paths[0][gen_i]

        for path in paths:
            if last_node == path[gen_i]:
                common = True
                last_node = path[gen_i]
            else:
                common = False
        if common:
            last_common = last_node
    return last_common
        

class SwitchCaseTermVisitor(c_ast.NodeVisitor):
    def __init__(self, asm_node):
        super().__init__()
        self._asm_node = asm_node
        self.hit = False

    def visit_Assignment(self, node):
        if node == self._asm_node:
            self.hit = True

class SwitchCaseTranVisitor(c_ast.NodeVisitor):
    def __init__(self, asm_node):
        super().__init__()
        self._asm_node = asm_node
        self.tran_table = []

    def visit_Case(self, node):
        sctv = SwitchCaseTermVisitor(self._asm_node)
        sctv.visit(node)
        if sctv.hit:
            self.tran_table.append((node.children()[0][1].name, self._asm_node.rvalue.name))

class SwitchCasePropertyVisitor(c_ast.NodeVisitor):
    def __init__(self, state_asmts):
        super().__init__()
        self._sas = state_asmts
        self.properties = []

    def visit_Assignment(self, node):
        if not(node in self._sas):
            prop = None
            if isinstance(node.lvalue, c_ast.StructRef):
                prop = f"{node.lvalue.children()[0][1].name}->{node.lvalue.children()[1][1].name}<={node.rvalue.name}";
            else:
                prop = f"{node.lvalue.name}<={node.rvalue.name}"
            self.properties.append(prop)


class SwitchCaseCodePropertyVisitor(c_ast.NodeVisitor):
    def __init__(self, case, state_asmts):
        super().__init__()
        self._case = case
        self._sas = state_asmts
        self.properties = []
        

    def visit_Case(self, node):
        label = node.children()[0][1]
        block = node
        if label.name == self._case:
            scpv = SwitchCasePropertyVisitor(self._sas)
            scpv.visit(block)
            self.properties += scpv.properties

if __name__ == "__main__":
    argparser = argparse.ArgumentParser('Create a Kripke Structure Model from C Code')
    argparser.add_argument('filename',
                           default='examples/c_files/basic.c',
                           nargs='?',
                           help='name of file to parse')
    argparser.add_argument('--coord', help='show coordinates in the dump',
                           action='store_true')

    argparser.add_argument('--func', help='process function')
    argparser.add_argument('--enum', help='state enum')
    argparser.add_argument('--initial', help='initial state')
    argparser.add_argument('--ltlfile', help='file containing LTL formulae')
    argparser.add_argument('-o', '--output', dest='output', help='output NuSMV file')
    args = argparser.parse_args()

    ast = parse_file(args.filename, use_cpp=False)

    fdv = FuncDefVisitor()
    fdv.visit(ast)
    func_table = fdv.func_table

    enumvis = TypedefVisitor()
    enumvis.visit(ast)
    
    state_asmts = []
    
    if args.enum in enum_table:
        ev = EnumVisitor()
        ev.visit(enum_table[args.enum])
        for ename in ev.enum_names:
            sav = StateAssignmentVisitor(ename)
            sav.visit(func_table[args.func])
            state_asmts += sav.assignments
    else:
        print(f"Initial State Enum '{args.enum}' not found")

    paths = []
    for asm in state_asmts:
        paths.append(asm[1])

    common = find_common_ancestor(paths)

    tran_table = []

    for sa in state_asmts:
        sctv = SwitchCaseTranVisitor(sa[0])
        sctv.visit(common)
        tran_table += sctv.tran_table

    comp_tt = {}
    print("Transitions")
    for t in tran_table:
        print(f"{t[0]}->{t[1]}")
        if t[0] in comp_tt:
            comp_tt[t[0]].append(t[1])
        else:
            comp_tt[t[0]] = [t[1]]
    print("")

    pure_sa = [x[0] for x in state_asmts]
    states = comp_tt.keys()
    print("States: ", ",".join(states))
    print("")

    props_by_state = {}
    print("Compact Transition Table:")
    for n,ms in comp_tt.items():
        sstr = ','.join(ms)
        print(f"{n}->{{{sstr}}}")

        # find properties
        sccpv = SwitchCaseCodePropertyVisitor(n, pure_sa)
        sccpv.visit(common)
        if len(sccpv.properties) > 0:
            props_by_state[n] = sccpv.properties
    print("")

    properties = {}
    for state,props in props_by_state.items():
        for prop in props:
            if prop in properties:
                properties[prop].append(state)
            else:
                properties[prop] = [state]
    
    print("Properties")
    for prop,pstates in properties.items():
        ss = ','.join(pstates)
        print(f" - '{prop}' when {ss}")
    print("")

    ltls = []
    if args.ltlfile is not None:
        with open(args.ltlfile) as f:
            ltls = [line.rstrip() for line in f]


    mod = ModelBuilder(states, args.initial, comp_tt, properties, ltls=ltls)
    nusmv = mod.generate()

    if args.output is not None:
        with open(args.output, "wt") as f:
            n = f.write(nusmv)
    else:
        print("-------------------")
        print(nusmv)