package sqlancer;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;
import java.util.function.BiConsumer;

import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.CaseExpression;
import net.sf.jsqlparser.expression.DoubleValue;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.NullValue;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.WhenClause;
import net.sf.jsqlparser.expression.operators.relational.Between;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.StatementVisitorAdapter;
import net.sf.jsqlparser.statement.select.GroupByElement;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SelectExpressionItem;
import net.sf.jsqlparser.statement.select.SelectVisitorAdapter;
import net.sf.jsqlparser.statement.select.SetOperationList;
import net.sf.jsqlparser.statement.select.SubSelect;
import net.sf.jsqlparser.statement.select.WithItem;
import sqlancer.common.query.Query;
import sqlancer.common.query.SQLQueryAdapter;

final class ExpressionTransformer {

    private static List<Expression> flattenChildren(BinaryExpression expr) {
        List<Expression> candidates = new ArrayList<>();
        Expression lhs = expr.getLeftExpression();
        Expression rhs = expr.getRightExpression();
        candidates.add(lhs);
        candidates.add(rhs);
        return candidates;
    }

    private static List<Expression> flattenChildren(Between expr) {
        Expression lhs = expr.getBetweenExpressionStart();
        Expression rhs = expr.getBetweenExpressionEnd();
        return List.of(lhs, rhs);
    }

    public static List<Expression> candidateExpressions(Expression expr) {
        if (expr instanceof Parenthesis) {
            // try removing a pair of brackets.
            Parenthesis paren = (Parenthesis) expr;
            return List.of(paren.getExpression());
        } else if (expr instanceof BinaryExpression) {
            return flattenChildren((BinaryExpression) expr);
        } else if (expr instanceof Between) {
            return flattenChildren((Between) expr);
        } else if (expr instanceof LongValue) {
            LongValue longValue = (LongValue) expr;
            if (String.valueOf(longValue).length() >= 4) {
                return List.of(new NullValue(), new LongValue(10), new LongValue(0), new LongValue(1));
            }
            return new ArrayList<>();
        } else if (expr instanceof DoubleValue) {
            DoubleValue doubleValue = (DoubleValue) expr;
            double literal = doubleValue.getValue();
            if (String.valueOf(literal).length() <= 4) {
                return new ArrayList<>();
            }
            double roundedValue = Math.round(literal * 10.0) / 10.0;
            return List.of(new NullValue(), new DoubleValue(String.valueOf(roundedValue)));
        } else if (expr instanceof StringValue) {
            StringValue sv = (StringValue) expr;
            String str = sv.getValue();
            if (str.length() > 4) {
                return List.of(new NullValue(), new StringValue(" "));
            }
            return new ArrayList<>();
        } else if (expr instanceof CaseExpression) {
            CaseExpression caseExpression = (CaseExpression) expr;
            return List.of(caseExpression.getSwitchExpression(), caseExpression.getElseExpression());
        } else {
            return new ArrayList<>();
        }
    }

    private ExpressionTransformer() throws Exception {
        throw new AssertionError("Do not initialize the util class");
    }
}

