/*
 * 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.comet.shims

import java.io.FileNotFoundException

import scala.util.matching.Regex

import org.apache.spark.{QueryContext, SparkDateTimeException, SparkException}
import org.apache.spark.sql.catalyst.trees.SQLQueryContext
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String

object ShimSparkErrorConverter {
  val ObjectLocationPattern: Regex = "Object at location (.+?) not found".r
}

/**
 * Spark 3.4 implementation for converting error types to proper Spark exceptions.
 *
 * Handles all error cases using the Spark 3.4 QueryExecutionErrors API. Differs from the 3.5
 * version in 4 places where the API changed between 3.4 and 3.5.
 */
trait ShimSparkErrorConverter {

  private def sqlCtx(context: Array[QueryContext]): SQLQueryContext =
    context.headOption.map(_.asInstanceOf[SQLQueryContext]).getOrElse(null)

  private def parseFloatLiteral(value: String): Float = {
    value.toLowerCase match {
      case "inf" | "+inf" | "infinity" | "+infinity" => Float.PositiveInfinity
      case "-inf" | "-infinity" => Float.NegativeInfinity
      case "nan" | "+nan" | "-nan" => Float.NaN
      case _ => value.toFloat
    }
  }

  private def parseDoubleLiteral(value: String): Double = {
    val normalized = value.toLowerCase.stripSuffix("d")
    normalized match {
      case "inf" | "+inf" | "infinity" | "+infinity" => Double.PositiveInfinity
      case "-inf" | "-infinity" => Double.NegativeInfinity
      case "nan" | "+nan" | "-nan" => Double.NaN
      case _ => normalized.toDouble
    }
  }

