import java.sql.{Connection, DriverManager, SQLException}

@throws(classOf[SQLException])
class MarkovGenerator(private var dbUrl: String, private val table: String = "markov"){


  private var connection: Connection = _

  if(!dbUrl.startsWith("jdbc:sqlite:")){
    dbUrl = "jdbc:sqlite:" + dbUrl
  }

  try{
    connection = DriverManager.getConnection(dbUrl)
    val sql = s"CREATE TABLE IF NOT EXISTS $table(" +
      " key TEXT PRIMARY KEY," +
      " value TEXT);"
    val stmt = connection.createStatement()
    stmt.execute(sql)
  }
  catch{
    case e: SQLException =>
      throw e
  }

  /**
    * Resets the sql table
    */
  def resetTable(): Unit ={
    var sql = s"DROP TABLE IF EXISTS $table"
    var stmt = connection.createStatement()
    stmt.execute(sql)

    sql = s"CREATE TABLE IF NOT EXISTS $table(" +
      " key TEXT PRIMARY KEY," +
      " value TEXT);"
    stmt = connection.createStatement()
    stmt.execute(sql)
  }

  /**
    * Prints the sql table
    */
  def printTable(): Unit ={
    val sql = s"SELECT key, value FROM $table"
    val rs = connection.createStatement().executeQuery(sql)
    while (rs.next()){
      println(rs.getString(1) + " === " + rs.getString(2) )
    }
  }

  private def hasChainedWords(wordPair: String): Boolean = {
    val sql = s"SELECT key FROM $table WHERE key = ?"
    val pstmt = connection.prepareStatement(sql)
    pstmt.setString(1, wordPair)
    val rs = pstmt.executeQuery()
    rs.next()
    rs.getRow > 0
  }

  private def getChainedWordsFromPair(wordPair: String): Array[String] = {
    val sql = s"SELECT key, value FROM $table WHERE key = ?"
    val pstmt = connection.prepareStatement(sql)
    pstmt.setString(1, wordPair)
    val rs = pstmt.executeQuery()
    rs.getString("value").split("\\s+")
  }

  private def chainWordToPair(wordPair: String, word: String): Unit ={
    val hasChained = hasChainedWords(wordPair)
    val sql = if(hasChained){
      s"UPDATE $table SET value = ? WHERE key = ? "
    }
    else{
      s"INSERT INTO $table(value, key) VALUES(?,?)"
    }

    var newWords = word

    if(hasChained){
      val chained = getChainedWordsFromPair(wordPair)
      if(chained.contains(word)) return
      for(chainedWord <- chained){
        newWords = newWords +  " " + chainedWord
      }

    }
    val pstmt = connection.prepareStatement(sql)
    pstmt.setString(1, newWords)
    pstmt.setString(2, wordPair)
    pstmt.executeUpdate()
  }



  /**
    * Parses a sentence to be included in the chain
    */
  def parseSentence(sentence: String): Unit ={
    val split = sentence.split("\\s+")
    if(split.length > 1){
      for(i <- split.indices){
        if(i < split.length - 2){
          val wordPair = split(i) + " " + split(i+1)
          chainWordToPair(wordPair, split(i+2))
        }
      }

    }


  }


}