using Revise
using LibPQ, JSON3, PrettyPrinting, UUIDs, DataFrames, DataStructures, Base64
using GeneralUtils, SQLLLM


config = copy(JSON3.read("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 text2textInstructLLM(prompt::String)
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:text2textinstruct][:mqtttopic];
    msgPurpose="inference",
    senderName="yiemagent",
    senderId=string(uuid4()),
    receiverName="text2textinstruct",
    mqttBrokerAddress=config[:mqttServerInfo][:broker],
    mqttBrokerPort=config[:mqttServerInfo][:port],
  )

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

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

  return response
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 addSQLVectorDB(state)
  # get embedding of the query
  query = [state[:thoughtHistory][:question]]
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:text2textinstruct][:mqtttopic];
    msgPurpose= "embedding",
    senderName= "yiemagent",
    senderId= string(uuid4()),
    receiverName= "text2textinstruct",
    mqttBrokerAddress= config[:mqttServerInfo][:broker],
    mqttBrokerPort= config[:mqttServerInfo][:port],
  )

  outgoingMsg = Dict(
    :msgMeta=> msgMeta,
    :payload=> Dict(
      :text=> query
      )
    )
  response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg)
  embedding = response[:response][:embeddings][1]

  # check whether there is close enough vector already store in vectorDB. if no, add, else skip 
  sql =
    """
    SELECT *, embedding <-> '$embedding' as distance
    FROM sql_statement_repository
    ORDER BY distance LIMIT 1;
    """
  response = executeSQLVectorDB(sql)
  df = DataFrame(response)
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]
  if row == 0 || distance > 10 # no close enough SQL stored in the database
    latestKey, _ = GeneralUtils.findHighestIndexKey(state[:thoughtHistory], :action_input)
    _sqlStatement = state[:thoughtHistory][latestKey]
    if occursin("SELECT", _sqlStatement)  # make sure it is an SQL statement before adding into DB
      sqlStatementBase64 = base64encode(_sqlStatement)
      sqlStatement = replace(_sqlStatement, "'"=>"")
      sql =
      """
      INSERT INTO sql_statement_repository (question, sql_statement, sql_statement_base64, embedding) VALUES ('$query', '$sqlStatement', '$sqlStatementBase64', '$embedding');
      """
      _ = executeSQLVectorDB(sql)
      println("--> added new SQL statement to vectorDB ", @__FILE__, " ", @__LINE__)
      println(sqlStatement)
    end
  end
end

function querySQLVectorDB(state)

  # provide similarSQL at the first time thinking only
  latestKey, _ = GeneralUtils.findHighestIndexKey(state[:thoughtHistory], :action_input)
  if latestKey === nothing
    # get embedding of the query
    query = [state[:thoughtHistory][:question]]
    msgMeta = GeneralUtils.generate_msgMeta(
      config[:externalservice][:text2textinstruct][:mqtttopic];
      msgPurpose= "embedding",
      senderName= "yiemagent",
      senderId= string(uuid4()),
      receiverName= "text2textinstruct",
      mqttBrokerAddress= config[:mqttServerInfo][:broker],
      mqttBrokerPort= config[:mqttServerInfo][:port],
    )

    outgoingMsg = Dict(
      :msgMeta=> msgMeta,
      :payload=> Dict(
        :text=> query
        )
      )
    response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg)
    embedding = response[:response][:embeddings][1]

    # check whether there is close enough vector already store in vectorDB. if no, add, else skip 
    sql =
    """
    SELECT *, embedding <-> '$embedding' as distance
    FROM sql_statement_repository
    ORDER BY distance LIMIT 1;
    """
    response = executeSQLVectorDB(sql)
    df = DataFrame(response)
    row, col = size(df)
    distance = row == 0 ? Inf : df[1, :distance]
    if row != 0 && distance < 100
      # if there is usable SQL, return it.
      sqlStatementBase64 = df[1, :sql_statement_base64]
      sqlStatement = String(base64decode(sqlStatementBase64))
      return sqlStatement
    else
      return nothing
    end
  end
  return nothing
