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

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



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


function executeSQL(sql::T) where {T<:AbstractString}
  DBconnection = LibPQ.Connection("host=192.168.88.12 port=10201 dbname=wineDB user=yiemtechnologies password=yiemtechnologies@Postgres_0.0")
  result = LibPQ.execute(DBconnection, sql)
  close(DBconnection)
  return result
end

function executeSQLVectorDB(sql)
  DBconnection = LibPQ.Connection("host=192.168.88.12 port=10203 dbname=SQLVectorDB user=yiemtechnologies password=yiemtechnologies@Postgres_0.0")
  result = LibPQ.execute(DBconnection, sql)
  close(DBconnection)
  return result
end

function text2textInstructLLM(prompt::String)
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:text2textinstruct][:mqtttopic];
    msgPurpose="inference",
    senderName="yiemagent",
    senderId=string(uuid4()),
    receiverName="text2textinstruct_small",
    mqttBrokerAddress=config[:mqttServerInfo][:broker],
    mqttBrokerPort=config[:mqttServerInfo][:port],
  )

  outgoingMsg = Dict(
    :msgMeta => msgMeta,
    :payload => Dict(
      :text => prompt,
      :kwargs => Dict(
        :num_ctx => 16384,
        :temperature => 0.2,
      )
    )
  )

  _response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=6000)
  response = _response[:response][:text]

  return response
end

# get text embedding from a LLM service
function getEmbedding(text::T) where {T<:AbstractString}
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:text2textinstruct][:mqtttopic];
    msgPurpose="embedding",
    senderName="yiemagent",
    senderId=string(uuid4()),
    receiverName="text2textinstruct_small",
    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=6000)
  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)
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]  
  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__)
    return (dict=output_str, distance=distance)
  else
    println("\n~~~ similar sql not found, max distance $maxdistance ", @__FILE__, " ", @__LINE__)
    return (dict=nothing, distance=nothing)
  end
end


function insertSQLVectorDB(query::T1, SQL::T2; maxdistance::Integer=1) 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__)
    println(sql)
    _ = executeSQLVectorDB(sql)
  end
end


function similarSommelierDecision(recentevents::T1; maxdistance::Integer=5
  )::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(JSON3.read(_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 = a.func[:getEmbedding](recentevents)[1]
    recentevents = replace(recentevents, "'" => "")
    decision_json = JSON3.write(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
  println("your respond: ")
  user_answer = readline()
  response = YiemAgent.conversation(a, Dict(:text=> user_answer))
  println("\n$response")
end


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