package com.github.vlsi.mat.calcite;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.Statement;
import java.util.regex.Pattern;

import javax.sql.rowset.CachedRowSet;
import javax.sql.rowset.RowSetFactory;
import javax.sql.rowset.RowSetProvider;

import org.eclipse.mat.query.IQuery;
import org.eclipse.mat.query.IResult;
import org.eclipse.mat.query.annotations.Argument;
import org.eclipse.mat.query.annotations.Category;
import org.eclipse.mat.query.annotations.CommandName;
import org.eclipse.mat.query.annotations.Name;
import org.eclipse.mat.query.results.TextResult;
import org.eclipse.mat.snapshot.ISnapshot;
import org.eclipse.mat.util.IProgressListener;

@Name("SQL via Calcite")
@Category("SQL")
@CommandName("calcite")
public class Query implements IQuery {
	public static final Pattern EXPLAIN_PLAN = Pattern.compile("explain\\s+plan\\s+for", Pattern.CASE_INSENSITIVE);
	public static final Pattern TRIM_ENUMERABLE = Pattern.compile("Enumerable(\\w+)Rel");
	public static final Pattern RENAME_INDEX = Pattern.compile("EnumerableTableScan\\(table=\\[\\[HEAP, \\$ids\\$:([^]]+)\\]\\]\\)");
	public static final Pattern RENAME_JOIN = Pattern.compile("EnumerableJoin");
	public static final Pattern RENAME_CALC = Pattern.compile("EnumerableCalc");

	@Argument
	public ISnapshot snapshot;

	@Argument(flag = Argument.UNFLAGGED)
	public String sql;

	@Override
	public IResult execute(IProgressListener listener) throws Exception {
		Connection con = null;
		Statement st = null;
		ResultSet rs = null;

		try {
			con = CalciteDataSource.getConnection(snapshot);

			st = con.createStatement();
			rs = st.executeQuery(sql);

			RowSetFactory rowSetFactory = RowSetProvider.newFactory();
			CachedRowSet rowSet = rowSetFactory.createCachedRowSet();
			rowSet.populate(rs);

			ResultSetMetaData md = rowSet.getMetaData();
			if (md.getColumnCount() == 1 && "PLAN".equals(md.getColumnName(1))
					&& rowSet.size() == 1
					&& EXPLAIN_PLAN.matcher(sql).find()) {
				rowSet.absolute(1);
				String plan = rowSet.getString(1);
				System.out.println("plan = " + plan);
				plan = RENAME_INDEX.matcher(plan).replaceAll("GetObjectIdsByClass (class=$1)");
				plan = RENAME_JOIN.matcher(plan).replaceAll("HashJoin ");
				plan = RENAME_CALC.matcher(plan).replaceAll("View ");
				plan = TRIM_ENUMERABLE.matcher(plan).replaceAll("$1 ");
				return new TextResult(plan, false);
			}

			return new RowSetTable(rowSet);
		} finally {
			CalciteDataSource.close(rs, st, con);
		}
	}

}

@@ -144,24 +144,56 @@
 .Statement;%0A
+import java.util.regex.Matcher;%0A
 import java.
@@ -212,16 +212,16 @@
 attern;%0A
-
 %0Aimport 
@@ -1050,24 +1050,86 @@
 %5C%5Cw+)Rel%22);%0A
+%09// EnumerableTableScan(table=%5B%5BHEAP, java, net, $ids$:URL%5D%5D)%0A
 %09public stat
@@ -1217,16 +1217,41 @@
 %5C%5BHEAP, 
+((?:%5B%5E%5C%5C$%5D%5B%5E,%5C%5C%5D%5D*+, )*+)
 %5C%5C$ids%5C%5C
@@ -2247,28 +2247,42 @@
  plan);%0A%09%09%09%09
-plan
+%7B%0A%09%09%09%09%09Matcher ind
  = RENAME_IN
@@ -2302,54 +2302,298 @@
 lan)
-.replaceAll(%22GetObjectIdsByClass (class=$1)%22);
+;%0A%09%09%09%09%09StringBuffer sb = new StringBuffer();%0A%09%09%09%09%09while (ind.find()) %7B%0A%09%09%09%09%09%09String className = ind.group(1).replace(%22, %22, %22.%22) + ind.group(2);%0A%09%09%09%09%09%09ind.appendReplacement(sb, %22GetObjectIdsByClass (class=%22 + className + %22)%22);%0A%09%09%09%09%09%7D%0A%09%09%09%09%09ind.appendTail(sb);%0A%09%09%09%09%09plan = sb.toString();%0A%09%09%09%09%7D
 %0A%09%09%09

Support EnumerableTableScan -> GetObjectIdsByClass when displaying plans for schema.schema.Class