end




# query = Dict(:text=> "How many wines from France do you have that can be paired with lamb?")
# query = "How many wines are from United States?"
query = "retailer: Yiem, wine_type: red, sweetness: 1-2, intensity: 4-5, wine price: 20-40"
# query = "wine_type: white, country: United States, sweetness: 1-2, tannin: 3, food to be served with wine: pizza"
# query = "wine_type: white, country: Austria, food to be served with wine: pork"
# query = "wine price: less than 25, wine_type: rose, country: France, sweetness: 2, tannin: 3, food to be served with wine: pizza"
# query = Dict(:text=> "wine_type: white, country: France, sweetness: 1")
result = SQLLLM.query(query, executeSQL, text2textInstructLLM;
                      addSQLVectorDB=addSQLVectorDB, 
                      querySQLVectorDB=querySQLVectorDB)

println(result)
error(555)



"""
CREATE TABLE sql_statement_repository (id bigserial PRIMARY KEY, question text, sql_statement text, sql_statement_base64 text, embedding vector(768));

SELECT * FROM wine WHERE wine_type = 'red' AND country = 'France' AND sweetness >= 1 AND sweetness <= 2 AND intensity >= 4 AND intensity <= 5 ORDER BY RANDOM() LIMIT 2;


"""


# sql = 
# """
# SELECT COUNT(*)  FROM wine_food  JOIN wine ON wine_food.wine_id = wine.wine_id  JOIN food ON wine_food.food_id = food.food_id  WHERE food.description LIKE '%lamb%';
# """
# response = SQLLLM.SQLexecution(executeSQL, sql);
# result = response[:result]
# userintention = 
# """
# Since this is the first round, there's no execution error to analyze. However, we can think about how to improve the query to achieve the desired result.
# 1) We need to join the wine_food table with the food table on the food_id column.\n  2) We want to filter the results to include only wines that can be paired with lamb by checking if the food_name or additional_search_term matches 'lamb'.\n  3) We'll use a COUNT(DISTINCT) function to count the number of unique wine_id values that meet the condition.
# """
# userintention_dict = Dict(:userintention=>userintention)



# sql = 
# """
# SELECT DISTINCT wf.wine_id, COUNT(wf.wine_id) AS wine_count FROM wine_food wf JOIN food f ON wf.food_id = f.food_id WHERE f.description LIKE '%lamb%' GROUP BY wf.wine_id ORDER BY wine_count DESC;
# """
# response = SQLLLM.SQLexecution(executeSQL, sql);
# result = response[:result]
# userintention = 
# """
# 1. Use TABLEINFO function to get information about the columns in the wine_food table.\n2. Use GETDATA function to retrieve data from the wine_food table that contains information about wines paired with lamb.\n3. Join the retrieved data with the wine table on the wine_id column to get information about the wines that can be paired with lamb.\n4. Count the number of unique wines associated with lamb through the wine_food junction table. 1. Use TABLEINFO function to get information about the columns in the wine_food table.\n2. Use GETDATA function to retrieve data from the wine_food table that contains information about wines paired with lamb.\n3. Join the retrieved data with the wine table on the wine_id column to get information about the wines that can be paired with lamb.\n4. Count the number of unique wines associated with lamb through the wine_food junction table.
# """
# userintention_dict = Dict(:userintention=>userintention)



