ES.scala

package com.gmei.up.utils

import scala.collection.mutable.ListBuffer
import scala.util.Try

import org.json4s.JsonDSL._
import org.json4s.jackson.JsonMethods._
import org.json4s.JObject

import org.apache.http.HttpHost
import org.apache.http.auth.{ AuthScope, UsernamePasswordCredentials }
import org.apache.http.impl.client.BasicCredentialsProvider
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder

import org.elasticsearch.client.RestClient
import org.elasticsearch.client.RestClientBuilder.HttpClientConfigCallback
import com.sksamuel.elastic4s.requests.searches.{ SearchRequest, SearchResponse }
import com.sksamuel.elastic4s.requests.searches.sort.{ FieldSort, SortOrder, ScriptSortType }
import com.sksamuel.elastic4s.ElasticDsl._
import com.sksamuel.elastic4s.{ HitReader, Hit, ElasticClient }
import com.sksamuel.elastic4s.http.JavaClient
import com.sksamuel.elastic4s.requests.searches.queries.Query
import com.sksamuel.elastic4s.requests.searches.queries.term.TermsQuery
import com.sksamuel.elastic4s.requests.script.{ Script, ScriptType }
import com.sksamuel.elastic4s.requests.searches.queries.funcscorer.{
  CombineFunction,
  FunctionScoreQueryScoreMode,
  WeightScore,
  ScoreFunction,
  ScriptScore
}

object ESClient {
  private class ElasticCredentials(username: String, password: String) extends HttpClientConfigCallback {
    val credentials = new UsernamePasswordCredentials(username, password)
    val credentialsProvider = new BasicCredentialsProvider()
    credentialsProvider.setCredentials(AuthScope.ANY, credentials)

    override def customizeHttpClient(httpClientBuilder: HttpAsyncClientBuilder): HttpAsyncClientBuilder =
      httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider)
  }

  def create(hostname: String, port: Int, username: String, password: String): ElasticClient = {
    val restClient = RestClient
      .builder(new HttpHost(hostname, port, "http"))
      .setHttpClientConfigCallback(new ElasticCredentials(username, password))
      .build()

    ElasticClient(JavaClient.fromRestClient(restClient))
  }
}

object ES {
  case class Character(id: String, index: String, projects: String)

  implicit object CharacterHitReader extends HitReader[Character] {
    override def read(hit: Hit): Try[Character] = {
      val source = hit.sourceAsMap
      Try(Character(source("id").toString, hit.index, source("tags_v3").toString))
    }
  }

  def test(
      diaryRequest: SearchRequest,
      tractateRequest: SearchRequest,
      answerRequest: SearchRequest,
      serviceDiaryReq: SearchRequest
  ) = {
    // TODO read from config
    val client: ElasticClient = ESClient.create("172.16.52.33", 9200, "elastic", "gengmei!@#")

    val multiReq = multi(
      diaryRequest,
      tractateRequest,
      answerRequest,
      serviceDiaryReq
    )

    println(multiReq.show)

    // TODO remove await
    val resp = client.execute(multiReq).await

    // println(resp)

    val a = resp.result.to[Character]

    println("%%%%%%%%%%%")
    println(a.size)
    a.foreach { x =>
      println(x.id, x.index, x.projects)
    }
    println("%%%%%%%%%%%")

    resp
  }

  def generateDiaryRequest(
      projects: List[String],
      secondDemands: List[String],
      cityId: Int = -1
  ): SearchRequest = {
    val size = 15
    val includes = List("id", "tags_v3")
    var scoreListBuffer: ListBuffer[ScoreFunction] = ListBuffer.empty[ScoreFunction]
    var shouldListBuffer: ListBuffer[Query] = ListBuffer.empty[Query]

    if (projects.size > 0) {
      scoreListBuffer += WeightScore(60).filter(
        must(termsQuery("tags_v3", projects), termQuery("is_promote", true))
      )
      scoreListBuffer += WeightScore(50).filter(termsQuery("tags_v3", projects))
      shouldListBuffer += TermsQuery("tags_v3", projects)
    }
    if (secondDemands.size > 0) {
      scoreListBuffer += WeightScore(40).filter(termsQuery("second_demands", secondDemands))
      shouldListBuffer += TermsQuery("second_demands", secondDemands)
    }

    val res = search("gm-dbmw-diary-read")
      .query {
        functionScoreQuery()
          .functions(scoreListBuffer.toList)
          .scoreMode(FunctionScoreQueryScoreMode.Max)
          .boostMode(CombineFunction.Replace)
          .query {
            boolQuery()
              .must(
                termQuery("is_online", true),
                termQuery("has_cover", true),
                termQuery("is_sink", false),
                termQuery("has_after_cover", true),
                termQuery("has_before_cover", true),
                rangeQuery("content_level").gte(3),
                termQuery("content_simi_bol_show", 0)
              )
              .not(termQuery("show_by_index", 2))
              .should(shouldListBuffer.toList)
              .minimumShouldMatch(1)
          }
      }
      .size(size)
      .sourceInclude(includes)
      .sortBy(
        scriptSort(
          script(
            "if (doc['doctor.hospital.city_tag_id'].size()!=0){if (doc['doctor.hospital.city_tag_id'].value ==params.user_city_tag_id){return 2;} else if (doc['nearby_city_tags_v1.tag_id'].contains(params.user_city_tag_id) ) {return 1;} else {return 0;} }else{return 0}"
          ).params(Map("user_city_tag_id" -> cityId)).lang("painless")
        ).typed("number").order(SortOrder.Desc),
        FieldSort("_score").desc(),
        FieldSort("has_video_cover").asc(),
        FieldSort("offline_score").desc(),
        FieldSort("good_click").desc(),
        FieldSort("last_update_time").desc()
      )
    // println("diary query")
    // println(res.show)
    res
  }

