using Revise
using JSON, JSON3, Dates, UUIDs, PrettyPrinting, LibPQ, Base64, DataFrames
using YiemAgent, GeneralUtils
using Base.Threads

# ---------------------------------------------- 100 --------------------------------------------- #



# load config
config = JSON.parsefile("/appfolder/app/dev/YiemAgent/test/config.json")
# config = copy(JSON.parsefile("../mountvolume/config.json"))


function executeSQL(sql::T) where {T<:AbstractString}
  host = config[:externalservice][:wineDB][:host]
  port = config[:externalservice][:wineDB][:port]
  dbname = config[:externalservice][:wineDB][:dbname]
  user = config[:externalservice][:wineDB][:user]
  password = config[:externalservice][:wineDB][:password]
  DBconnection = LibPQ.Connection("host=$host port=$port dbname=$dbname user=$user password=$password")
  result = LibPQ.execute(DBconnection, sql)
  close(DBconnection)
  return result
end

function executeSQLVectorDB(sql)
  host = config[:externalservice][:SQLVectorDB][:host]
  port = config[:externalservice][:SQLVectorDB][:port]
  dbname = config[:externalservice][:SQLVectorDB][:dbname]
  user = config[:externalservice][:SQLVectorDB][:user]
  password = config[:externalservice][:SQLVectorDB][:password]
  DBconnection = LibPQ.Connection("host=$host port=$port dbname=$dbname user=$user password=$password")
  result = LibPQ.execute(DBconnection, sql)
  close(DBconnection)
  return result
end

function text2textInstructLLM(prompt::String; maxattempt::Integer=3, modelsize::String="medium",
    llmkwargs=Dict(
      :num_ctx => 32768,
      :temperature => 0.1,
    )
  )
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:loadbalancer][:mqtttopic];
    msgPurpose="inference",
    senderName="yiemagent",
    senderId=sessionId,
    receiverName="text2textinstruct_$modelsize",
    mqttBrokerAddress=config[:mqttServerInfo][:broker],
    mqttBrokerPort=config[:mqttServerInfo][:port],
  )

  outgoingMsg = Dict(
    :msgMeta => msgMeta,
    :payload => Dict(
      :text => prompt,
      :kwargs => llmkwargs
    )
  )

  response = nothing
  for attempts in 1:maxattempt
    _response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=180, maxattempt=maxattempt)
    payload = _response[:response]
    if _response[:success] && payload[:text] !== nothing
      response = _response[:response][:text]
      break
    else
      println("\n<text2textInstructLLM()> attempt $attempts/$maxattempt failed ", @__FILE__, ":", @__LINE__, " $(Dates.now())")
      pprintln(outgoingMsg)
      println("</text2textInstructLLM()> attempt $attempts/$maxattempt failed ", @__FILE__, ":", @__LINE__, " $(Dates.now())\n")
      sleep(3)
    end
  end

  return response
end

# get text embedding from a LLM service
function getEmbedding(text::T) where {T<:AbstractString}
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:loadbalancer][:mqtttopic];
    msgPurpose="embedding",
    senderName="yiemagent",
    senderId=sessionId,
    receiverName="textembedding",
    mqttBrokerAddress=config[:mqttServerInfo][:broker],
    mqttBrokerPort=config[:mqttServerInfo][:port],
  )

  outgoingMsg = Dict(
    :msgMeta => msgMeta,
    :payload => Dict(
      :text => [text] # must be a vector of string
    )
  )

  response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=120, maxattempt=3)
  embedding = response[:response][:embeddings]
  return embedding
end

function findSimilarTextFromVectorDB(text::T1, tablename::T2, embeddingColumnName::T3,
  vectorDB::Function; limit::Integer=1
  )::DataFrame where {T1<:AbstractString, T2<:AbstractString, T3<:AbstractString}
  # get embedding from LLM service
  embedding = getEmbedding(text)[1]
  # check whether there is close enough vector already store in vectorDB. if no, add, else skip 
  sql = """
        SELECT *, $embeddingColumnName <-> '$embedding' as distance
        FROM $tablename
        ORDER BY distance LIMIT $limit;
        """
  response = vectorDB(sql)
  df = DataFrame(response)
  return df
end

function similarSQLVectorDB(query; maxdistance::Integer=100)
  tablename = "sqlllm_decision_repository"
  # get embedding of the query
  df = findSimilarTextFromVectorDB(query, tablename, 
                                  "function_input_embedding", executeSQLVectorDB)
  # println(df[1, [:id, :function_output]])
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]  
  # distance = 100  # CHANGE this is for testing only
  if row != 0 && distance < maxdistance
    # if there is usable SQL, return it.
    output_b64 = df[1, :function_output_base64]   # pick the closest match
    output_str = String(base64decode(output_b64))
    rowid = df[1, :id]
    println("\n~~~ found similar sql. row id $rowid, distance $distance ", @__FILE__, ":", @__LINE__, " $(Dates.now())")
    return (dict=output_str, distance=distance)
  else
    println("\n~~~ similar sql not found, max distance $maxdistance ", @__FILE__, ":", @__LINE__, " $(Dates.now())")
    return (dict=nothing, distance=nothing)
  end