# sql = 
# """
# SELECT COUNT(DISTINCT w.wine_name) FROM (SELECT * FROM wine_food wf JOIN food f ON wf.food_id = f.food_id WHERE f.food_name = 'lamb') AS temp_table JOIN wine w ON temp_table.wine_id = w.wine_id;
# """
# response = SQLLLM.SQLexecution(executeSQL, sql);
# result = response[:result]
# userintention = 
# """
# 1. Join the wine_food table with the food table using the food_id column in both tables.\n2. Filter the results to only include rows where the associated food is 'lamb'.\n3. Join the resulting table with the wine table using the wine_id column in both tables.\n4. Count the number of unique wines that can be paired with lamb. 1. Join the wine_food table with the food table using the food_id column in both tables.\n2. Filter the results to only include rows where the associated food is 'lamb'.\n3. Join the resulting table with the wine table using the wine_id column in both tables.\n4. Count the number of unique wines that can be paired with lamb.
# """
# userintention_dict = Dict(:userintention=>userintention)



# sql = 
# """
# SELECT * FROM wine WHERE country = 'France' AND sweetness = 1 AND wine_type = 'white' LIMIT 2;
# """
# response = SQLLLM.SQLexecution(executeSQL, sql);
# result = response[:result]
# userintention = 
# """
# "- Identify the primary key in the wine table.\n- Filter the results to only include wines with type white, from France and level of sweetness 1.\n- Retrieve the information about wines that match the specified criteria. - Identify the primary key in the wine table.\n- Filter the results to only include wines with type white, from France and level of sweetness 1.\n- Retrieve the information about wines that match the specified criteria.
# """
# userintention_dict = Dict(:userintention=>userintention)















# readout = SQLLLM.extractContent_dataframe(result, userintention_dict, text2textInstructLLM)

# println("runtest.jl is done")






# sql = 
# """
# SELECT * FROM wine WHERE country = 'France' AND sweetness = 1 AND wine_type = 'white' LIMIT 2;
# """
# _result = executeSQL(sql)
# df2 = DataFrame(_result)











# state = Dict(
#   :isterminal => true,
#   :lesson => nothing,
#   :reward => 1,
#   :evaluation =>
#       "The user's question is to search the database for wines that have a type of \"white\", are from \"France\", and have a sweetness level of 1. The thought is correct in identifying the conditions needed to filter the wine table. The action taken is to execute a SQL query to retrieve the desired data, which is also correct. The observation provides a search summary and two search results that match the user's question. Each result includes details about the wine such as ID, name, brand, manufacturer, region, country, type, grape variety, serving temperature, intensity, sweetness, tannin, and acidity.",
#   :accepted_as_answer => "Yes",
#   :thoughtHistory =>
#       OrderedDict{Symbol, Any}(:question => "Search the database for wine_type: white, country: France, sweetness: 1", :thought_1 => "The user wants to search the database for wines that have a type of \"white\", are from \"France\", and have a sweetness level of 1. To achieve this, we need to filter the wine table based on these conditions.", :action_name_1 => "GETDATA", :action_input_1 => "SELECT * FROM wine WHERE wine.wine_type = 'white' AND wine.country = 'France' AND wine.sweetness = 1;", :observation_1 => "\"Search summary: The resulting table represents wines.\\nSearch result: 1) wine_id: 5b6b6df9-d87c-4f33-8995-7249c2ecc917, wine_name: corton-charlemagne grand cru, brand: domaine des croix, manufacturer: domaine des croix, region: bourgogne, country: France, wine_type: white, grape_variety: cote de beaune blanc, serving_temperature: 11 to 13 Celsius, intensity: 4, sweetness: 1, tannin: missing, acidity: 3, fizziness: missing\\n2) wine_id: 1ad27d16-ef64-4907-acf1-40631630c143, wine_name: puligny-montrachet 1er cru 'les demoiselles', brand: amiot guy, manufacturer: amiot guy, region: bourgogne, country: France, wine_type: white, grape_variety: cote de beaune blanc, serving_temperature: 11 to 13 Celsius, intensity: 4, sweetness: 1, tannin: missing, acidity: 3, fizziness: missing\\n\\n\""),
#   :evaluationscore => 9,
#   :select => nothing,
#   :suggestion => "None")



# result = SQLLLM.evaluator(state, text2textInstructLLM)












println("runtest.jl done")




