package de.mpg.mpi_inf.ambiversenlu.nlu.entitylinking.datapreparation.yago.yago3.util;

import de.mpg.mpi_inf.ambiversenlu.nlu.entitylinking.EntityLinkingManager;
import de.mpg.mpi_inf.ambiversenlu.nlu.tools.basics3.Fact;
import edu.stanford.nlp.util.StringUtils;
import org.apache.commons.lang.StringEscapeUtils;

import java.io.IOException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.LinkedList;
import java.util.List;

public class YAGO3DBReader implements YAGO3Reader {

  private String tableName;

  public YAGO3DBReader(String tableName) {
    this.tableName = tableName;
  }

  @Override public List<Fact> getFacts(String relation) throws IOException {
    List<Fact> triples = new LinkedList<>();
    Statement stmt;
    ResultSet rs;

    try {
      Connection con = EntityLinkingManager.getConnectionForDatabase(EntityLinkingManager.DB_YAGO);
      con.setAutoCommit(false);

      stmt = con.createStatement();
      stmt.setFetchSize(10000000);
      String sql = "SELECT id, subject, object " + "FROM " + tableName + " WHERE predicate='" + relation + "'";
      rs = stmt.executeQuery(sql);

      while (rs.next()) {
        String id = rs.getString("id");
        String subject = rs.getString("subject");
        String object = rs.getString("object");
        triples.add(new Fact(id, subject, relation, object));
      }
      rs.close();
      con.setAutoCommit(true);

      EntityLinkingManager.releaseConnection(con);
    } catch (SQLException e) {
      throw new IOException(e);
    }

    return triples;
  }

  public List<Fact> getFacts(String subject, String relation, String object) throws SQLException {
    List<Fact> triples = new LinkedList<>();
    Statement stmt;
    ResultSet rs;

    Connection con = EntityLinkingManager.getConnectionForDatabase(EntityLinkingManager.DB_YAGO);

    List<String> queryCondition = new LinkedList<>();
    if (subject != null) {
      queryCondition.add("subject='" + StringEscapeUtils.escapeSql(subject) + "'");
    }
    if (relation != null) {
      queryCondition.add("predicate='" + StringEscapeUtils.escapeSql(relation) + "'");
    }

    if (object != null) {
      queryCondition.add("object='" + StringEscapeUtils.escapeSql(object) + "'");
    }

    String conditionString = StringUtils.join(queryCondition, " AND ");
    stmt = con.createStatement();
    String sql = "SELECT id, subject, object, predicate  FROM " + tableName + " WHERE " + conditionString;
    rs = stmt.executeQuery(sql);

    while (rs.next()) {
      triples.add(new Fact(rs.getString("id"), rs.getString("subject"), rs.getString("predicate"), rs.getString("object")));
    }
    rs.close();

    EntityLinkingManager.releaseConnection(con);

    return triples;
  }
}