  def convertErrorType(
      errorType: String,
      errorClass: String,
      params: Map[String, Any],
      context: Array[QueryContext],
      summary: String): Option[Throwable] = {
    val _ = (errorClass, summary)

    errorType match {

      case "DivideByZero" =>
        Some(QueryExecutionErrors.divideByZeroError(sqlCtx(context)))

      case "RemainderByZero" =>
        Some(QueryExecutionErrors.divideByZeroError(sqlCtx(context)))

      case "IntervalDividedByZero" =>
        Some(QueryExecutionErrors.intervalDividedByZeroError(sqlCtx(context)))

      case "BinaryArithmeticOverflow" =>
        // Spark 3.x does not take functionName parameter
        Some(
          QueryExecutionErrors.binaryArithmeticCauseOverflowError(
            params("value1").toString.toShort,
            params("symbol").toString,
            params("value2").toString.toShort))

      case "ArithmeticOverflow" =>
        val fromType = params("fromType").toString
        Some(
          QueryExecutionErrors
            .arithmeticOverflowError(fromType + " overflow", "", sqlCtx(context)))

      case "IntegralDivideOverflow" =>
        Some(QueryExecutionErrors.overflowInIntegralDivideError(sqlCtx(context)))

      case "DecimalSumOverflow" =>
        // Spark 3.x takes SQLQueryContext, not QueryContext
        Some(QueryExecutionErrors.overflowInSumOfDecimalError(sqlCtx(context)))

      case "NumericValueOutOfRange" =>
        val decimal = Decimal(params("value").toString)
        Some(
          QueryExecutionErrors.cannotChangeDecimalPrecisionError(
            decimal,
            params("precision").toString.toInt,
            params("scale").toString.toInt,
            sqlCtx(context)))

      case "DatetimeOverflow" =>
        Some(
          new SparkException(
            errorClass = "DATETIME_OVERFLOW",
            messageParameters = params.map { case (k, v) => (k, v.toString) },
            cause = null))

      case "InvalidArrayIndex" =>
        Some(
          QueryExecutionErrors.invalidArrayIndexError(
            params("indexValue").toString.toInt,
            params("arraySize").toString.toInt,
            sqlCtx(context)))

      case "InvalidElementAtIndex" =>
        Some(
          QueryExecutionErrors.invalidElementAtIndexError(
            params("indexValue").toString.toInt,
            params("arraySize").toString.toInt,
            sqlCtx(context)))

      case "InvalidIndexOfZero" =>
        Some(QueryExecutionErrors.invalidIndexOfZeroError(sqlCtx(context)))

      case "InvalidBitmapPosition" =>
        // invalidBitmapPositionError does not exist in Spark 3.4
        Some(
          new SparkException(
            errorClass = "INVALID_BITMAP_POSITION",
            messageParameters = params.map { case (k, v) => (k, v.toString) },
            cause = null))

      case "DuplicatedMapKey" =>
        Some(QueryExecutionErrors.duplicateMapKeyFoundError(params("key")))

      case "NullMapKey" =>
        Some(QueryExecutionErrors.nullAsMapKeyNotAllowedError())

      case "MapKeyValueDiffSizes" =>
        Some(QueryExecutionErrors.mapDataKeyArrayLengthDiffersFromValueArrayLengthError())

      case "ExceedMapSizeLimit" =>
        Some(QueryExecutionErrors.exceedMapSizeLimitError(params("size").toString.toInt))

      case "CollectionSizeLimitExceeded" =>
        // createArrayWithElementsExceedLimitError takes (count: Any) in Spark 3.4
        Some(
          QueryExecutionErrors.createArrayWithElementsExceedLimitError(
            params("numElements").toString.toLong))

      case "NotNullAssertViolation" =>
        Some(
          QueryExecutionErrors.foundNullValueForNotNullableFieldError(
            params("fieldName").toString))

      case "ValueIsNull" =>
        Some(
          QueryExecutionErrors.fieldCannotBeNullError(
            params.getOrElse("rowIndex", 0).toString.toInt,
            params("fieldName").toString))

      case "CannotParseTimestamp" =>
        Some(
          QueryExecutionErrors.ansiDateTimeParseError(new Exception(params("message").toString)))

      case "InvalidFractionOfSecond" =>
        Some(QueryExecutionErrors.invalidFractionOfSecondError())

      case "CastInvalidValue" =>
        val str = UTF8String.fromString(params("value").toString)
        val targetType = getDataType(params("toType").toString)
        Some(
          QueryExecutionErrors
            .invalidInputInCastToNumberError(targetType, str, sqlCtx(context)))

      case "InvalidInputInCastToDatetime" =>
        val expression =
          s"'${params("value").toString.replace("\\", "\\\\").replace("'", "\\'")}'"
        val sourceType = s""""${params("fromType").toString}""""
        val targetType = s""""${params("toType").toString}""""
        Some(
          new SparkDateTimeException(
            errorClass = "CAST_INVALID_INPUT",
            messageParameters = Map(
              "expression" -> expression,
              "sourceType" -> sourceType,
              "targetType" -> targetType,
              "ansiConfig" -> "\"spark.sql.ansi.enabled\""),
            context = context,
            summary = summary))

      case "CastOverFlow" =>
        val fromType = getDataType(params("fromType").toString)
        val toType = getDataType(params("toType").toString)
        val valueStr = params("value").toString

        val typedValue: Any = fromType match {
          case _: DecimalType =>
            val cleanStr = if (valueStr.endsWith("BD")) valueStr.dropRight(2) else valueStr
            Decimal(cleanStr)
          case ByteType =>
            val cleanStr = if (valueStr.endsWith("T")) valueStr.dropRight(1) else valueStr
            cleanStr.toByte
          case ShortType =>
            val cleanStr = if (valueStr.endsWith("S")) valueStr.dropRight(1) else valueStr
            cleanStr.toShort
          case IntegerType => valueStr.toInt
          case LongType =>
            val cleanStr = if (valueStr.endsWith("L")) valueStr.dropRight(1) else valueStr
            cleanStr.toLong
          case FloatType => parseFloatLiteral(valueStr)
          case DoubleType => parseDoubleLiteral(valueStr)
          case StringType => UTF8String.fromString(valueStr)
          case _ => valueStr
        }

        Some(QueryExecutionErrors.castingCauseOverflowError(typedValue, fromType, toType))

      case "CannotParseDecimal" =>
        Some(QueryExecutionErrors.cannotParseDecimalError())

      case "InvalidUtf8String" =>
        // invalidUTF8StringError does not exist in Spark 3.x; use generic fallback
        Some(
          new SparkException(
            errorClass = "INVALID_UTF8_STRING",
            messageParameters = params.map { case (k, v) => (k, v.toString) },
            cause = null))

      case "UnexpectedPositiveValue" =>
        Some(
          QueryExecutionErrors.unexpectedValueForStartInFunctionError(
            params("parameterName").toString))

      case "UnexpectedNegativeValue" =>
        Some(
          QueryExecutionErrors.unexpectedValueForLengthInFunctionError(
            params("parameterName").toString))

      case "InvalidRegexGroupIndex" =>
        // invalidRegexGroupIndexError does not exist in Spark 3.4
        Some(
          new SparkException(
            errorClass = "INVALID_REGEX_GROUP_INDEX",
            messageParameters = params.map { case (k, v) => (k, v.toString) },
            cause = null))

      case "DatatypeCannotOrder" =>
        // orderedOperationUnsupportedByDataTypeError takes DataType in Spark 3.4, not String
        Some(
          QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError(
            getDataType(params("dataType").toString)))

      case "ScalarSubqueryTooManyRows" =>
        // multipleRowScalarSubqueryError was renamed to multipleRowSubqueryError in Spark 3.x
        Some(QueryExecutionErrors.multipleRowSubqueryError(sqlCtx(context)))

      case "IntervalArithmeticOverflowWithSuggestion" =>
        // Spark 3.x uses a single intervalArithmeticOverflowError method
        Some(
          QueryExecutionErrors.intervalArithmeticOverflowError(
            "Interval arithmetic overflow",
            params.get("functionName").map(_.toString).getOrElse(""),
            sqlCtx(context)))

      case "IntervalArithmeticOverflowWithoutSuggestion" =>
        Some(
          QueryExecutionErrors
            .intervalArithmeticOverflowError("Interval arithmetic overflow", "", sqlCtx(context)))

      case "DuplicateFieldCaseInsensitive" =>
        Some(
          QueryExecutionErrors.foundDuplicateFieldInCaseInsensitiveModeError(
            params("requiredFieldName").toString,
            params("matchedOrcFields").toString))

      case "DuplicateFieldByFieldId" =>
        // Mirror Spark's `ParquetReadSupport.matchIdField` which calls
        // `foundDuplicateFieldInFieldIdLookupModeError` when more than one Parquet field
        // shares an id requested by the read schema.
        Some(
          QueryExecutionErrors.foundDuplicateFieldInFieldIdLookupModeError(
            params("requiredId").toString.toInt,
            params("matchedFields").toString))

      case "ParquetMissingFieldIds" =>
        // Mirror Spark's `ParquetReadSupport.inferSchema`, which throws a plain
        // `RuntimeException` (not a SparkException) when the read schema requests field
        // ids and the file carries none.
        Some(
          new RuntimeException(
            "Spark read schema expects field Ids, but Parquet file schema doesn't " +
              "contain any field Ids. Please remove the field ids from Spark schema or " +
              "ignore missing ids by setting " +
              "`spark.sql.parquet.fieldId.read.ignoreMissing = true`"))

      case "ParquetSchemaConvert" =>
        // Mirror Spark's FileScanRDD: wrap the SchemaColumnConvertNotSupportedException
        // in a SparkException whose message is "Parquet column cannot be converted in
        // file <path>...". The native side may not have the file path; an empty path
        // still produces a message that contains "Parquet column cannot be converted in
        // file" (which is what Spark's own SQL tests assert).
        val column = params("column").toString
        val physicalType = params("physicalType").toString
        val logicalType = params("sparkType").toString
        val filePath = params.get("filePath").map(_.toString).getOrElse("")
        val cause =
          new SchemaColumnConvertNotSupportedException(column, physicalType, logicalType)
        Some(
          QueryExecutionErrors.unsupportedSchemaColumnConvertError(
            filePath,
            column,
            logicalType,
            physicalType,
            cause))

      case "FileNotFound" =>
        val msg = params("message").toString
        // Extract file path from native error message and format like Hadoop's
        // FileNotFoundException: "File <path> does not exist"
        val path = ShimSparkErrorConverter.ObjectLocationPattern
          .findFirstMatchIn(msg)
          .map(_.group(1))
          .getOrElse(msg)
        Some(
          QueryExecutionErrors.readCurrentFileNotFoundError(
            new FileNotFoundException(s"File $path does not exist")))

      case _ =>
        None
    }
  }

