// Copyright (C) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in project root for information.

package com.microsoft.azure.synapse.ml.services.form

import com.microsoft.azure.synapse.ml.services._
import com.microsoft.azure.synapse.ml.services.vision.{BasicAsyncReply, HasImageInput, ReadLine}
import com.microsoft.azure.synapse.ml.logging.{FeatureNames, SynapseMLLogging}
import com.microsoft.azure.synapse.ml.param.ServiceParam
import com.microsoft.azure.synapse.ml.stages.UDFTransformer
import org.apache.http.client.methods.{HttpGet, HttpRequestBase}
import org.apache.http.entity.{AbstractHttpEntity, ByteArrayEntity, ContentType, StringEntity}
import org.apache.spark.injections.UDFUtils
import org.apache.spark.ml.ComplexParamsReadable
import org.apache.spark.ml.util.Identifiable
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.{DataType, StringType}
import spray.json.DefaultJsonProtocol._
import spray.json._

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
abstract class FormRecognizerBase(override val uid: String) extends CognitiveServicesBaseNoHandler(uid)
  with HasCognitiveServiceInput with HasInternalJsonOutputParser with BasicAsyncReply
  with HasImageInput with HasSetLocation with HasSetLinkedService {

  override protected def prepareEntity: Row => Option[AbstractHttpEntity] = {
    r =>
      getValueOpt(r, imageUrl)
        .map(url => new StringEntity(Map("source" -> url).toJson.compactPrint, ContentType.APPLICATION_JSON))
        .orElse(getValueOpt(r, imageBytes)
          .map(bytes => new ByteArrayEntity(bytes, ContentType.APPLICATION_OCTET_STREAM))
        ).orElse(throw new IllegalArgumentException(
        "Payload needs to contain image bytes or url. This code should not run"))
  }

}

trait HasPages extends HasServiceParams {

  val pages = new ServiceParam[String](this, "pages", "The page selection only leveraged for" +
    " multi-page PDF and TIFF documents. Accepted input include single pages (e.g.'1, 2' -> pages 1 and 2 will be " +
    "processed), finite (e.g. '2-5' -> pages 2 to 5 will be processed) and open-ended ranges (e.g. '5-' -> all the" +
    " pages from page 5 will be processed; e.g. '-10' -> pages 1 to 10 will be processed). All of these can be mixed" +
    " together and ranges are allowed to overlap (eg. '-5, 1, 3, 5-10' - pages 1 to 10 will be processed). The" +
    " service will accept the request if it can process at least one page of the document (e.g. using '5-100' on a " +
    "5 page document is a valid input where page 5 will be processed). If no page range is provided, the entire" +
    " document will be processed.", isURLParam = true)

  def setPages(v: String): this.type = setScalarParam(pages, v)

  def setPagesCol(v: String): this.type = setVectorParam(pages, v)

  def getPages: String = getScalarParam(pages)

  def getPagesCol: String = getVectorParam(pages)

}

trait HasTextDetails extends HasServiceParams {
  val includeTextDetails = new ServiceParam[Boolean](this, "includeTextDetails",
    "Include text lines and element references in the result.", isURLParam = true)

  def setIncludeTextDetails(v: Boolean): this.type = setScalarParam(includeTextDetails, v)

  def setIncludeTextDetailsCol(v: String): this.type = setVectorParam(includeTextDetails, v)

  def getIncludeTextDetails: Boolean = getScalarParam(includeTextDetails)

  def getIncludeTextDetailsCol: String = getVectorParam(includeTextDetails)

}

trait HasModelID extends HasServiceParams {
  val modelId = new ServiceParam[String](this, "modelId", "Model identifier.", isRequired = true)

  def setModelId(v: String): this.type = setScalarParam(modelId, v)

  def setModelIdCol(v: String): this.type = setVectorParam(modelId, v)

  def getModelId: String = getScalarParam(modelId)

  def getModelIdCol: String = getVectorParam(modelId)

}

trait HasLocale extends HasServiceParams {
  val locale = new ServiceParam[String](this, "locale", "Locale of the receipt. Supported" +
    " locales: en-AU, en-CA, en-GB, en-IN, en-US.", {
    case Left(s) => Set("en-AU", "en-CA", "en-GB", "en-IN", "en-US")(s)
    case Right(_) => true
  }, isURLParam = true)

  def setLocale(v: String): this.type = setScalarParam(locale, v)

  def setLocaleCol(v: String): this.type = setVectorParam(locale, v)

  def getLocale: String = getScalarParam(locale)

