/*
 * Copyright (C) 2013 Matthew Dawson <matthew@mjdsystems.ca>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package ca.mcmaster.cas.smtlibchecker;

import ca.mcmaster.cas.matlab2smt.SMTLIBGenerator;
import ca.mcmaster.cas.matlab2smt.Variable;
import ca.mcmaster.cas.matlab2smt.VariableParser;
import ca.mcmaster.cas.tablularexpression.HierarchcialGridCheckerWalkerGenerator;
import ca.mcmaster.cas.tablularexpression.HierarchicalGrid;
import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.misc.NotNull;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.ParseTreeWalker;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Reader;
import java.io.UnsupportedEncodingException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 *
 * @author Matthew Dawson <matthew@mjdsystems.ca>
 */
final public class CheckerRunner {
    public static List<CheckerRunnerResult> RunZ3(HierarchicalGrid grid, VariableParser vars) {
        List<CheckerRunnerResult> ret = new ArrayList<CheckerRunnerResult>();

        // Get the queries
        HierarchicalGridSMTLIBGenerator queryGenerator = new HierarchicalGridSMTLIBGenerator(vars);
        HierarchcialGridCheckerWalkerGenerator.GenerateCheckerFromGrid(grid, queryGenerator);

        // Start Z3
        ProcessBuilder pb;
        Process p;
        try {
            pb = new ProcessBuilder("z3", "-smt2", "-in");
            p = pb.start();
        } catch(IOException ex) {
            throw new RuntimeException("Failed to start Z3!", ex);
        }

        // Input/Output refer to the direction from Z3's pov.  So output is Z3's stdout.
        BufferedReader z3Output;
        Writer z3Input;
        try {
            // Sigh ... Java has horrible naming conversions here ...
            z3Output = new BufferedReader(new InputStreamReader(p.getInputStream(), "ISO-8859-1"));
            z3Input = new OutputStreamWriter(p.getOutputStream());
        } catch(UnsupportedEncodingException ex) {
            throw new RuntimeException("Failed to get ISO-8859-1 charset for talking to Z3!", ex);
        }

        StringBuilder GetValueRequestBuilder = new StringBuilder();
        GetValueRequestBuilder.append("(get-value (");
        for(Variable var:vars.getVarList()) {
            GetValueRequestBuilder.append(var.name()).append(' ');
        }
        GetValueRequestBuilder.append("))\n");
        String GetValueRequestString = GetValueRequestBuilder.toString();

        try {
            String VarDef = vars.getCheckerOutput(new SMTLIBGenerator());
            z3Input.write(VarDef);
            for(String query : queryGenerator.getFinalQueries()) {
                z3Input.write("(push 1)");
                z3Input.write(query);
                z3Input.write("(check-sat)\n");
                z3Input.flush();

                String res = z3Output.readLine();
                if(res.equals("sat")) { // Failed query.  Get a the result.
                    z3Input.write(GetValueRequestString);
                    z3Input.flush();

                    CheckerRunnerResult result = new CheckerRunnerResult();
                    result.Query = query;
                    result.SampleValues = parseSMTLIBOutput(z3Output);
                    ret.add(result);
                } else if(!res.equals("unsat")) {
                    // Something went wrong
                    throw new IllegalStateException("Z3 outputed something wrong (" + res + ")!");
                }

                z3Input.write("(pop 1)\r");
            }
        } catch(IOException ex) {
            throw new RuntimeException("Something bad happened ...", ex);
        } finally {
            try {
                z3Input.write("(exit)\n");
                z3Input.flush();

                p.destroy();
                z3Input.close();
                z3Output.close();
            } catch(Exception e) {
                // Don't care here, so just print it out.
                System.out.println(e);
            }
        }

        return ret;
    }

    private static Map<String, String> parseSMTLIBOutput(Reader reader) throws IOException {
        if(reader.markSupported() == false) {
            throw new IllegalArgumentException("Reader must support marks!");
        }
        final Map<String, String> ret = new HashMap<String, String>();

        StringBuilder input = new StringBuilder();

        int bracketMismatch = 1; // This starts with a (, so we immediately have a a mismatch.
        char nextChar = getNextChar(reader);
        if (nextChar != '(') { // Should start with a ( for output.
            throw new IllegalStateException("Bad input, wanted a ( to start, got " + nextChar);
        }
        input.append(nextChar);
        while(bracketMismatch != 0) {
            nextChar = getNextChar(reader);
            switch (nextChar) {
                case '(':
                    bracketMismatch++;
                    break;
                case ')':
                    bracketMismatch--;
                    break;
            }
            input.append(nextChar);
        }

        // Finally, clear the new line, so the output is now clear.
        do {
            reader.mark(5);
            nextChar = getNextChar(reader);
        } while(isWhitespace(nextChar) && nextChar != '\n');
        if(!isWhitespace(nextChar)) {
            reader.reset();
        }

        ANTLRInputStream antlrInput = new ANTLRInputStream(input.toString());

        // create a lexer that feeds off of input CharStream
        SMTLibOutParserLexer lexer = new SMTLibOutParserLexer(antlrInput);

        // create a buffer of tokens pulled from the lexer
        CommonTokenStream tokens = new CommonTokenStream(lexer);

        // create a parser that feeds off the tokens buffer
        SMTLibOutParserParser parser = new SMTLibOutParserParser(tokens);

        ParseTree tree = parser.valuelist();
        ParseTreeWalker walker = new ParseTreeWalker();
        SMTLibOutParserListener listener = new SMTLibOutParserBaseListener() {
            @Override
            public void enterValue(@NotNull SMTLibOutParserParser.ValueContext ctx) {
                id = ctx.ID().getText();
                value = "";
            }

            @Override
            public void enterNegative(@NotNull SMTLibOutParserParser.NegativeContext ctx) {
                value = "(- ";
            }

            @Override
            public void exitSingle(@NotNull SMTLibOutParserParser.SingleContext ctx) {
                value += ctx.getText();
            }

            @Override
            public void exitFraction(@NotNull SMTLibOutParserParser.FractionContext ctx) {
                value += "(" + ctx.RATIONAL(0).getText() + " / " + ctx.RATIONAL(1).getText() + ")";
            }

            @Override
            public void exitNegative(@NotNull SMTLibOutParserParser.NegativeContext ctx) {
                value += ')';
            }

            @Override
            public void exitExpr(@NotNull SMTLibOutParserParser.ExprContext ctx) {
                ret.put(id, value);
            }

            String id = "";
            String value;
        };
        walker.walk(listener, tree);
        return ret;
    }

    private static char getNextChar(Reader reader) throws IOException {
        char[] c = new char[1];
        reader.read(c);
        return c[0];
    }

    private static boolean isWhitespace(char toTest) {
        if(toTest == ' ' ||
                toTest == '\n' ||
                toTest == '\r' ||
                toTest == '\t') {
            return true;
        }
        return false;
    }
}