@SuppressWarnings("unchecked")
public class ASTBasedReducer<G extends GlobalState<O, ?, C>, O extends DBMSSpecificOptions<?>, C extends SQLancerDBConnection>
        implements Reducer<G> {

    private final DatabaseProvider<G, O, C> provider;

    @SuppressWarnings("unused")
    private G state;
    private G newGlobalState;
    private Reproducer<G> reproducer;
    private int reduceTargetIndex;
    private Statement targetStatement;

    // statement after reduction.
    private List<Query<C>> reducedStatements;

    public ASTBasedReducer(DatabaseProvider<G, O, C> provider) {
        this.provider = provider;
    }

    private void updateStatements() {
        String queryString = targetStatement.toString();
        boolean couldAffectSchema = queryString.contains("CREATE TABLE") || queryString.contains("EXPLAIN");
        reducedStatements.set(reduceTargetIndex, (Query<C>) new SQLQueryAdapter(queryString, couldAffectSchema));
    }

    public <P> void expressionReduce(P parent, Expression subExpr, // NOPMD
            BiConsumer<P, Expression> setter) {
        boolean observeChange;
        do {
            observeChange = false;
            List<Expression> candidates = ExpressionTransformer.candidateExpressions(subExpr);
            for (Expression candidate : candidates) {
                try {
                    setter.accept(parent, candidate);
                    if (bugStillTriggers()) {
                        subExpr = candidate;
                        observeChange = true;
                    }
                } catch (Exception ignoredException) {
                }
            }
            setter.accept(parent, subExpr);
        } while (observeChange);
    }

    public <P, T> void listElementRemovingReduce(P parent, List<T> elms, // NOPMD
            BiConsumer<P, List<T>> setter) {
        // TODO: For AST-Reducer, is delta-debugging needed ? Or just use the naive approach ?
        boolean observeChange;
        do {
            observeChange = false;
            for (int i = elms.size() - 1; i >= 0; i--) {
                List<T> reducedElms = new ArrayList<>(elms);
                reducedElms.subList(i, i + 1).clear();
                setter.accept(parent, reducedElms);
                try {
                    if (bugStillTriggers()) {
                        elms = reducedElms;
                        observeChange = true;
                    }
                } catch (Exception e) {
                    System.out.println("An error occurred when trying executing reduced statements");
                    e.printStackTrace();
                }
            }
            setter.accept(parent, elms);
        } while (observeChange);

    }

    ExpressionVisitorAdapter expressionReducerVisitor = new ExpressionVisitorAdapter() {

        @Override
        public void visit(InExpression expr) {
            Expression rhs = expr.getRightExpression();

            if (rhs instanceof SubSelect) {
                SubSelect subSelect = (SubSelect) rhs;
                subSelect.getSelectBody().accept(selectReducerVisitor);
            } else {
                ItemsList itemslist = expr.getRightItemsList();
                itemslist.accept(this);
            }
        }

        @Override
        protected void visitBinaryExpression(BinaryExpression expr) {
            Expression lhs = expr.getLeftExpression();
            Expression rhs = expr.getRightExpression();
            expressionReduce(expr, lhs, (expression, candidate) -> {
                expression.setLeftExpression(candidate);
                updateStatements();
            });
            expressionReduce(expr, rhs, (expression, candidate) -> {
                expression.setRightExpression(candidate);
                updateStatements();
            });
            lhs.accept(this);
            rhs.accept(this);
        }

        // @Override
        // public void visit(DateValue value) {
        // super.visit(value);
        // }
        //
        // @Override
        // public void visit(TimeValue value) {
        // super.visit(value);
        // }
        //
        // @Override
        // public void visit(LikeExpression expr) {
        // super.visit(expr);
        // }

        @Override
        public void visit(CaseExpression expr) {
            Expression switchExpr = expr.getSwitchExpression();
            Expression elseExpr = expr.getElseExpression();

            expressionReduce(expr, switchExpr, (parent, sw) -> {
                parent.setSwitchExpression(sw);
                updateStatements();
            });
            expressionReduce(expr, elseExpr, (parent, els) -> {
                parent.setElseExpression(els);
                updateStatements();
            });

            super.visit(expr);
        }

        @Override
        public void visit(WhenClause whenClause) {
            Expression when = whenClause.getWhenExpression();
            Expression then = whenClause.getThenExpression();

            expressionReduce(whenClause, when, (wc, w) -> {
                wc.setWhenExpression(w);
                updateStatements();
            });

            expressionReduce(whenClause, then, (wc, t) -> {
                wc.setThenExpression(t);
                updateStatements();
            });

            super.visit(whenClause);
        }

        @Override
        public void visit(Parenthesis parenthesis) {
            Expression closedExpr = parenthesis.getExpression();
            expressionReduce(parenthesis, closedExpr, (p, s) -> {
                p.setExpression(s);
                updateStatements();
            });
            closedExpr.accept(this);
        }

        // @Override
        // public void visit(Function function) {
        // super.visit(function);
        // }

        @Override
        public void visit(ExpressionList expressionList) {
            List<Expression> expressions = expressionList.getExpressions();
            listElementRemovingReduce(expressionList, expressions, (l, es) -> {
                l.setExpressions(es);
                updateStatements();
            });
            expressions = expressionList.getExpressions();
            for (int i = 0; i < expressions.size(); i++) {
                Expression expr = expressions.get(i);
                int index = i;
                expressionReduce(expressions, expr, (l, e) -> {
                    l.set(index, e);
                    updateStatements();
                });
            }
            super.visit(expressionList);
        }

        @Override
        public void visit(SelectExpressionItem selectExpressionItem) {
            Expression expr = selectExpressionItem.getExpression();
            expressionReduce(selectExpressionItem, expr, (item, e) -> {
                item.setExpression(e);
                updateStatements();
            });
            super.visit(selectExpressionItem);
        }

        @Override
        public void visit(SubSelect subSelect) {
            subSelect.getSelectBody().accept(selectReducerVisitor);
        }

    };

    SelectVisitorAdapter selectReducerVisitor = new SelectVisitorAdapter() {

        // Clauses that would be tried removing.
        // examples:
        // Remove when: select * from table when 1 -> select * from table
        // Remove limit: select * from table limit 1 -> select * from table
        private final String[] removeList = { "Limit", "Offset", "Where", "Having", "GroupBy", "Distinct",
                "OrderByElements", "Joins" };

        // Clauses that would be tried transforming.
        // examples:
        // select * from t where a + b < c + d
        // where clause might become one of the statement below after transformation:
        // -> select * from table where c + d
        // -> select * from table where a + b

        private final String[] transformList = { "Where", "Having", "FromItem", "SelectItems", "GroupBy", "Joins" };

        // Clauses that would be visited for further reduction.
        // example:
        // select * from t where a + b < c + d
        // Assuming that the coexistence of a and c would trigger the bug.
        // The where clause : a + b < c + d would be visited and a + b, c + d would be reduced respectively.
        // a + b < c + d might become a + c
        private final String[] descendList = { "Where", "Having", "FromItem", "SelectItems", "GroupBy" };

        private String getterName(String astNodeName) {
            return "get" + astNodeName;
        }

        private String setterName(String astNodeName) {
            if (astNodeName.equals("GroupBy")) {
                return "set" + astNodeName + "Element";
            } else {
                return "set" + astNodeName;
            }
        }

        @Override
        public void visit(WithItem withItem) {
            // withItem.getItemsList();
        }

        @Override
        public void visit(PlainSelect plainSelect) {
            // transform section. Lists defined above would be iterated to get the corresponding clause name. Reflection
            // is used to avoid repetitive code. e.g. The current astNodeName is When `getWhen`, `setWhen` would be
            // called.
            for (String astNodeName : removeList) {
                try {
                    Method nodeGetter = plainSelect.getClass().getMethod(getterName(astNodeName));
                    Object astNode = nodeGetter.invoke(plainSelect);
                    if (astNode == null) {
                        continue;
                    }
                    Method nodeSetter = plainSelect.getClass().getMethod(setterName(astNodeName),
                            nodeGetter.getReturnType());
                    nodeSetter.invoke(plainSelect, new Object[] { null });
                    updateStatements();
                    if (!bugStillTriggers()) {
                        nodeSetter.invoke(plainSelect, astNode);
                        updateStatements();
                    }
                } catch (Exception e) {
                    throw new AssertionError(e);
                }
            }

            // Pull Up Section
            for (String astNodeName : transformList) {
                try {
                    Method nodeGetter = plainSelect.getClass().getMethod(getterName(astNodeName));
                    Object astNode = nodeGetter.invoke(plainSelect);
                    if (astNode == null) {
                        continue;
                    }
                    Method nodeSetter = plainSelect.getClass().getMethod(setterName(astNodeName),
                            nodeGetter.getReturnType());

                    if (astNode instanceof Expression) {
                        Expression expr = (Expression) astNode;
                        expressionReduce(plainSelect, expr, (select, expression) -> {
                            try {
                                nodeSetter.invoke(select, expression);
                                updateStatements();
                            } catch (IllegalAccessException | InvocationTargetException e) {
                                e.printStackTrace();
                            }
                        });
                    } else if (astNode instanceof List) {
                        List<?> elms = (List<?>) astNode;
                        if (elms.size() <= 1) {
                            continue;
                        }
                        listElementRemovingReduce(plainSelect, elms, (select, items) -> {
                            try {
                                nodeSetter.invoke(select, items);
                                updateStatements();
                            } catch (IllegalAccessException | InvocationTargetException e) {
                                e.printStackTrace();
                            }
                        });
                    } else if (astNode instanceof GroupByElement) {
                        GroupByElement groupByElement = (GroupByElement) astNode;
                        ExpressionList expressionList = groupByElement.getGroupByExpressionList();
                        if (expressionList == null) {
                            groupByElement.getGroupingSets();
                            // TODO: TO BE IMPLEMENTED.
                        } else {
                            List<Expression> elms = expressionList.getExpressions();
                            if (elms.size() <= 1) {
                                continue;
                            }
                            listElementRemovingReduce(groupByElement, elms, (select, items) -> {
                                groupByElement.setGroupByExpressionList(new ExpressionList(items));
                                updateStatements();
                            });
                        }
                    }
                } catch (Exception e) {
                    throw new AssertionError(e);
                }
            }

            for (String astNodeName : descendList) {
                try {
                    Method nodeGetter = plainSelect.getClass().getMethod(getterName(astNodeName));
                    Object astNode = nodeGetter.invoke(plainSelect);
                    if (astNode == null) {
                        continue;
                    }
                    if (astNode instanceof Expression) {
                        ((Expression) astNode).accept(expressionReducerVisitor);
                    } else if (astNode instanceof List) {
                        // Really hacky... Some other ways to simplify it ?
                        List<?> elms = (List<?>) astNode;
                        for (Object obj : elms) {
                            if (obj instanceof SelectExpressionItem) {
                                ((SelectExpressionItem) obj).accept(expressionReducerVisitor);
                            } else if (obj instanceof Expression) {
                                ((Expression) obj).accept(expressionReducerVisitor);
                            }
                        }
                    } else if (astNode instanceof GroupByElement) {
                        GroupByElement groupByElement = (GroupByElement) astNode;
                        ExpressionList expressionList = groupByElement.getGroupByExpressionList();
                        if (expressionList != null) {
                            expressionList.accept(expressionReducerVisitor);
                        }
                        // TODO: groupByElement.getGroupingSets() TO BE IMPLEMENTED

                    }
                } catch (Exception e) {
                    throw new AssertionError(e);
                }
            }
        }

        @Override
        public void visit(SetOperationList setOpList) {
            List<SelectBody> selectBodies = setOpList.getSelects();
            listElementRemovingReduce(setOpList, selectBodies, (optionList, selects) -> {
                optionList.setSelects(selects);
                updateStatements();
            });
            for (SelectBody selectBody : selectBodies) {
                if (selectBody instanceof PlainSelect) {
                    visit((PlainSelect) selectBody);
                }
            }
        }

    };

    StatementVisitorAdapter statementReducerVisitor = new StatementVisitorAdapter() {
        @Override
        public void visit(Select select) {
            SelectBody selectBody = select.getSelectBody();
            if (selectBody != null) {
                selectBody.accept(selectReducerVisitor);
            }
        }
    };

    @Override
    public void reduce(G state, Reproducer<G> reproducer, G newGlobalState) throws Exception {
        this.state = state;
        this.newGlobalState = newGlobalState;
        this.reproducer = reproducer;

        List<Query<?>> initialBugInducingStatements = state.getState().getStatements();
        reducedStatements = new ArrayList<>();
        for (Query<?> query : initialBugInducingStatements) {
            reducedStatements.add((Query<C>) query);
        }

        for (int i = 0; i < reducedStatements.size(); i++) {
            reduceTargetIndex = i;
            Query<C> query = reducedStatements.get(reduceTargetIndex);
            targetStatement = CCJSqlParserUtil.parse(query.getQueryString());
            targetStatement.accept(statementReducerVisitor);
        }

        newGlobalState.getState().setStatements(new ArrayList<>(reducedStatements));
    }

    private boolean bugStillTriggers() throws Exception {
        try (C con2 = provider.createDatabase(newGlobalState)) {
            newGlobalState.setConnection(con2);
            List<Query<C>> candidateStatements = new ArrayList<>(reducedStatements);
            newGlobalState.getState().setStatements(new ArrayList<>(candidateStatements));

            for (Query<C> s : candidateStatements) {
                try {
                    s.execute(newGlobalState);
                } catch (Throwable ignoredException) {
                    // ignore
                }
            }
            try {
                if (reproducer.bugStillTriggers(newGlobalState)) {
                    return true;
                }
            } catch (Throwable ignoredException) {

            }
        }
        return false;
    }
}