  def getLocaleCol: String = getVectorParam(locale)

}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object FormsFlatteners {

  import FormsJsonProtocol._

  def flattenReadResults(inputCol: String, outputCol: String): UDFTransformer = {
    val fromRow = AnalyzeResponse.makeFromRowConverter

    def extractText(lines: Seq[ReadLine]): String = {
      lines.map(_.text).mkString(" ")
    }

    new UDFTransformer()
      .setUDF(UDFUtils.oldUdf(
        { r: Row =>
          Option(r).map(fromRow).map(
            _.analyzeResult.readResults.map(_.lines.map(extractText).mkString("")).mkString(" ")).mkString("")
        },
        StringType))
      .setInputCol(inputCol)
      .setOutputCol(outputCol)
  }

  def flattenPageResults(inputCol: String, outputCol: String): UDFTransformer = {
    val fromRow = AnalyzeResponse.makeFromRowConverter

    def extractTableText(pageResults: Seq[PageResult]): String = {
      pageResults.map(_.tables.map(_.cells.map(_.text).mkString(" | ")).mkString("\n")).mkString("\n\n")
    }

    def generateKeyValuePairs(keyValuePair: KeyValuePair): String = {
      "key: " + keyValuePair.key.text + " value: " + keyValuePair.value.text
    }

    def extractKeyValuePairs(pageResults: Seq[PageResult]): String = {
      pageResults.map(_.keyValuePairs.map(_.map(generateKeyValuePairs).mkString("\n"))).mkString("\n\n")
    }

    def extractAllText(pageResults: Seq[PageResult]): String = {
      "KeyValuePairs: " + extractKeyValuePairs(pageResults) + "\n\n\n" + "Tables: " + extractTableText(pageResults)
    }

    new UDFTransformer()
      .setUDF(UDFUtils.oldUdf(
        { r: Row =>
          Option(r).map(fromRow).map(
            _.analyzeResult.pageResults.map(extractAllText).mkString(" "))
        },
        StringType))
      .setInputCol(inputCol)
      .setOutputCol(outputCol)
  }

  def flattenDocumentResults(inputCol: String, outputCol: String): UDFTransformer = {
    val fromRow = AnalyzeResponse.makeFromRowConverter

    def extractFields(documentResults: Seq[DocumentResult]): String = {
      documentResults.map(_.fields.toJson.compactPrint).mkString("\n")
    }

    new UDFTransformer()
      .setUDF(UDFUtils.oldUdf(
        { r: Row =>
          Option(r).map(fromRow).map(
            _.analyzeResult.documentResults.map(extractFields).mkString("")).mkString("")
        },
        StringType))
      .setInputCol(inputCol)
      .setOutputCol(outputCol)
  }

  def flattenModelList(inputCol: String, outputCol: String): UDFTransformer = {
    val fromRow = ListCustomModelsResponse.makeFromRowConverter
    new UDFTransformer()
      .setUDF(UDFUtils.oldUdf(
        { r: Row =>
          Option(r).map(fromRow).map(
            _.modelList.map(_.modelId).mkString(" "))
        },
        StringType))
      .setInputCol(inputCol)
      .setOutputCol(outputCol)
  }
}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object AnalyzeLayout extends ComplexParamsReadable[AnalyzeLayout]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class AnalyzeLayout(override val uid: String) extends FormRecognizerBase(uid)
  with SynapseMLLogging with HasPages {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("AnalyzeLayout"))

  def urlPath: String = "/formrecognizer/v2.1/layout/analyze"

  val language = new ServiceParam[String](this, "language", "The BCP-47 language code of the " +
    "text in the document. Layout supports auto language identification and multilanguage documents, so only provide" +
    " a language code if you would like to force the documented to be processed as that specific language.",
    isURLParam = true)

  def setLanguage(v: String): this.type = setScalarParam(language, v)

  def setLanguageCol(v: String): this.type = setVectorParam(language, v)

  val readingOrder = new ServiceParam[String](this, "readingOrder", "Optional parameter to " +
    "specify which reading order algorithm should be applied when ordering the extract text elements. Can be either" +
    " 'basic' or 'natural'. Will default to basic if not specified", isURLParam = true)

  def setReadingOrder(v: String): this.type = setScalarParam(readingOrder, v)

  def setReadingOrderCol(v: String): this.type = setVectorParam(readingOrder, v)

  setDefault(readingOrder -> Left("basic"))

  override protected def responseDataType: DataType = AnalyzeResponse.schema

}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object AnalyzeReceipts extends ComplexParamsReadable[AnalyzeReceipts]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class AnalyzeReceipts(override val uid: String) extends FormRecognizerBase(uid)
  with SynapseMLLogging with HasPages with HasTextDetails with HasLocale {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("AnalyzeReceipts"))

  def urlPath: String = "/formrecognizer/v2.1/prebuilt/receipt/analyze"

  override protected def responseDataType: DataType = AnalyzeResponse.schema

}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object AnalyzeBusinessCards extends ComplexParamsReadable[AnalyzeBusinessCards]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class AnalyzeBusinessCards(override val uid: String) extends FormRecognizerBase(uid)
  with SynapseMLLogging with HasPages with HasTextDetails with HasLocale {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("AnalyzeBusinessCards"))

  def urlPath: String = "/formrecognizer/v2.1/prebuilt/businessCard/analyze"

  override protected def responseDataType: DataType = AnalyzeResponse.schema

}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object AnalyzeInvoices extends ComplexParamsReadable[AnalyzeInvoices]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class AnalyzeInvoices(override val uid: String) extends FormRecognizerBase(uid)
  with SynapseMLLogging with HasPages with HasTextDetails with HasLocale {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("AnalyzeInvoices"))

  def urlPath: String = "/formrecognizer/v2.1/prebuilt/invoice/analyze"

  override protected def responseDataType: DataType = AnalyzeResponse.schema

}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object AnalyzeIDDocuments extends ComplexParamsReadable[AnalyzeIDDocuments]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class AnalyzeIDDocuments(override val uid: String) extends FormRecognizerBase(uid)
  with SynapseMLLogging with HasPages with HasTextDetails {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("AnalyzeIDDocuments"))

  def urlPath: String = "formrecognizer/v2.1/prebuilt/idDocument/analyze"

  override protected def responseDataType: DataType = AnalyzeResponse.schema

}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object ListCustomModels extends ComplexParamsReadable[ListCustomModels]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class ListCustomModels(override val uid: String) extends CognitiveServicesBase(uid)
  with HasCognitiveServiceInput with HasInternalJsonOutputParser
  with HasSetLocation with HasSetLinkedService with SynapseMLLogging {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("ListCustomModels"))

  def urlPath: String = "formrecognizer/v2.1/custom/models"

  override protected def prepareMethod(): HttpRequestBase = new HttpGet()

  val op = new ServiceParam[String](this, "op",
    "Specify whether to return summary or full list of models.", isURLParam = true)

  def setOp(v: String): this.type = setScalarParam(op, v)

  def setOpCol(v: String): this.type = setVectorParam(op, v)

  override protected def prepareEntity: Row => Option[AbstractHttpEntity] = { _ => None }

  override protected def responseDataType: DataType = ListCustomModelsResponse.schema
}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object GetCustomModel extends ComplexParamsReadable[GetCustomModel]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class GetCustomModel(override val uid: String) extends CognitiveServicesBase(uid)
  with HasCognitiveServiceInput with HasInternalJsonOutputParser
  with HasSetLocation with HasSetLinkedService with SynapseMLLogging with HasModelID {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("GetCustomModel"))

  def urlPath: String = "formrecognizer/v2.1/custom/models"

  override protected def prepareUrlRoot: Row => String = { row =>
    getUrl + s"/${getValue(row, modelId)}"
  }

  override protected def prepareMethod(): HttpRequestBase = new HttpGet()

  val includeKeys = new ServiceParam[Boolean](this, "includeKeys",
    "Include list of extracted keys in model information.", isURLParam = true)

  def setIncludeKeys(v: Boolean): this.type = setScalarParam(includeKeys, v)

  def setIncludeKeysCol(v: String): this.type = setVectorParam(includeKeys, v)

  override protected def prepareEntity: Row => Option[AbstractHttpEntity] = { _ => None }

  override protected def responseDataType: DataType = GetCustomModelResponse.schema
}

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
object AnalyzeCustomModel extends ComplexParamsReadable[AnalyzeCustomModel]

@deprecated("The Form Recognition v2.1 API is deprecated please use " +
  "com.microsoft.azure.synapse.ml.services.form.AnalyzeDocument", "v1.0.4")
class AnalyzeCustomModel(override val uid: String) extends FormRecognizerBase(uid)
  with SynapseMLLogging with HasTextDetails with HasModelID {
  logClass(FeatureNames.AiServices.Form)

  def this() = this(Identifiable.randomUID("AnalyzeCustomModel"))

  def urlPath: String = "formrecognizer/v2.1/custom/models"

  override protected def prepareUrlRoot: Row => String = {row =>
    getUrl + s"/${getValue(row, modelId)}/analyze"
  }

  override protected def responseDataType: DataType = AnalyzeResponse.schema
}
