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

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('--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/')

    p.parse(source)
    oh = io.StringIO()
    p.write(oh)
    prep_source = oh.getvalue()
    
    #print(prep_source)
    #exit()

    parser = CParser()
    ast = parser.parse(prep_source)
    #ast = parse_file(args.filename, use_cpp=False)
    
    assign_table = []
    state_enums = []
    func_table = {}
    enum_table = {}
    state_asmts = []
    
    fdv = FuncDefVisitor()
    fdv.visit(ast)
    func_table = fdv.func_table

    etv = EnumTypedefVisitor()
    etv.visit(ast)
    enum_table = etv.enums
    
    if not(args.func in func_table):
        raise Exception(f"Function name '{args.func}' not found!")

    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)
    
    if args.dot is not None:
        g = gv.Digraph('G')
        
        for state in states:
            if state in props_by_state:
                pstr = ",".join(props_by_state[state])
                if state == args.initial:
                    g.attr('node', shape='doublecircle')
                else:
                    g.attr('node', shape='circle')
                g.node(state, label=state, xlabel=f"{{{pstr}}}")
              
                
        
        for t in tran_table:
            g.edge(t[0], t[1])
        g.render(filename=args.dot)