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

import java.time.{Duration, Period, Year}

import org.apache.spark.sql.YearUDT
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils
import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog
import org.apache.spark.sql.execution.HiveResult._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession}
import org.apache.spark.sql.types.{YearMonthIntervalType, YearMonthIntervalType => YM}


class HiveResultSuite extends SharedSparkSession {
  import testImplicits._

  test("date formatting in hive result") {
    DateTimeTestUtils.outstandingTimezonesIds.foreach { zoneId =>
      withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> zoneId) {
        val dates = Seq("2018-12-28", "1582-10-03", "1582-10-04", "1582-10-15")
        val df = dates.toDF("a").selectExpr("cast(a as date) as b")
        val executedPlan1 = df.queryExecution.executedPlan
        val result = hiveResultString(executedPlan1)
        assert(result == dates)
        val executedPlan2 = df.selectExpr("array(b)").queryExecution.executedPlan
        val result2 = hiveResultString(executedPlan2)
        assert(result2 == dates.map(x => s"[$x]"))
      }
    }
  }

  test("time formatting in hive result") {
    val times = Seq(
      "00:00:00",
      "00:01:02.003004",
      "12:13:14.999999",
      "23:59:59.1234")
    val df = times.toDF("a").selectExpr("cast(a as time) as b")
    val executedPlan1 = df.queryExecution.executedPlan
    val result = hiveResultString(executedPlan1)
    assert(result == times)
    val executedPlan2 = df.selectExpr("array(b)").queryExecution.executedPlan
    val result2 = hiveResultString(executedPlan2)
    assert(result2 == times.map(x => s"[$x]"))
  }

  test("timestamp formatting in hive result") {
    val timestamps = Seq(
      "2018-12-28 01:02:03",
      "1582-10-03 01:02:03",
      "1582-10-04 01:02:03",
      "1582-10-15 01:02:03")
    val df = timestamps.toDF("a").selectExpr("cast(a as timestamp) as b")
    val executedPlan1 = df.queryExecution.executedPlan
    val result = hiveResultString(executedPlan1)
    assert(result == timestamps)
    val executedPlan2 = df.selectExpr("array(b)").queryExecution.executedPlan
    val result2 = hiveResultString(executedPlan2)
    assert(result2 == timestamps.map(x => s"[$x]"))
  }

  test("toHiveString correctly handles UDTs") {
    val point = new ExamplePoint(50.0, 50.0)
    val tpe = new ExamplePointUDT()
    assert(toHiveString((point, tpe), false, getTimeFormatters, getBinaryFormatter) ===
      "(50.0, 50.0)")
  }

  test("decimal formatting in hive result") {
    val df = Seq(new java.math.BigDecimal("1")).toDS()
    Seq(2, 6, 18).foreach { scala =>
      val executedPlan =
        df.selectExpr(s"CAST(value AS decimal(38, $scala))").queryExecution.executedPlan
      val result = hiveResultString(executedPlan)
      assert(result.head.split("\\.").last.length === scala)
    }

    val executedPlan = Seq(java.math.BigDecimal.ZERO).toDS()
      .selectExpr(s"CAST(value AS decimal(38, 8))").queryExecution.executedPlan
    val result = hiveResultString(executedPlan)
    assert(result.head === "0.00000000")
  }

  test("SHOW TABLES in hive result") {
    withSQLConf("spark.sql.catalog.testcat" -> classOf[InMemoryTableCatalog].getName) {
      Seq(("testcat.ns", "tbl", "foo"), ("spark_catalog.default", "tbl", "csv")).foreach {
        case (ns, tbl, source) =>
          withTable(s"$ns.$tbl") {
            spark.sql(s"CREATE TABLE $ns.$tbl (id bigint) USING $source")
            val df = spark.sql(s"SHOW TABLES FROM $ns")
            val executedPlan = df.queryExecution.executedPlan
            assert(hiveResultString(executedPlan).head == tbl)
          }
      }
    }
  }

  test("DESCRIBE TABLE in hive result") {
    withSQLConf("spark.sql.catalog.testcat" -> classOf[InMemoryTableCatalog].getName) {
      Seq(("testcat.ns", "tbl", "foo"), ("spark_catalog.default", "tbl", "csv")).foreach {
        case (ns, tbl, source) =>
          withTable(s"$ns.$tbl") {
            spark.sql(s"CREATE TABLE $ns.$tbl (id bigint COMMENT 'col1') USING $source")
            val df = spark.sql(s"DESCRIBE $ns.$tbl")
            val executedPlan = df.queryExecution.executedPlan
            val expected = "id                  " +
              "\tbigint              " +
              "\tcol1                "
            assert(hiveResultString(executedPlan).head == expected)
          }
      }
    }
  }

  test("SPARK-34984, SPARK-35016: year-month interval formatting in hive result") {
    val df = Seq(Period.ofYears(-10).minusMonths(1)).toDF("i")
    val plan1 = df.queryExecution.executedPlan
    assert(hiveResultString(plan1) === Seq("-10-1"))
    val plan2 = df.selectExpr("array(i)").queryExecution.executedPlan
    assert(hiveResultString(plan2) === Seq("[-10-1]"))
  }

  test("SPARK-49208: negative month intervals") {
    Seq(
      "0-0" -> (11, YM.YEAR, YM.YEAR),
      "0-0" -> (-11, YM.YEAR, YM.YEAR),
      "0-11" -> (11, YM.YEAR, YM.MONTH),
      "-0-11" -> (-11, YM.YEAR, YM.MONTH),
      "0-11" -> (11, YM.MONTH, YM.MONTH),
      "-0-11" -> (-11, YM.MONTH, YM.MONTH),
      "1-0" -> (12, YM.YEAR, YM.YEAR),
      "-1-0" -> (-12, YM.YEAR, YM.YEAR),
      "1-0" -> (12, YM.YEAR, YM.MONTH),
      "-1-0" -> (-12, YM.YEAR, YM.MONTH),
      "1-0" -> (12, YM.MONTH, YM.MONTH),
      "-1-0" -> (-12, YM.MONTH, YM.MONTH),
      "1-0" -> (13, YM.YEAR, YM.YEAR),
      "-1-0" -> (-13, YM.YEAR, YM.YEAR),
      "1-1" -> (13, YM.YEAR, YM.MONTH),
      "-1-1" -> (-13, YM.YEAR, YM.MONTH),
      "1-1" -> (13, YM.MONTH, YM.MONTH),
      "-1-1" -> (-13, YM.MONTH, YM.MONTH)
    ).foreach { case (hiveString, (months, startField, endField)) =>
      assert(toHiveString((Period.ofMonths(months), YearMonthIntervalType(startField, endField)),
        false,
        getTimeFormatters,
        getBinaryFormatter) === hiveString)
    }
  }

  test("SPARK-34984, SPARK-35016: day-time interval formatting in hive result") {
    val df = Seq(Duration.ofDays(5).plusMillis(10)).toDF("i")
    val plan1 = df.queryExecution.executedPlan
    assert(hiveResultString(plan1) === Seq("5 00:00:00.010000000"))
    val plan2 = df.selectExpr("array(i)").queryExecution.executedPlan
    assert(hiveResultString(plan2) === Seq("[5 00:00:00.010000000]"))
  }

  test("geometry formatting in hive result") {
    // Geometry with no SRID.
    val df = spark.sql(
      "SELECT ST_GeomFromWKB(X'0101000000000000000000f03f0000000000000040')")
    val result = hiveResultString(df.queryExecution.executedPlan)
    assert(result === Seq("POINT(1 2)"))
    // Geometry with SRID 0.
    val dfSrid0 = spark.sql(
      "SELECT ST_GeomFromWKB(X'0101000000000000000000f03f0000000000000040', 0)")
    val resultSrid0 = hiveResultString(dfSrid0.queryExecution.executedPlan)
    assert(resultSrid0 === Seq("POINT(1 2)"))
    // Geometry with SRID 4326.
    val dfSrid4326 = spark.sql(
      "SELECT ST_GeomFromWKB(X'0101000000000000000000f03f0000000000000040', 4326)")
    val resultSrid4326 = hiveResultString(dfSrid4326.queryExecution.executedPlan)
    assert(resultSrid4326 === Seq("SRID=4326;POINT(1 2)"))
  }

  test("nested geometry formatting in hive result") {
    // Geometry with no SRID.
    val df = spark.sql(
      "SELECT array(ST_GeomFromWKB(X'0101000000000000000000f03f0000000000000040'))")
    val result = hiveResultString(df.queryExecution.executedPlan)
    assert(result === Seq("[\"POINT(1 2)\"]"))
    // Geometry with SRID 4326.
    val dfSrid4326 = spark.sql(
      "SELECT array(ST_GeomFromWKB(X'0101000000000000000000f03f0000000000000040', 4326))")
    val resultSrid4326 = hiveResultString(dfSrid4326.queryExecution.executedPlan)
    assert(resultSrid4326 === Seq("[\"SRID=4326;POINT(1 2)\"]"))
  }

  test("geography formatting in hive result") {
    val df = spark.sql(
      "SELECT ST_GeogFromWKB(X'0101000000000000000000f03f0000000000000040')")
    val result = hiveResultString(df.queryExecution.executedPlan)
    assert(result === Seq("SRID=4326;POINT(1 2)"))
  }

  test("nested geography formatting in hive result") {
    val df = spark.sql(
      "SELECT array(ST_GeogFromWKB(X'0101000000000000000000f03f0000000000000040'))")
    val result = hiveResultString(df.queryExecution.executedPlan)
    assert(result === Seq("[\"SRID=4326;POINT(1 2)\"]"))
  }

  test("SPARK-52650: Use stringifyValue to get UDT string representation") {
    val year = Year.of(18)
    val tpe = new YearUDT()
    assert(toHiveString((year, tpe),
      nested = false, getTimeFormatters, getBinaryFormatter) === "18")
    val tpe2 = new YearUDT() {
      override def stringifyValue(obj: Any): String = {
        f"${obj.asInstanceOf[Year].getValue}%04d"
      }
    }
    assert(toHiveString((year, tpe2),
      nested = false, getTimeFormatters, getBinaryFormatter) === "0018")
  }
}