  def generateTractateRequest(projects: List[String], secondDemands: List[String]): SearchRequest = {
    val size = 15
    val includes = List("id", "tags_v3")
    var scoreListBuffer: ListBuffer[ScoreFunction] = ListBuffer.empty[ScoreFunction]
    var shouldListBuffer: ListBuffer[Query] = ListBuffer.empty[Query]

    if (projects.size > 0) {
      scoreListBuffer += WeightScore(60).filter(
        must(termsQuery("tags_v3", projects), termQuery("is_promote", true))
      )
      scoreListBuffer += WeightScore(50).filter(termsQuery("tags_v3", projects))
      shouldListBuffer += TermsQuery("tags_v3", projects)
    }
    if (secondDemands.size > 0) {
      scoreListBuffer += WeightScore(40).filter(termsQuery("second_demands", secondDemands))
      shouldListBuffer += TermsQuery("second_demands", secondDemands)
    }

    val res = search("gm-dbmw-tractate-read")
      .query {
        functionScoreQuery()
          .functions(scoreListBuffer.toList)
          .scoreMode(FunctionScoreQueryScoreMode.Max)
          .boostMode(CombineFunction.Replace)
          .query {
            boolQuery()
              .must(
                termQuery("is_online", true),
                rangeQuery("content_level").gte(3)
              )
              .not(termQuery("show_by_index", 2), termQuery("status", 4))
              .should(shouldListBuffer.toList)
              .minimumShouldMatch(1)
          }
      }
      .size(size)
      .sourceInclude(includes)
      .sortBy(
        FieldSort("_score").desc()
      )
    // println("tractate query")
    // println(res.show)
    res
  }

  def generateAnswerRequest(projects: List[String], secondDemands: List[String]): SearchRequest = {
    val size = 15
    val includes = List("id", "tags_v3")
    var scoreListBuffer: ListBuffer[ScoreFunction] = ListBuffer.empty[ScoreFunction]
    var shouldListBuffer: ListBuffer[Query] = ListBuffer.empty[Query]

    if (projects.size > 0) {
      scoreListBuffer += WeightScore(60).filter(
        must(termsQuery("tags_v3", projects), termQuery("is_promote", true))
      )
      scoreListBuffer += WeightScore(50).filter(termsQuery("tags_v3", projects))
      shouldListBuffer += TermsQuery("tags_v3", projects)
    }
    if (secondDemands.size > 0) {
      scoreListBuffer += WeightScore(40).filter(termsQuery("second_demands", secondDemands))
      shouldListBuffer += TermsQuery("second_demands", secondDemands)
    }

    val res = search("gm-dbmw-answer-read")
      .query {
        functionScoreQuery()
          .functions(scoreListBuffer.toList)
          .scoreMode(FunctionScoreQueryScoreMode.Max)
          .boostMode(CombineFunction.Replace)
          .query {
            boolQuery()
              .must(
                termQuery("is_online", true),
                rangeQuery("content_level").gte(3),
                rangeQuery("content_length").gte(30)
              )
              .not(termQuery("show_by_index", 2))
              .should(shouldListBuffer.toList)
              .minimumShouldMatch(1)
          }
      }
      .size(size)
      .sourceInclude(includes)
      .sortBy(
        FieldSort("_score").desc(),
        FieldSort("has_picture").desc(),
        FieldSort("smart_rank_v2").desc(),
        FieldSort("good_click").desc()
      )
    // println("answer query")
    // println(res.show)
    res
  }

  def generateServiceDiaryRequest(
      projects: List[String],
      secondDemands: List[String],
      cityId: Int = -1
  ): SearchRequest = {
    val size = 15
    val includes = List("id", "tags_v3")
    var shouldListBuffer: ListBuffer[Query] = ListBuffer.empty[Query]

    if (projects.size > 0) {
      shouldListBuffer += TermsQuery("tags_v3", projects)
    }
    if (secondDemands.size > 0) {
      shouldListBuffer += TermsQuery("second_demands", secondDemands)
    }

    val res = search("gm-dbmw-diary-read")
      .query {
        functionScoreQuery(
          )
          .functions(
            ScriptScore(
              "return (doc.containsKey('new_smr')&&doc['new_smr'].value>0?doc['new_smr'].value:0.0) * ((doc.containsKey('is_promote')&&doc['is_promote'].value?5.0:0.0)+(doc.containsKey('has_before_cover')&&doc['has_before_cover'].value?2.0:0.0) + (doc.containsKey('has_after_cover')&&doc['has_after_cover'].value?2.0:0.0)+(doc.containsKey('has_after_cover')&&doc['has_after_cover'].value?8.0:0.0))"
            ),
            WeightScore(1000000).filter(termQuery("doctor.hospital.city_tag_id", cityId)),
            WeightScore(1000000).filter(termQuery("service_nearby_city_tag_ids", cityId))
          )
          .scoreMode(FunctionScoreQueryScoreMode.Sum)
          .boostMode(CombineFunction.Sum)
          .query {
            boolQuery()
              .must(
                termQuery("is_online", true),
                termQuery("has_service", true),
                termQuery("service.is_promote", true),
                rangeQuery("content_level").gte(3)
              )
              .not(termQuery("show_by_index", 2), termQuery("doctor.hospital.id", "shyxylmryy2203"))
              .should(shouldListBuffer.toList)
              .minimumShouldMatch(1)
          }
      }
      .size(size)
      .sourceInclude(includes)
    // println("service diary query")
    // println(res.show)
    res
  }
}