end

function insertSQLVectorDB(query::T1, SQL::T2; maxdistance::Integer=3) where {T1<:AbstractString, T2<:AbstractString}
  tablename = "sqlllm_decision_repository"
  # get embedding of the query
  # query = state[:thoughtHistory][:question]
  df = findSimilarTextFromVectorDB(query, tablename, 
                                  "function_input_embedding", executeSQLVectorDB)
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]
  if row == 0 || distance > maxdistance # no close enough SQL stored in the database
    query_embedding = getEmbedding(query)[1]
    query = replace(query, "'" => "")
    sql_base64 = base64encode(SQL)
    sql_ = replace(SQL, "'" => "")

    sql = """
          INSERT INTO $tablename (function_input, function_output, function_output_base64, function_input_embedding) VALUES ('$query', '$sql_', '$sql_base64', '$query_embedding');
          """
    # println("\n~~~ added new decision to vectorDB ", @__FILE__, ":", @__LINE__, " $(Dates.now())")
    # println(sql)
    _ = executeSQLVectorDB(sql)
  end
end


function similarSommelierDecision(recentevents::T1; maxdistance::Integer=3
  )::Union{AbstractDict, Nothing} where {T1<:AbstractString}
  tablename = "sommelier_decision_repository"
  # find similar
  println("\n~~~ search vectorDB for this: $recentevents ", @__FILE__, " ", @__LINE__)
  df = findSimilarTextFromVectorDB(recentevents, tablename, 
                                  "function_input_embedding", executeSQLVectorDB)
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]
  if row != 0 && distance < maxdistance
    # if there is usable decision, return it.
    rowid = df[1, :id]
    println("\n~~~ found similar decision. row id $rowid, distance $distance ", @__FILE__, " ", @__LINE__)
    output_b64 = df[1, :function_output_base64]   # pick the closest match
    _output_str = String(base64decode(output_b64))
    output = copy(JSON.parsefile(_output_str))
    return output
  else
    println("\n~~~ similar decision not found, max distance $maxdistance ", @__FILE__, " ", @__LINE__)
    return nothing
  end
end


function insertSommelierDecision(recentevents::T1, decision::T2; maxdistance::Integer=5
  ) where {T1<:AbstractString, T2<:AbstractDict}
  tablename = "sommelier_decision_repository"
  # find similar
  df = findSimilarTextFromVectorDB(recentevents, tablename, 
                                  "function_input_embedding", executeSQLVectorDB)
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]
  if row == 0 || distance > maxdistance # no close enough SQL stored in the database
    recentevents_embedding = getEmbedding(recentevents)[1]
    recentevents = replace(recentevents, "'" => "")
    decision_json = JSON.json(decision)
    decision_base64 = base64encode(decision_json)
    decision = replace(decision_json, "'" => "")
    
    sql = """
          INSERT INTO $tablename (function_input, function_output, function_output_base64, function_input_embedding) VALUES ('$recentevents', '$decision', '$decision_base64', '$recentevents_embedding');
          """
    println("\n~~~ added new decision to vectorDB ", @__FILE__, " ", @__LINE__)
    println(sql)
    _ = executeSQLVectorDB(sql)
  else
    println("~~~ similar decision previously cached, distance $distance ", @__FILE__, " ", @__LINE__)
  end
end


sessionId = "12345"

externalFunction = (
    getEmbedding=getEmbedding,
    text2textInstructLLM=text2textInstructLLM,
    executeSQL=executeSQL,
    similarSQLVectorDB=similarSQLVectorDB,
    insertSQLVectorDB=insertSQLVectorDB,
    similarSommelierDecision=similarSommelierDecision,
    insertSommelierDecision=insertSommelierDecision,
  )



a = YiemAgent.sommelier(
  externalFunction;
  name="Ton",
  id=sessionId,  # agent instance id
  retailername="Yiem",
)

while true
  print("\nyour respond: ")
  user_answer = readline()
  response = YiemAgent.conversation(agent;
                                  userinput=Dict(:text=> user_answer),
                                  maximumMsg=50)
  println("\n$response")
end


# response = YiemAgent.conversation(a, Dict(:text=> "I want to get a French red wine under 100."))


"""
hello I want to get a bottle of red wine for my boss. I have a budget around 50 dollars. Show me some options.

I have no idea about his wine taste but he likes spicy food.


"""






