  private def getDataType(typeName: String): DataType = {
    typeName.toUpperCase match {
      case "BYTE" | "TINYINT" => ByteType
      case "SHORT" | "SMALLINT" => ShortType
      case "INT" | "INTEGER" => IntegerType
      case "LONG" | "BIGINT" => LongType
      case "FLOAT" | "REAL" => FloatType
      case "DOUBLE" => DoubleType
      case "DECIMAL" => DecimalType.SYSTEM_DEFAULT
      case "STRING" | "VARCHAR" => StringType
      case "BINARY" => BinaryType
      case "BOOLEAN" => BooleanType
      case "DATE" => DateType
      case "TIMESTAMP" => TimestampType
      case _ =>
        try {
          DataType.fromDDL(typeName)
        } catch {
          case _: Exception =>
            // fromDDL rejects types that are syntactically invalid in SQL DDL, such as
            // DECIMAL(p,s) with a negative scale (valid when allowNegativeScaleOfDecimal=true).
            // Parse those manually rather than silently falling back to StringType.
            if (typeName.toUpperCase.startsWith("DECIMAL(") && typeName.endsWith(")")) {
              val inner = typeName.substring("DECIMAL(".length, typeName.length - 1)
              val parts = inner.split(",")
              if (parts.length == 2) {
                try {
                  DataTypes.createDecimalType(parts(0).trim.toInt, parts(1).trim.toInt)
                } catch {
                  case _: Exception => StringType
                }
              } else StringType
            } else StringType
        }
    }
  }
}
