from __future__ import print_function
import argparse
import os,sys,io
from pycparser import parse_file, c_ast, CParser
from pcpp import Preprocessor
import graphviz as gv
import html

from modelbuilder import *
from utils import *
from astvisitors import *

if __name__ == "__main__":
    argparser = argparse.ArgumentParser('Create a Kripke Structure Model from C Code')
    argparser.add_argument('filename',
                           nargs='?',
                           help='name of file to parse')

    argparser.add_argument('-I', '--includedirs', nargs='+', default=[])
    argparser.add_argument('-D', '--defines', nargs='+', default=[])
    argparser.add_argument('-p', '--preprocess', help='output preprocessor file')
    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')
    argparser.add_argument('--dot', help='output dot file')
    args = argparser.parse_args()
    
    source = ""
    with open(args.filename, "r") as f:     
        #read whole file to a string
        source = f.read()
    
    p = Preprocessor()
    p.add_path('/usr/lib/gcc/x86_64-linux-gnu/12/include/')
    for inc in args.includedirs:
        print(f"Include-Dir: {inc}")
        p.add_path(inc)
    
    for define in args.defines:
        name,value = define.split('=')
        print(f"Define: {name}={value}")
        p.define(f"{name} {value}")

    p.parse(source)
    oh = io.StringIO()
    p.write(oh)
    prep_source = oh.getvalue()#unicode(oh.getvalue(), errors='ignore')
    
    if args.preprocess is not None:
        with open(args.preprocess, "wt") as f:
            n = f.write(prep_source)

    parser = CParser()
    ast = parser.parse(prep_source)
    #ast = parse_file(args.filename, use_cpp=False)
    #ast.show()

    
    initial_state = args.initial
    assign_table = []
    state_enums = []
    func_table = {}
    enum_table = {}
    state_asmts = []
    fsm_funcs = []

    fdv = FuncDefVisitor()
    fdv.visit(ast)
    func_table = fdv.func_table
    proc_func = None
    if not(args.func in func_table):
        raise Exception(f"Function name '{args.func}' not found!")
    else:
        proc_func = func_table[args.func]

    etv = EnumTypedefVisitor()
    etv.visit(ast)
    enum_table = etv.enums
    states = []
    
    def discover_cals(func, visited=None):
        if visited is None:
            visited = []
        funcs = []
        child_funcs = []
        fcv = FuncCallVisitor()
        fcv.visit(func)
        visited.append(proc_func)
        for fc in fcv.func_calls:
            if fc in func_table:
                funcs.append(func_table[fc])
                child_funcs += discover_cals(func_table[fc], visited)
        return funcs + child_funcs

    # discover FSM functions
    fsm_funcs.append(proc_func)
    fsm_funcs += discover_cals(proc_func)
    #fsm_funcs += [ func_table[x] for x in fcv.func_calls ]
    
    print("Function Table")
    for f in fsm_funcs:
        print(f" - {f.decl.name} {'<<< entry' if (f.decl.name == args.func) else ''}")
    print("")

    print("Enum Table")
    if args.enum in enum_table:
        ev = EnumVisitor()
        ev.visit(enum_table[args.enum])
        for ename in ev.enum_names:
            print(" - ",ename)
            states.append(ename)
            #for f in fsm_funcs:
            sav = StateAssignmentVisitor(ast, ename)
            #sav.visit(f)
            sav.visit(proc_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:
        for f in fsm_funcs:
            sctv = SwitchCaseTranVisitor(sa[0], sa[1], states, sa[2], sa[3], sa[4])
            sctv.visit(f)
            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:
            if t[1] not in comp_tt[t[0]]:
                comp_tt[t[0]].append(t[1])
        else:
            comp_tt[t[0]] = [t[1]]
    print("")

    pure_sa = [x[0] for x in state_asmts]
    #todo recomment once fixed
    #states = comp_tt.keys()
    print("States: ", ",".join(states))
    print("")
    
    prop_func_blacklist = [f.decl.name for f in fsm_funcs]
    reachable_states = [initial_state]
    props_by_state = {}
    print("Compact Transition Table:")
    for n,ms in comp_tt.items():
        sstr = ','.join(ms)
        print(f"{n}->{{{sstr}}}")
        for s in ms:
            if not(s in reachable_states):
                reachable_states.append(s)
        # find properties
        for f in fsm_funcs:
            sccpv = SwitchCaseCodePropertyVisitor(n, pure_sa, prop_func_blacklist)
            sccpv.visit(f)
            if len(sccpv.properties) > 0:
                props_by_state[n] = sccpv.properties
    print("")

    states_by_property = {}
    for state,props in props_by_state.items():
        for prop in props:
            if prop in states_by_property:
                states_by_property[prop].append(state)
            else:
                states_by_property[prop] = [state]
    
    properties = {}
    property_alias = {}
    for i,(prop,pstates) in enumerate(states_by_property.items()):
        alias = base_10_to_alphabet(i + 1)
        property_alias[prop] = alias
        properties[alias] = (pstates, prop)

    print("Properties")
    for prop,(pstates,full_prop) in properties.items():
        ss = ','.join(pstates)
        print(f" - '{full_prop}' when {ss}")
    print("")
    
    print("States shortform:")
    states_prefix = os.path.commonprefix(states)
    state_to_short = {}
    for s in states:
        state_to_short[s] = remove_prefix(s, states_prefix)
    ltls = []
    if args.ltlfile is not None:
        with open(args.ltlfile) as f:
            ltls = [line.rstrip() for line in f]
    
    
    mod = ModelBuilder(reachable_states, initial_state, 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)
    
    if args.dot is not None:
        g = gv.Digraph('G')
        # add states
        for state in reachable_states:
            state_short = state_to_short[state]
            shape = 'oval'
            if state == initial_state:
                shape = 'doubleoctagon'

            if state in props_by_state:
                pstr = ",".join([property_alias[x] for x in props_by_state[state]])
                g.node(state_short, label=f"{state_short}\n\n{{{pstr}}}", shape=shape)
            else:
                g.node(state_short, label=state_short, shape=shape)
        
        # add transitions
        for n,ms in comp_tt.items():
            for m in ms:
                g.edge(state_to_short[n], state_to_short[m])

        # add property table
        """
        tabattr = 'border="0px"'
        table_rows = []
        for prop,alias in property_alias.items():
            alias = html.escape(alias)
            prop = html.escape(prop)
            table_rows.append(f"<tr><td>{alias}</td><td>{prop}</td></tr>")
        table_rows = ''.join(table_rows)
        html = f"<table {tabattr}>{table_rows}</table>"
        print(html)
        g.node("Properties", label=f"<{html}>",shape="none", rank="sink")
        """
        g.render(filename=args.dot)