import os
import time
import json

from pysmt.smtlib.parser import SmtLibParser
from pysmt.shortcuts import Solver
from io import StringIO
from pysmt.exceptions import PysmtSyntaxError, PysmtTypeError, UnknownSmtLibCommandError

from utils import solve
import utils.utils as utils
import utils.auto_utils as auto_utils

error_tuples = (AttributeError, TypeError, AssertionError, PysmtSyntaxError, PysmtTypeError, NotImplementedError, UnknownSmtLibCommandError)

def check(data, index, client, logger, args):
    Q, A = data['question'], data['answer']
    file_name = "problem_%s" %(args.index[0] + index)
    save_file_path = os.path.join(args.save_folder_path, file_name+'.json')
    if not os.path.exists(save_file_path):
        json_data = dict()
        answer = auto_utils.parse_answer(A)
        thy_series = {'problem': Q, 
                      'solution': A, 
                      'answer': answer}
        with open(save_file_path, 'w') as file:  
            json.dump(thy_series, file, indent=4)
    else:
        with open(save_file_path) as file:
            thy_series = json.load(file)
        Q = thy_series['problem']
        answer = thy_series['answer']
    #### check the formal problem
    if args.save_iter == -1:
        check = thy_series.get("check", "")
        formal_problem = thy_series.get("formal problem", "")
    else:
        thy = thy_series.get(f'generation_{max(args.save_iter,0)}', {})
        check = thy.get("answer", "")
        formal_problem = thy.get("formal problem", "")
    if check != "" and check is not None and args.overwrite == 0:
        logger.info("skip")
    else:
        t0 = time.time()
        try:
            if formal_problem == "" or formal_problem is None: raise TypeError("no formal problem exists")
            ok, solution = solve.solve(formal_problem)
            if ok == True:
                tmp_sol = ",".join(solution)
            else:
                tmp_sol = "unsat"
            correct = utils.is_equiv(answer, tmp_sol)
            #### save check result
            if args.save_iter == -1:
                thy_series['check'] = tmp_sol
            else:
                thy_series[f'generation_{max(args.save_iter,0)}']['answer'] = tmp_sol
            with open(save_file_path, 'w') as file:  
                json.dump(thy_series, file, indent=4)
                ''' 
                check the answer 
                note: if the iter is -1, the answer is the original answer, else the answer is the generated answer
                '''
                msg = formal_problem + '\nTO ========>>>\n' + tmp_sol + "\nThe original answer is : " + answer + '\t' + f"({correct})"
        except error_tuples as e:
            msg = formal_problem + '\nTO ========>>>\n' + "Compile error %s; reset the formalization" %(e)
            ok = False
        except Exception as e:
            raise e
        if ok == False:
            if args.save_iter == -1:
                thy_series['formal problem'] = ""
                thy_series['check'] = ""
            else:
                thy_series[f'generation_{max(args.save_iter,0)}']['formal problem'] = ""
                thy_series[f'generation_{max(args.save_iter,0)}']['answer'] = ""
            thy_series['formal problem'] = ""
            with open(save_file_path, 'w') as file:  
                    json.dump(thy_series, file, indent=4)
        t1 = time.time()
        logger.info('success! In (%s)s' % (t1-t0))
        logger.info(msg)

        
                        
