/*
 * 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.codegen


import java.util.{ArrayList => JArrayList, Collection => JCollection}
import org.apache.calcite.rex.{RexLiteral, RexNode, RexProgram}
import org.apache.flink.api.common.functions.FlatMapFunction
import org.apache.flink.api.common.typeinfo.TypeInformation
import org.apache.flink.api.java.typeutils.RowTypeInfo
import org.apache.flink.streaming.api.functions.async.{AsyncFunction, ResultFuture}
import org.apache.flink.table.api.scala._
import org.apache.flink.table.api.{TableConfig, TableConfigOptions}
import org.apache.flink.table.api.functions.{AsyncTableFunction, TableFunction}
import org.apache.flink.table.api.types.{DataTypes, InternalType, RowType}
import org.apache.flink.table.calcite.FlinkTypeFactory
import org.apache.flink.table.codegen.CodeGenUtils._
import org.apache.flink.table.dataformat.{BaseRow, GenericRow, JoinedRow}
import org.apache.flink.table.plan.schema.BaseRowSchema
import org.apache.flink.table.runtime.collector.TableFunctionCollector
import org.apache.flink.table.runtime.conversion.DataStructureConverters.RowConverter
import org.apache.flink.types.Row
import org.apache.flink.util.{Collector, Preconditions}

object TemporalJoinCodeGenerator {

  /**
    * Generates a lookup function ([[TableFunction]])
    */
  def generateLookupFunction(
    config: TableConfig,
    typeFactory: FlinkTypeFactory,
    inputType: InternalType,
    returnType: InternalType,
    tableReturnTypeInfo: TypeInformation[_],
    tableReturnClass: Class[_],
    lookupKeyInOrder: Array[Int],
    lookupKeysFromInput: Map[Int, Int], // lookup key index -> input field index
    lookupKeysFromConstant: Map[Int, RexLiteral],  // lookup key index -> constant value
    lookupFunction: TableFunction[_],
    enableObjectReuse: Boolean)
  : GeneratedFunction[FlatMapFunction[BaseRow, BaseRow], BaseRow] = {

    val ctx = CodeGeneratorContext(config)
    val (prepareCode, parameters) = prepareParameters(
      ctx,
      config,
      typeFactory,
      inputType,
      lookupKeyInOrder,
      lookupKeysFromInput,
      lookupKeysFromConstant,
      enableObjectReuse)

    val lookupFunctionTerm = ctx.addReusableFunction(lookupFunction)

    val setCollectorCode = if (tableReturnClass == classOf[Row]) {
      val converterCollector =
        new RowToBaseRowCollector(tableReturnTypeInfo.asInstanceOf[RowTypeInfo])
      val term = ctx.addReusableObject(converterCollector, "collector")
      s"""
         |$term.setCollector(${CodeGeneratorContext.DEFAULT_COLLECTOR_TERM});
         |$lookupFunctionTerm.setCollector($term);
         """.stripMargin
    } else {
      s"$lookupFunctionTerm.setCollector(${CodeGeneratorContext.DEFAULT_COLLECTOR_TERM});"
    }

    val body =
      s"""
        |$prepareCode
        |$setCollectorCode
        |$lookupFunctionTerm.eval($parameters);
      """.stripMargin

    FunctionCodeGenerator.generateFunction(
      ctx,
      "LookupFunction",
      classOf[FlatMapFunction[BaseRow, BaseRow]],
      body,
      returnType,
      inputType,
      config)
  }

  /**
    * Generates a async lookup function ([[AsyncTableFunction]])
    */
  def generateAsyncLookupFunction(
    config: TableConfig,
    typeFactory: FlinkTypeFactory,
    inputType: InternalType,
    returnType: InternalType,
    tableReturnTypeInfo: TypeInformation[_],
    tableReturnClass: Class[_],
    lookupKeyInOrder: Array[Int],
    lookupKeysFromInput: Map[Int, Int], // lookup key index -> input field index
    lookupKeysFromConstant: Map[Int, RexLiteral],
    asyncLookupFunction: AsyncTableFunction[_])
  : GeneratedFunction[AsyncFunction[BaseRow, BaseRow], BaseRow] = {

    val ctx = CodeGeneratorContext(config)
    val (prepareCode, parameters) = prepareParameters(
      ctx,
      config,
      typeFactory,
      inputType,
      lookupKeyInOrder,
      lookupKeysFromInput,
      lookupKeysFromConstant,
      fieldCopy = true) // always copy input field because of async buffer

    val lookupFunctionTerm = ctx.addReusableFunction(asyncLookupFunction)

    var futureTerm: String = null
    val setFutureCode = if (tableReturnClass == classOf[Row]) {
      val converterFuture =
        new RowToBaseRowResultFuture(tableReturnTypeInfo.asInstanceOf[RowTypeInfo])
      futureTerm = ctx.addReusableObject(converterFuture, "future")
      s"$futureTerm.setFuture(${CodeGeneratorContext.DEFAULT_COLLECTOR_TERM});"
    } else {
      futureTerm = CodeGeneratorContext.DEFAULT_COLLECTOR_TERM
      ""
    }

    val body =
      s"""
         |$prepareCode
         |$setFutureCode
         |$lookupFunctionTerm.eval($futureTerm, $parameters);
      """.stripMargin

    FunctionCodeGenerator.generateFunction(
      ctx,
      "LookupFunction",
      classOf[AsyncFunction[BaseRow, BaseRow]],
      body,
      returnType,
      inputType,
      config)
  }

  private def prepareParameters(
    ctx: CodeGeneratorContext,
    config: TableConfig,
    typeFactory: FlinkTypeFactory,
    inputType: InternalType,
    lookupKeyInOrder: Array[Int],
    lookupKeysFromInput: Map[Int, Int], // lookup key index -> input field index
    lookupKeysFromConstant: Map[Int, RexLiteral],
    fieldCopy: Boolean): (String, String) = {

    // the total number of lookupKeys should equal to fromInput plus fromConstant
    Preconditions.checkArgument(
      lookupKeyInOrder.length == lookupKeysFromInput.size + lookupKeysFromConstant.size)
    val inputFieldExprs = for (i <- lookupKeyInOrder) yield {
      if (lookupKeysFromInput.contains(i)) {
        generateInputAccess(
          ctx,
          inputType,
          CodeGeneratorContext.DEFAULT_INPUT1_TERM,
          lookupKeysFromInput(i),
          nullableInput = false,
          config.getNullCheck,
          fieldCopy)
      } else if (lookupKeysFromConstant.contains(i)) {
        val literal = lookupKeysFromConstant(i)
        val resultType = FlinkTypeFactory.toInternalType(literal.getType)
        val value = literal.getValue3
        generateLiteral(ctx, literal.getType, resultType, value, config.getNullCheck)
      } else {
        throw new CodeGenException("This should never happen!")
      }
    }
    val codeAndArg = inputFieldExprs
        .map { e =>
          val bType = boxedTypeTermForType(e.resultType)
          val newTerm = newName("arg")
          val code =
            s"""
               |$bType $newTerm = null;
               |if (!${e.nullTerm}) {
               |  $newTerm = ${e.resultTerm};
               |}
             """.stripMargin
          (code, newTerm)
        }
    (codeAndArg.map(_._1).mkString("\n"), codeAndArg.map(_._2).mkString(", "))
  }

  /**
    * Generates async collector for async temporal join ([[ResultFuture]])
    */
  def generateAsyncCollector(
    config: TableConfig,
    inputType: RowType,
    tableType: RowType,
    joinCondition: Option[RexNode]): GeneratedCollector = {

    val inputTerm = CodeGeneratorContext.DEFAULT_INPUT1_TERM
    val tableInputTerm = CodeGeneratorContext.DEFAULT_INPUT2_TERM

    val ctx = CodeGeneratorContext(config)
    val exprGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck)
      .bindInput(tableType, inputTerm = tableInputTerm)

    val tableResultExpr = exprGenerator.generateConverterResultExpression(
      tableType, classOf[GenericRow])

    val body =
      s"""
         |${tableResultExpr.code}
         |getCollector().complete(java.util.Collections.singleton(${tableResultExpr.resultTerm}));
      """.stripMargin

    val collectorCode = if (joinCondition.isEmpty) {
      body
    } else {

      val filterGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck)
        .bindInput(inputType, inputTerm)
        .bindSecondInput(tableType, tableInputTerm)
      val filterCondition = filterGenerator.generateExpression(joinCondition.get)

      s"""
         |${filterCondition.code}
         |if (${filterCondition.resultTerm}) {
         |  $body
         |} else {
         |  getCollector().complete(java.util.Collections.emptyList());
         |}
         |""".stripMargin
    }

    CollectorCodeGenerator.generateTableAsyncCollector(
      ctx,
      "TableAsyncCollector",
      collectorCode,
      inputType,
      tableType,
      config)
  }

  /**
    * Generates collector for temporal join ([[Collector]])
    *
    * Differs from CommonCorrelate.generateCollector which has no real condition because of
    * FLINK-7865, here we should deal with outer join type when real conditions filtered result.
    */
  def generateCollector(
    ctx: CodeGeneratorContext,
    config: TableConfig,
    inputType: RowType,
    udtfTypeInfo: RowType,
    resultType: RowType,
    condition: Option[RexNode],
    pojoFieldMapping: Option[Array[Int]],
    retainHeader: Boolean = true): GeneratedCollector = {

    val inputTerm = CodeGeneratorContext.DEFAULT_INPUT1_TERM
    val udtfInputTerm = CodeGeneratorContext.DEFAULT_INPUT2_TERM

    val exprGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck)
      .bindInput(udtfTypeInfo, inputTerm = udtfInputTerm, inputFieldMapping = pojoFieldMapping)

    val udtfResultExpr = exprGenerator.generateConverterResultExpression(
      udtfTypeInfo, classOf[GenericRow])

    val joinedRowTerm = CodeGenUtils.newName("joinedRow")
    ctx.addOutputRecord(resultType, classOf[JoinedRow], joinedRowTerm)

    val header = if (retainHeader) {
      s"$joinedRowTerm.setHeader($inputTerm.getHeader());"
    } else {
      ""
    }

    val body =
      s"""
         |${udtfResultExpr.code}
         |$joinedRowTerm.replace($inputTerm, ${udtfResultExpr.resultTerm});
         |$header
         |getCollector().collect($joinedRowTerm);
         |super.collect(record);
      """.stripMargin

    val collectorCode = if (condition.isEmpty) {
      body
    } else {

      val filterGenerator = new ExprCodeGenerator(ctx, false, config.getNullCheck)
        .bindInput(inputType, inputTerm)
        .bindSecondInput(udtfTypeInfo, udtfInputTerm, pojoFieldMapping)
      val filterCondition = filterGenerator.generateExpression(condition.get)

      s"""
         |${filterCondition.code}
         |if (${filterCondition.resultTerm}) {
         |  $body
         |}
         |""".stripMargin
    }

    generateTableFunctionCollectorForJoinTable(
      ctx,
      "JoinTableFuncCollector",
      collectorCode,
      inputType,
      udtfTypeInfo,
      config,
      inputTerm = inputTerm,
      collectedTerm = udtfInputTerm)
  }

  /**
    * The only differences against CollectorCodeGenerator.generateTableFunctionCollector is
    * "super.collect" call is binding with collect join row in "body" code
    */
  private def generateTableFunctionCollectorForJoinTable(
    ctx: CodeGeneratorContext,
    name: String,
    bodyCode: String,
    inputType: RowType,
    collectedType: RowType,
    config: TableConfig,
    inputTerm: String = CodeGeneratorContext.DEFAULT_INPUT1_TERM,
    collectedTerm: String = CodeGeneratorContext.DEFAULT_INPUT2_TERM)
  : GeneratedCollector = {

    val className = newName(name)
    val input1TypeClass = boxedTypeTermForType(inputType)
    val input2TypeClass = boxedTypeTermForType(collectedType)

    val unboxingCodeSplit = generateSplitFunctionCalls(
      ctx.reusableInputUnboxingExprs.values.map(_.code).toSeq,
      config.getConf.getInteger(TableConfigOptions.SQL_CODEGEN_LENGTH_MAX),
      "inputUnbox",
      "private final void",
      ctx.reuseFieldCode().length,
      defineParams = s"$input1TypeClass $inputTerm, $input2TypeClass $collectedTerm",
      callingParams = s"$inputTerm, $collectedTerm"
    )


    val funcCode = if (unboxingCodeSplit.isSplit) {
      s"""
      public class $className extends ${classOf[TableFunctionCollector[_]].getCanonicalName} {

        ${ctx.reuseMemberCode()}
        ${ctx.reuseFieldCode()}

        public $className() throws Exception {
          ${ctx.reuseInitCode()}
        }

        @Override
        public void collect(Object record) throws Exception {
          $input1TypeClass $inputTerm = ($input1TypeClass) getInput();
          $input2TypeClass $collectedTerm = ($input2TypeClass) record;
          ${unboxingCodeSplit.callings.mkString("\n")}
          $bodyCode
        }

        ${unboxingCodeSplit.definitions.zip(unboxingCodeSplit.bodies) map {
        case (define, body) =>
          s"""
             |$define throws Exception {
             |  $body
             |}
          """.stripMargin
      } mkString "\n"
      }

        @Override
        public void close() {
        }
      }
    """.stripMargin
    } else {
      s"""
      public class $className extends ${classOf[TableFunctionCollector[_]].getCanonicalName} {

        ${ctx.reuseMemberCode()}

        public $className() throws Exception {
          ${ctx.reuseInitCode()}
        }

        @Override
        public void collect(Object record) throws Exception {
          $input1TypeClass $inputTerm = ($input1TypeClass) getInput();
          $input2TypeClass $collectedTerm = ($input2TypeClass) record;
          ${ctx.reuseFieldCode()}
          ${ctx.reuseInputUnboxingCode()}
          $bodyCode
        }

        @Override
        public void close() {
        }
      }
    """.stripMargin
    }

    GeneratedCollector(className, funcCode)
  }

  /**
    * Genrates calculate flatmap function for temporal join which is used
    * to projection/filter the dimension table results
    */
  def generateCalcMapFunction(
    config: TableConfig,
    calcProgram: Option[RexProgram],
    tableSourceSchema: BaseRowSchema)
  : GeneratedFunction[FlatMapFunction[BaseRow, BaseRow], BaseRow] = {

    val program = calcProgram.get
    val condition = if (program.getCondition != null) {
      Some(program.expandLocalRef(program.getCondition))
    } else {
      None
    }
    CalcCodeGenerator.generateFunction(
      tableSourceSchema.internalType,
      "TableCalcMapFunction",
      FlinkTypeFactory.toInternalRowType(program.getOutputRowType),
      classOf[GenericRow],
      program,
      condition,
      config,
      classOf[FlatMapFunction[BaseRow, BaseRow]])
  }


  // ----------------------------------------------------------------------------------------
  //                                Utility Classes
  // ----------------------------------------------------------------------------------------

  class RowToBaseRowCollector(rowTypeInfo: RowTypeInfo)
    extends TableFunctionCollector[Row] with Serializable {

    private val converter =
      RowConverter(rowTypeInfo.toInternalType.asInstanceOf[RowType])

    override def collect(record: Row): Unit = {
      super.collect(record)
      val result = converter.toInternalImpl(record)
      getCollector.asInstanceOf[Collector[BaseRow]].collect(result)
    }

    override def reset(): Unit = {
      super.reset()
      getCollector.asInstanceOf[TableFunctionCollector[_]].reset()
    }

    override def close(): Unit = getCollector.close()
  }

  class RowToBaseRowResultFuture(rowTypeInfo: RowTypeInfo)
    extends ResultFuture[Row] with Serializable {

    private val converter =
      RowConverter(rowTypeInfo.toInternalType.asInstanceOf[RowType])
    private var future: ResultFuture[BaseRow] = _

    def setFuture(future: ResultFuture[BaseRow]): Unit = {
      this.future = future
    }

    override def complete(result: JCollection[Row]): Unit = {
      if (result == null) {
        this.future.complete(null)
      } else {
        val baseRowResult = new JArrayList[BaseRow]
        val iter = result.iterator()
        while (iter.hasNext) {
          baseRowResult.add(converter.toInternalImpl(iter.next()))
        }
        this.future.complete(baseRowResult)
      }
    }

    override def completeExceptionally(error: Throwable): Unit = {
      this.future.completeExceptionally(error)
    }
  }
}