/*
 * 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.paimon.spark.sql

import org.apache.paimon.spark.PaimonSparkTestWithRestCatalogBase

import org.apache.spark.sql.Row

abstract class PaimonSQLFunctionTestBase extends PaimonSparkTestWithRestCatalogBase {

  test("Paimon SQL Function: scalar function with expression body") {
    withUserDefinedFunction("area" -> false) {
      sql("""
            |CREATE FUNCTION area(width DOUBLE, height DOUBLE)
            |RETURNS DOUBLE
            |RETURN width * height
            |""".stripMargin)
      checkAnswer(sql("SELECT area(2, 3)"), Row(6.0))

      withTable("t") {
        sql("CREATE TABLE t (a DOUBLE, b DOUBLE)")
        sql("INSERT INTO t VALUES (1, 2), (3, 4)")
        checkAnswer(sql("SELECT area(a, b) FROM t ORDER BY a"), Seq(Row(2.0), Row(12.0)))
      }
    }
  }

  test("Paimon SQL Function: create with qualified name and call across catalog") {
    withUserDefinedFunction("area" -> false) {
      sql("""
            |CREATE FUNCTION test.area(width DOUBLE, height DOUBLE)
            |RETURNS DOUBLE
            |RETURN width * height
            |""".stripMargin)
      checkAnswer(sql("SELECT test.area(2, 3)"), Row(6.0))
      checkAnswer(sql("SELECT paimon.test.area(2, 4)"), Row(8.0))
    }
  }

  test("Paimon SQL Function: create or replace / if not exists") {
    withUserDefinedFunction("inc" -> false) {
      sql("CREATE FUNCTION inc(x INT) RETURNS INT RETURN x + 1")
      checkAnswer(sql("SELECT inc(10)"), Row(11))

      // create again should fail
      intercept[Exception] {
        sql("CREATE FUNCTION inc(x INT) RETURNS INT RETURN x + 1")
      }

      // if not exists: no-op, keeps the old definition
      sql("CREATE FUNCTION IF NOT EXISTS inc(x INT) RETURNS INT RETURN x + 100")
      checkAnswer(sql("SELECT inc(10)"), Row(11))

      // or replace: new definition takes effect
      sql("CREATE OR REPLACE FUNCTION inc(x INT) RETURNS INT RETURN x + 100")
      checkAnswer(sql("SELECT inc(10)"), Row(110))
    }
  }

  test("Paimon SQL Function: scalar function with query body referencing a table") {
    withUserDefinedFunction("dept_total" -> false) {
      withTable("emp") {
        sql("CREATE TABLE emp (id INT, dept_id INT, salary INT)")
        sql("INSERT INTO emp VALUES (1, 10, 100), (2, 10, 200), (3, 20, 300)")
        sql("""
              |CREATE FUNCTION dept_total(d INT) RETURNS INT
              |RETURN SELECT SUM(salary) FROM emp WHERE dept_id = d
              |""".stripMargin)
        checkAnswer(sql("SELECT dept_total(10)"), Row(300))
        checkAnswer(sql("SELECT dept_total(20)"), Row(300))
      }
    }
  }

  test("Paimon SQL Function: parameter with DEFAULT value") {
    withUserDefinedFunction("addd" -> false) {
      sql("CREATE FUNCTION addd(x INT, y INT DEFAULT 10) RETURNS INT RETURN x + y")
      checkAnswer(sql("SELECT addd(5)"), Row(15))
      checkAnswer(sql("SELECT addd(5, 1)"), Row(6))
    }
  }

  test("Paimon SQL Function: describe function") {
    withUserDefinedFunction("area" -> false) {
      sql("CREATE FUNCTION area(width DOUBLE, height DOUBLE) RETURNS DOUBLE RETURN width * height")

      val desc = sql("DESCRIBE FUNCTION area").collect().map(_.getString(0))
      assert(desc.exists(_.contains("SCALAR")), desc.mkString("\n"))
      assert(desc.exists(_.contains("Input")), desc.mkString("\n"))
      assert(desc.exists(_.contains("width")), desc.mkString("\n"))
      assert(desc.exists(_.contains("DOUBLE")), desc.mkString("\n"))

      val descExt = sql("DESCRIBE FUNCTION EXTENDED area").collect().map(_.getString(0))
      assert(descExt.exists(_.contains("Deterministic")), descExt.mkString("\n"))
      assert(descExt.exists(_.contains("width * height")), descExt.mkString("\n"))
    }
  }

  test("Paimon SQL Function: describe function with comment") {
    withUserDefinedFunction("inc" -> false) {
      sql("CREATE FUNCTION inc(x INT) RETURNS INT COMMENT 'increment by one' RETURN x + 1")

      val desc = sql("DESCRIBE FUNCTION inc").collect().map(_.getString(0))
      assert(desc.exists(_.contains("SCALAR")), desc.mkString("\n"))

      val descExt = sql("DESCRIBE FUNCTION EXTENDED inc").collect().map(_.getString(0))
      assert(descExt.exists(_.contains("increment by one")), descExt.mkString("\n"))
      assert(descExt.exists(_.contains("x + 1")), descExt.mkString("\n"))
    }
  }

  test("Paimon SQL Function: show functions lists the created function") {
    withUserDefinedFunction("area" -> false) {
      sql("CREATE FUNCTION area(w DOUBLE, h DOUBLE) RETURNS DOUBLE RETURN w * h")
      val names = sql("SHOW USER FUNCTIONS").collect().map(_.getString(0).toLowerCase)
      assert(names.exists(_.contains("area")), names.mkString(", "))
    }
  }

  test("Paimon SQL Function: drop function") {
    withUserDefinedFunction("area" -> false) {
      sql("CREATE FUNCTION area(w DOUBLE, h DOUBLE) RETURNS DOUBLE RETURN w * h")
      checkAnswer(sql("SELECT area(2, 3)"), Row(6.0))

      sql("DROP FUNCTION area")
      intercept[Exception] {
        sql("SELECT area(2, 3)")
      }
      sql("DROP FUNCTION IF EXISTS area")
    }
  }

  test("Paimon SQL Function: SQL configs captured at creation time") {
    assume(gteqSpark4_1)
    withUserDefinedFunction("div_func" -> false) {
      // Create with ANSI enabled — division by zero should throw at query time.
      sql("SET spark.sql.ansi.enabled=true")
      sql("CREATE FUNCTION div_func(x INT) RETURNS DOUBLE RETURN 1 / x")
      sql("SET spark.sql.ansi.enabled=false")

      // Even though ANSI is now disabled in the session, the function was created with ANSI=true,
      // so division by zero should still throw ArithmeticException.
      val e = intercept[Exception] {
        sql("SELECT div_func(0)").collect()
      }
      assert(
        e.getMessage.contains("Division by zero") ||
          e.getMessage.contains("ArithmeticException") ||
          e.getMessage.contains("DIVIDE_BY_ZERO"))

      sql("RESET spark.sql.ansi.enabled")
    }
  }

  test("Paimon SQL Function: non-deterministic function body") {
    assume(gteqSpark4_1)
    withUserDefinedFunction("rnd" -> false) {
      sql("CREATE FUNCTION rnd() RETURNS DOUBLE RETURN rand()")
      val r1 = sql("SELECT rnd()").collect()(0).getDouble(0)
      val r2 = sql("SELECT rnd()").collect()(0).getDouble(0)
      assert(r1 >= 0.0 && r1 < 1.0)
      assert(r2 >= 0.0 && r2 < 1.0)
    }
  }

  test("Paimon SQL Function: reject aggregate in scalar function body") {
    assume(gteqSpark4_1)
    val e = intercept[Exception] {
      sql("CREATE FUNCTION bad_agg(x INT) RETURNS INT RETURN SUM(x)")
    }
    assert(e.getMessage.contains("CANNOT_CONTAIN_COMPLEX_FUNCTIONS"))
  }

  test("Paimon SQL Function: reject window function in scalar function body") {
    assume(gteqSpark4_1)
    val e = intercept[Exception] {
      sql("CREATE FUNCTION bad_win(x INT) RETURNS INT RETURN ROW_NUMBER() OVER (ORDER BY x)")
    }
    assert(e.getMessage.contains("CANNOT_CONTAIN_COMPLEX_FUNCTIONS"))
  }

  test("Paimon SQL Function: reject duplicate parameter names") {
    assume(gteqSpark4_1)
    val e = intercept[Exception] {
      sql("CREATE FUNCTION bad_dup(x INT, x INT) RETURNS INT RETURN x + x")
    }
    assert(e.getMessage.toLowerCase.contains("duplicate"))
  }

  test("Paimon SQL Function: reject non-trailing defaults") {
    assume(gteqSpark4_1)
    val e = intercept[Exception] {
      sql("CREATE FUNCTION bad_def(x INT DEFAULT 1, y INT) RETURNS INT RETURN x + y")
    }
    assert(e.getMessage.toLowerCase.contains("default"))
  }

  test("Paimon SQL Function: omitting RETURNS clause") {
    assume(gteqSpark4_1)
    withUserDefinedFunction("inc" -> false) {
      sql("CREATE FUNCTION inc(x INT) RETURN x + 1")
      checkAnswer(sql("SELECT inc(10)"), Row(11))
    }
  }

  test("Paimon SQL Function: table function is not supported yet") {
    val e = intercept[Exception] {
      sql("""
            |CREATE FUNCTION rows_of(x INT)
            |RETURNS TABLE(a INT)
            |RETURN SELECT x + 1 AS a
            |""".stripMargin)
    }
    assert(e.getMessage.contains("does not support creating SQL table functions"))
  }
}
