/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.flink.table.planner.runtime.stream.sql

import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment
import org.apache.flink.streaming.api.functions.timestamps.BoundedOutOfOrdernessTimestampExtractor
import org.apache.flink.table.api._
import org.apache.flink.table.api.bridge.scala._
import org.apache.flink.table.planner.runtime.utils.{StreamingEnvUtil, StreamingWithStateTestBase, TestingAppendSink}
import org.apache.flink.table.planner.runtime.utils.StreamingWithStateTestBase.StateBackendMode
import org.apache.flink.table.planner.utils.TableTestUtil
import org.apache.flink.testutils.junit.extensions.parameterized.ParameterizedTestExtension

import org.assertj.core.api.Assertions.assertThat
import org.junit.jupiter.api.TestTemplate
import org.junit.jupiter.api.extension.ExtendWith

import java.sql.Timestamp
import java.time.Duration

import scala.collection.mutable

@ExtendWith(Array(classOf[ParameterizedTestExtension]))
class TemporalTableFunctionJoinITCase(state: StateBackendMode)
  extends StreamingWithStateTestBase(state) {

  /**
   * Because of nature of the processing time, we can not (or at least it is not that easy) validate
   * the result here. Instead of that, here we are just testing whether there are no exceptions in a
   * full blown ITCase. Actual correctness is tested in unit tests.
   */
  @TestTemplate
  def testProcessTimeInnerJoin(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    val tEnv = StreamTableEnvironment.create(env, TableTestUtil.STREAM_SETTING)
    env.setParallelism(1)

    val sqlQuery =
      """
        |SELECT
        |  o.amount * r.rate AS amount
        |FROM
        |  Orders AS o,
        |  LATERAL TABLE (Rates(o.proctime)) AS r
        |WHERE r.currency = o.currency
        |""".stripMargin

    val ordersData = new mutable.MutableList[(Long, String)]
    ordersData.+=((2L, "Euro"))
    ordersData.+=((1L, "US Dollar"))
    ordersData.+=((50L, "Yen"))
    ordersData.+=((3L, "Euro"))
    ordersData.+=((5L, "US Dollar"))

    val ratesHistoryData = new mutable.MutableList[(String, Long)]
    ratesHistoryData.+=(("US Dollar", 102L))
    ratesHistoryData.+=(("Euro", 114L))
    ratesHistoryData.+=(("Yen", 1L))
    ratesHistoryData.+=(("Euro", 116L))
    ratesHistoryData.+=(("Euro", 119L))

    val orders = StreamingEnvUtil
      .fromCollection(env, ordersData)
      .toTable(tEnv, 'amount, 'currency, 'proctime.proctime)
    val ratesHistory = StreamingEnvUtil
      .fromCollection(env, ratesHistoryData)
      .toTable(tEnv, 'currency, 'rate, 'proctime.proctime)

    tEnv.createTemporaryView("Orders", orders)
    tEnv.createTemporaryView("RatesHistory", ratesHistory)
    tEnv.createTemporarySystemFunction(
      "Rates",
      ratesHistory.createTemporalTableFunction($"proctime", $"currency"))

    val result = tEnv.sqlQuery(sqlQuery).toDataStream
    result.addSink(new TestingAppendSink)
    env.execute()
  }

  @TestTemplate
  def testProcessTimeInnerJoinWithConstantTable(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    val tEnv = StreamTableEnvironment.create(env, TableTestUtil.STREAM_SETTING)
    env.setParallelism(1)
    val result = tEnv
      .sqlQuery(
        "SELECT amount, currency, proctime() as proctime " +
          "FROM (VALUES (1, 2.0)) AS T(amount, currency)")
      .toDataStream
    result.addSink(new TestingAppendSink)
    env.execute()
  }

  @TestTemplate
  def testProcessTimeInnerJoinUnionAll(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    val tEnv = StreamTableEnvironment.create(env, TableTestUtil.STREAM_SETTING)
    env.setParallelism(1)

    val sqlQuery =
      """
        |SELECT
        |  o.amount * r.rate AS amount
        |FROM
        |  Orders AS o,
        |  LATERAL TABLE (Rates(o.proctime)) AS r
        |WHERE r.currency = o.currency
        |""".stripMargin

    val ordersData = new mutable.MutableList[(Long, String)]
    ordersData.+=((2L, "Euro"))
    ordersData.+=((1L, "US Dollar"))
    ordersData.+=((50L, "Yen"))
    ordersData.+=((3L, "Euro"))
    ordersData.+=((5L, "US Dollar"))

    val ratesHistoryData = new mutable.MutableList[(String, Long)]
    ratesHistoryData.+=(("US Dollar", 102L))
    ratesHistoryData.+=(("Euro", 114L))
    ratesHistoryData.+=(("Yen", 1L))
    ratesHistoryData.+=(("Euro", 116L))
    ratesHistoryData.+=(("Euro", 119L))

    val orders1 = StreamingEnvUtil
      .fromCollection(env, ordersData)
      .toTable(tEnv, 'amount, 'currency, 'proctime.proctime)
    val orders2 = StreamingEnvUtil
      .fromCollection(env, ordersData)
      .toTable(tEnv, 'amount, 'currency, 'proctime.proctime)
    val ratesHistory = StreamingEnvUtil
      .fromCollection(env, ratesHistoryData)
      .toTable(tEnv, 'currency, 'rate, 'proctime.proctime)

    tEnv.createTemporaryView("Orders1", orders1)
    tEnv.createTemporaryView("Orders2", orders2)
    tEnv.createTemporaryView("RatesHistory", ratesHistory)

    tEnv.createTemporaryFunction(
      "Rates",
      ratesHistory.createTemporalTableFunction($"proctime", $"currency"))
    tEnv.createTemporaryView(
      "Orders",
      tEnv.sqlQuery("SELECT * FROM Orders1 UNION ALL SELECT * FROM Orders2"))
    val result = tEnv.sqlQuery(sqlQuery).toDataStream
    result.addSink(new TestingAppendSink)
    env.execute()
  }

  @TestTemplate
  def testEventTimeInnerJoin(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    val tEnv = StreamTableEnvironment.create(env, TableTestUtil.STREAM_SETTING)
    env.setParallelism(1)

    val sqlQuery =
      """
        |SELECT
        |  o.amount * r.rate AS amount
        |FROM
        |  Orders AS o,
        |  LATERAL TABLE (Rates(o.rowtime)) AS r
        |WHERE r.currency = o.currency
        |""".stripMargin

    val ordersData = new mutable.MutableList[(Long, String, Timestamp)]
    ordersData.+=((2L, "Euro", new Timestamp(2L)))
    ordersData.+=((1L, "US Dollar", new Timestamp(3L)))
    ordersData.+=((50L, "Yen", new Timestamp(4L)))
    ordersData.+=((3L, "Euro", new Timestamp(5L)))

    val ratesHistoryData = new mutable.MutableList[(String, Long, Timestamp)]
    ratesHistoryData.+=(("US Dollar", 102L, new Timestamp(1L)))
    ratesHistoryData.+=(("Euro", 114L, new Timestamp(1L)))
    ratesHistoryData.+=(("Yen", 1L, new Timestamp(1L)))
    ratesHistoryData.+=(("Euro", 116L, new Timestamp(5L)))
    ratesHistoryData.+=(("Euro", 119L, new Timestamp(7L)))

    var expectedOutput = new mutable.HashSet[String]()
    expectedOutput += (2 * 114).toString
    expectedOutput += (3 * 116).toString

    val orders = StreamingEnvUtil
      .fromCollection(env, ordersData)
      .assignTimestampsAndWatermarks(new TimestampExtractor[(Long, String, Timestamp)]())
      .toTable(tEnv, 'amount, 'currency, 'rowtime.rowtime)
    val ratesHistory = StreamingEnvUtil
      .fromCollection(env, ratesHistoryData)
      .assignTimestampsAndWatermarks(new TimestampExtractor[(String, Long, Timestamp)]())
      .toTable(tEnv, 'currency, 'rate, 'rowtime.rowtime)

    tEnv.createTemporaryView("Orders", orders)
    tEnv.createTemporaryView("RatesHistory", ratesHistory)
    tEnv.createTemporaryView(
      "FilteredRatesHistory",
      tEnv.sqlQuery("SELECT * FROM RatesHistory WHERE rate > 110"))
    tEnv.createTemporarySystemFunction(
      "Rates",
      tEnv
        .from("FilteredRatesHistory")
        .createTemporalTableFunction($"rowtime", $"currency"))
    tEnv.createTemporaryView("TemporalJoinResult", tEnv.sqlQuery(sqlQuery))

    // Scan from registered table to test for interplay between
    // LogicalCorrelateToTemporalTableJoinRule and TableScanRule
    val result = tEnv.from("TemporalJoinResult").toDataStream
    val sink = new TestingAppendSink
    result.addSink(sink)
    env.execute()

    assertThat(sink.getAppendResults.toSet).isEqualTo(expectedOutput)
  }

  @TestTemplate
  def testNestedTemporalJoin(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    val tEnv = StreamTableEnvironment.create(env, TableTestUtil.STREAM_SETTING)

    val sqlQuery =
      """
        |SELECT
        |  o.orderId,
        |  (o.amount * p.price * r.rate) as total_price
        |FROM
        |  Orders AS o,
        |  LATERAL TABLE (Prices(o.rowtime)) AS p,
        |  LATERAL TABLE (Rates(o.rowtime)) AS r
        |WHERE
        |  o.productId = p.productId AND
        |  r.currency = p.currency
        |""".stripMargin

    val ordersData = new mutable.MutableList[(Long, String, Long, Timestamp)]
    ordersData.+=((1L, "A1", 2L, new Timestamp(2L)))
    ordersData.+=((2L, "A2", 1L, new Timestamp(3L)))
    ordersData.+=((3L, "A4", 50L, new Timestamp(4L)))
    ordersData.+=((4L, "A1", 3L, new Timestamp(5L)))
    val orders = StreamingEnvUtil
      .fromCollection(env, ordersData)
      .assignTimestampsAndWatermarks(new TimestampExtractor[(Long, String, Long, Timestamp)]())
      .toTable(tEnv, 'orderId, 'productId, 'amount, 'rowtime.rowtime)

    val ratesHistoryData = new mutable.MutableList[(String, Long, Timestamp)]
    ratesHistoryData.+=(("US Dollar", 102L, new Timestamp(1L)))
    ratesHistoryData.+=(("Euro", 114L, new Timestamp(1L)))
    ratesHistoryData.+=(("Yen", 1L, new Timestamp(1L)))
    ratesHistoryData.+=(("Euro", 116L, new Timestamp(5L)))
    ratesHistoryData.+=(("Euro", 119L, new Timestamp(7L)))
    val ratesHistory = StreamingEnvUtil
      .fromCollection(env, ratesHistoryData)
      .assignTimestampsAndWatermarks(new TimestampExtractor[(String, Long, Timestamp)]())
      .toTable(tEnv, 'currency, 'rate, 'rowtime.rowtime)

    val pricesHistoryData = new mutable.MutableList[(String, String, Double, Timestamp)]
    pricesHistoryData.+=(("A2", "US Dollar", 10.2d, new Timestamp(1L)))
    pricesHistoryData.+=(("A1", "Euro", 11.4d, new Timestamp(1L)))
    pricesHistoryData.+=(("A4", "Yen", 1d, new Timestamp(1L)))
    pricesHistoryData.+=(("A1", "Euro", 11.6d, new Timestamp(5L)))
    pricesHistoryData.+=(("A1", "Euro", 11.9d, new Timestamp(7L)))
    val pricesHistory = StreamingEnvUtil
      .fromCollection(env, pricesHistoryData)
      .assignTimestampsAndWatermarks(new TimestampExtractor[(String, String, Double, Timestamp)]())
      .toTable(tEnv, 'productId, 'currency, 'price, 'rowtime.rowtime)

    tEnv.createTemporaryView("Orders", orders)
    tEnv.createTemporaryView("RatesHistory", ratesHistory)
    tEnv.createTemporarySystemFunction(
      "Rates",
      ratesHistory.createTemporalTableFunction($"rowtime", $"currency"))
    tEnv.createTemporarySystemFunction(
      "Prices",
      pricesHistory.createTemporalTableFunction($"rowtime", $"productId"))

    tEnv.createTemporaryView("TemporalJoinResult", tEnv.sqlQuery(sqlQuery))

    // Scan from registered table to test for interplay between
    // LogicalCorrelateToTemporalTableJoinRule and TableScanRule
    val result = tEnv.from("TemporalJoinResult").toDataStream
    val sink = new TestingAppendSink
    result.addSink(sink)
    env.execute()

    val expected = List(
      s"1,${2 * 114 * 11.4}",
      s"2,${1 * 102 * 10.2}",
      s"3,${50 * 1 * 1.0}",
      s"4,${3 * 116 * 11.6}")
    assertThat(sink.getAppendResults.sorted).isEqualTo(expected.sorted)
  }
}

class TimestampExtractor[T <: Product]
  extends BoundedOutOfOrdernessTimestampExtractor[T](Duration.ofSeconds(10)) {
  override def extractTimestamp(element: T): Long = element match {
    case (_, _, ts: Timestamp) => ts.getTime
    case (_, _, _, ts: Timestamp) => ts.getTime
    case _ =>
      throw new IllegalArgumentException(
        "Expected the last element in a tuple to be of a Timestamp type.")
  }
}
