module llmfunction

export  virtualWineUserChatbox, jsoncorrection, winestock,
        virtualWineUserRecommendbox, userChatbox, userRecommendbox

using HTTP, JSON3, URIs, Random, PrettyPrinting, UUIDs
using GeneralUtils, SQLLLM
using ..type, ..util

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

"""

# Arguments
  
# Return

# Example
```jldoctest
julia> 
```

# TODO
  - [] update docstring
  - [WORKING] implement the function

# Signature
"""
function userChatbox(a::T1, input::T2) where {T1<:agent, T2<:AbstractString}
  error("--> userChatbox")

  # put in model format
  virtualWineCustomer = a.config[:externalservice][:virtualWineCustomer_1]
  llminfo = virtualWineCustomer[:llminfo]
  formattedinput = 
  if llminfo[:name] == "llama3instruct"
    formatLLMtext_llama3instruct("assistant", input)
  else
    error("llm model name is not defied yet $(@__LINE__)")
  end

  # send formatted input to user using GeneralUtils.sendReceiveMqttMsg


  # return response

end


"""

# Arguments
  
# Return

# Example
```jldoctest
julia> 
```

# TODO
  - [] update docstring
  - [PENDING] implement the function

# Signature
"""
function userRecommendbox(a::T1, input::T2) where {T1<:agent, T2<:AbstractString}
  error("--> userRecommendbox")

  # put in model format
  virtualWineCustomer = a.config[:externalservice][:virtualWineCustomer_1]
  llminfo = virtualWineCustomer[:llminfo]
  formattedinput = 
  if llminfo[:name] == "llama3instruct"
    formatLLMtext_llama3instruct("assistant", input)
  else
    error("llm model name is not defied yet $(@__LINE__)")
  end

  # send formatted input to user using GeneralUtils.sendReceiveMqttMsg


  # return response

end


""" Chatbox for chatting with virtual wine customer.

# Arguments
  - `a::T1`
    one of Yiem's agent
  - `input::T2`
    text to be send to virtual wine customer
  
# Return
  - `response::String`
    response of virtual wine customer
# Example
```jldoctest
julia> 
```

# TODO
  - [] update docstring
  - [] add reccommend() to compare wine

# Signature
"""
function virtualWineUserRecommendbox(a::T1, input
  )::Union{Tuple{String, Number, Number, Bool}, Tuple{String, Nothing, Number, Bool}} where {T1<:agent}

  # put in model format
  virtualWineCustomer = a.config[:externalservice][:virtualWineCustomer_1]
  llminfo = virtualWineCustomer[:llminfo]
  prompt = 
  if llminfo[:name] == "llama3instruct"
    formatLLMtext_llama3instruct("assistant", input)
  else
    error("llm model name is not defied yet $(@__LINE__)")
  end

  # send formatted input to user using GeneralUtils.sendReceiveMqttMsg
  msgMeta = GeneralUtils.generate_msgMeta(
    virtualWineCustomer[:mqtttopic],
    senderName= "virtualWineUserRecommendbox",
    senderId= a.id,
    receiverName= "virtualWineCustomer",
    mqttBroker= a.config[:mqttServerInfo][:broker],
    mqttBrokerPort= a.config[:mqttServerInfo][:port],
    msgId = "dummyid" #CHANGE remove after testing finished
    )

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

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

  return (response[:text], response[:select], response[:reward], response[:isterminal])
end



""" Chatbox for chatting with virtual wine customer.

# Arguments
  - `a::T1`
    one of Yiem's agent
  - `input::T2`
    text to be send to virtual wine customer
  
# Return
  - `response::String`
    response of virtual wine customer
# Example
```jldoctest
julia> 
```

# TODO
  - [] update docs 
  - [x] write a prompt for virtual customer 

# Signature
"""
function virtualWineUserChatbox(config::T1, input::T2, virtualCustomerChatHistory
  )::Union{Tuple{String, Number, Number, Bool}, Tuple{String, Nothing, Number, Bool}} where {T1<:AbstractDict, T2<:AbstractString}

  previouswines = 
  """
  You have the following wines previously:

  """

  systemmsg = 
  """
  You find yourself in a well-stocked wine store, engaged in a conversation with the store's knowledgeable sommelier. 
  You're on a quest to find a bottle of wine that aligns with your specific preferences and requirements. 
  
  The ideal wine you're seeking should meet the following criteria:
  1. It should fit within your budget.
  2. It should be suitable for the occasion you're planning.
  3. It should pair well with the food you intend to serve.
  4. It should be of a particular type of wine you prefer.
  5. It should possess certain characteristics, including:
    - The level of sweetness.
    - The intensity of its flavor.
    - The amount of tannin it contains.
    - Its acidity level.

  Here's the criteria details:
  {
    "budget": 50,
    "occasion": "graduation ceremony",
    "food pairing": "Thai food",
    "type of wine": "red",
    "wine sweetness level": "dry",
    "wine intensity level": "full-bodied",
    "wine tannin level": "low",
    "wine acidity level": "medium", 
  }

  You should only respond with "text", "select", "reward", "isterminal" steps.
  "text" is your conversation.
  "select" is an integer. Choose an option when presented with choices, or leave it null if none of the options satisfy you or if no choices are available.
  "reward" is an integer, it can be three number:
    1) 1 if you find the right wine.
    2) 0 if you don’t find the ideal wine.
    3) -1 if you’re dissatisfied with the sommelier’s response.
  "isterminal" can be false if you still want to talk with the sommelier, true otherwise.

  You should only respond in JSON format as describe below:
  {
    "text": "your conversation",
    "select": null,
    "reward": 0,
    "isterminal": false
  }

  Here are some examples:

  sommelier: "What's your budget?
  you:
  {
    "text": "My budget is 30 USD.",
    "select": null,
    "reward": 0,
    "isterminal": false
  }

  sommelier: "The first option is Zena Crown and the second one is Buano Red."
  you:
  {
    "text": "I like the 2nd option.",
    "select": 2,
    "reward": 1,
    "isterminal": true
  }

  Let's begin!
  """

  pushfirst!(virtualCustomerChatHistory, Dict(:name=> "system", :text=> systemmsg))
  
  # replace the :user key in chathistory to allow the virtual wine customer AI roleplay
  chathistory::Vector{Dict{Symbol, Any}} = Vector{Dict{Symbol, Any}}()
  for i in virtualCustomerChatHistory
    newdict = Dict()
    newdict[:name] =
    if i[:name] == "user"
      "you"
    elseif i[:name] == "assistant"
      "sommelier"
    else
      i[:name]
    end

    newdict[:text] = i[:text]
    push!(chathistory, newdict)
  end

  push!(chathistory, Dict(:name=> "assistant", :text=> input))

  # put in model format
  prompt = formatLLMtext(chathistory, "llama3instruct")
  prompt *=
    """
    <|start_header_id|>you<|end_header_id|>
    {"text"
    """

  pprint(prompt)
  externalService = config[:externalservice][:text2textinstruct]

  # send formatted input to user using GeneralUtils.sendReceiveMqttMsg
  msgMeta = GeneralUtils.generate_msgMeta(
    externalService[:mqtttopic],
    senderName= "virtualWineUserChatbox",
    senderId= string(uuid4()),
    receiverName= "text2textinstruct",
    mqttBroker= config[:mqttServerInfo][:broker],
    mqttBrokerPort= config[:mqttServerInfo][:port],
    msgId = string(uuid4()) #CHANGE remove after testing finished
  )

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

  attempt = 0
  for attempt in 1:5
    try
      response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=120)
      _responseJsonStr = response[:response][:text]
      expectedJsonExample =
      """
      Here is an expected JSON format:
      {
        "text": "...",
        "select": "...",
        "reward": "...",
        "isterminal": "..."
      }
      """
      responseJsonStr = jsoncorrection(config, _responseJsonStr, expectedJsonExample)
      responseDict = copy(JSON3.read(responseJsonStr)) 

      text::AbstractString = responseDict[:text]
      select::Union{Nothing, Number} = responseDict[:select] == "null" ? nothing : responseDict[:select]
      reward::Number = responseDict[:reward]
      isterminal::Bool = responseDict[:isterminal]

      if text != ""
        # pass test
      else
        error("virtual customer not answer correctly")
      end
      
      return (text, select, reward, isterminal)
    catch e
      io = IOBuffer()
      showerror(io, e)
      errorMsg = String(take!(io))
      st = sprint((io, v) -> show(io, "text/plain", v), stacktrace(catch_backtrace()))
      println("")
      @warn "Error occurred: $errorMsg\n$st"
      println("")
    end
  end
  error("virtualWineUserChatbox failed to get a response")
end

""" Search wine in stock.

# Arguments
  - `a::T1` 
    one of ChatAgent's agent.
  - `input::T2`
# Return
  A JSON string of available wine

# Example
```jldoctest
julia> using ChatAgent
julia> agent = ChatAgent.agentReflex("Jene")
julia> input = "{\"food\": \"pizza\", \"occasion\": \"anniversary\"}"
julia> result = winestock(agent, input)
"{"wine 1": {\"Winery\": \"Pichon Baron\", \"wine name\": \"Pauillac (Grand Cru Classé)\", \"grape variety\": \"Cabernet Sauvignon\", \"year\": 2010, \"price\": \"125 USD\", \"stock ID\": \"ar-17\"}, }"
```

# TODO
  - [] update docs
  - [WORKING] implement the function

# Signature
"""
function winestock(a::T1, input::T2
  )::Union{Tuple{String, Number, Number, Bool}, Tuple{String, Nothing, Number, Bool}} where {T1<:agent, T2<:AbstractString}

  wineattributes = extractWineAttributes(a, input)
  result = SQLLLM.query(Dict(:text=> wineattributes), a.executeSQL, a.text2textInstructLLM)
  return result
end


"""

# Arguments
  - `v::Integer`
    dummy variable
    
# Return

# Example
```jldoctest
julia> 
```

# TODO
  - [] update docstring
  - [x] implement the function

# Signature
"""
function extractWineAttributes(a::T1, input::T2
  )::String where {T1<:agent, T2<:AbstractString}

  converstiontable =
  """
  Conversion Table:
  Intensity level:
    Level 1: May correspond to "light-bodied" or a similar description.
    Level 2: May correspond to "med light", "medium light" or a similar description.
    Level 3: May correspond to "medium" or a similar description.
    Level 4: May correspond to "med full", "medium full" or a similar description.
    Level 5: May correspond to "full" or a similar description.
  Sweetness level:
    Level 1: May correspond to "dry", "no sweet" or a similar description.
    Level 2: May correspond to "off dry", "less sweet" or a similar description.
    Level 3: May correspond to "semi sweet" or a similar description.
    Level 4: May correspond to "sweet" or a similar description.
    Level 5: May correspond to "very sweet" or a similar description.
  Tannin level:
    Level 1: May correspond to "low tannin" or a similar description.
    Level 2: May correspond to "semi low tannin" or a similar description.
    Level 3: May correspond to "medium tannin" or a similar description.
    Level 4: May correspond to "semi high tannin" or a similar description.
    Level 5: May correspond to "high tannin" or a similar description.
  Acidity level:
    Level 1: May correspond to "low acidity" or a similar description.
    Level 2: May correspond to "semi low acidity" or a similar description.
    Level 3: May correspond to "medium acidity" or a similar description.
    Level 4: May correspond to "semi high acidity" or a similar description.
    Level 5: May correspond to "high acidity" or a similar description.
  """
  
  systemmsg =
  """
  As an attentive sommelier, your task is to determine the user's wine preferred levels of sweetness, intensity, tannin, acidity and other criteria based on the user query. 

  At each round of conversation, the user will give you the current situation:
  Conversion Table: ...
  Query: ...

  You should follow the following guidelines:
    - If specific information is unavailable, please use "NA" to indicate this.
    - Use converstion table to convert sweetness, acidity, tannin, intensity describing word into an integer.
    - Do not generate other comments.

  You should then respond to the user with the following points:
  - wine_type: Can be one of: red, white, sparkling, rose, dessert or fortified
  - price: ...
  - occasion: ...
  - food_paired: food that will be served with wine
  - country: wine's country of origin
  - grape_variety: ...
  - tasting_notes: wine's flavors
  - sweetness: S where S is an integer indicating sweetness level
  - acidity: A where A is an integer indicating acidity level
  - tannin: T where T is an integer indicating tannin level
  - intensity: I where I is an integer indicating intensity level

  You should only respond in format as described below:
  repeat: repeat the user's query
  wine_type: ...
  price: ...
  occasion: ...
  food_paired: ...
  country: ...
  grape_variety: ...
  tasting_notes: ...
  sweetness: 
  acidity: ...
  tannin: ...
  intensity: ...

  Here are some examples:

  user: "price < 25, for wedding party, full-bodied white wine with sweetness level 2, apple and honey notes, low tannin level and medium acidity level, Pizza, France, Riesling"
  assistant: repeat: ... \n wine_type: white\n price: less than 25\n occasion: wedding party\n food_pairing: Pizza\n country: France\n grape_variety: Riesling\n tasting_notes: apple, honey\n sweetness: 2\n acidity: 3\n tannin: 1\n intensity: 5

  user: body=full-bodied, dry, acidity=medium and low tannin
  assistant: repeat: ... \n wine_type: NA\n price: NA\n occasion: NA\n food_pairing: NA\n country: NA\n grape_variety: NA\n tasting_notes: NA\n sweetness: 1\n acidity: 3\n tannin: 1\n intensity: 5

  Let's begin!
  """

  usermsg =
  """
  Conversion Table: $converstiontable
  Query: $input
  """

  _prompt = 
  [
    Dict(:name=> "system", :text=> systemmsg),
    Dict(:name=> "user", :text=> usermsg)
  ]

  # put in model format
  prompt = GeneralUtils.formatLLMtext(_prompt, "llama3instruct")
  prompt *=
    """
    <|start_header_id|>assistant<|end_header_id|>
    """

  attributes = ["repeat", "wine_type", "price", "occasion", "food_paired", "country", "grape_variety", "sweetness", "acidity", "tannin", "intensity"]

  for attempt in 1:5
    try
      response = a.text2textInstructLLM(prompt)
      responsedict = GeneralUtils.textToDict(response, attributes, rightmarker=":", symbolkey=true)

      for i ∈ attributes
        if length(JSON3.write(responsedict[Symbol(i)])) == 0
          error("$i is empty ", @__LINE__)
        end
      end

      # check if there are more than 1 key per categories
      for i ∈ attributes
        matchkeys = GeneralUtils.findMatchingDictKey(responsedict, i)
        if length(matchkeys) > 1
          error("generatechat has more than one key per categories")
        end
      end

      result = ""
      for (k, v) in responsedict
        if k != :repeat && !occursin("NA", v)
          result *= "$k: $v, "
        end
      end

      result = result[1:end-2]  # remove the ending ", "

      # replace because SQLLLM didn't know what food_paired means
      result = replace(result, "food_paired" => "food_to_be_paired_with_wine")

      return result
    catch e
      io = IOBuffer()
      showerror(io, e)
      errorMsg = String(take!(io))
      st = sprint((io, v) -> show(io, "text/plain", v), stacktrace(catch_backtrace()))
      println("")
      println("Attempt $attempt. Error occurred: $errorMsg\n$st")
      println("")
    end
  end
  error("wineattributes_wordToNumber() failed to get a response")
end

""" Attemp to correct LLM response's incorrect JSON response.

# Arguments
  - `a::T1`
    one of Yiem's agent
  - `input::T2`
    text to be send to virtual wine customer
  
# Return
  - `correctjson::String`
    corrected json string  

# Example
```jldoctest
julia> 
```

# Signature
"""
function jsoncorrection(config::T1, input::T2, correctJsonExample::T3;
  maxattempt::Integer=3
  ) where {T1<:AbstractDict, T2<:AbstractString, T3<:AbstractString}

  incorrectjson = deepcopy(input)
  correctjson = nothing

  for attempt in 1:maxattempt
    try
      d = copy(JSON3.read(incorrectjson))
      correctjson = incorrectjson
      return correctjson
    catch e
      @warn "Attempting to correct JSON string. Attempt $attempt"
      e = """$e"""
      if occursin("EOF", e)
        e = split(e, "EOF")[1] * "EOF"
      end
      incorrectjson = deepcopy(input)
      _prompt =
      """
      Your goal are:
      1) Use the expected JSON format as a guideline to check why the given JSON string failed to load and provide a corrected version that can be loaded by Python's json.load function.
      2) Provide Corrected JSON string only. Do not provide any other info.

      $correctJsonExample

      Let's begin!
      Given JSON string: $incorrectjson
      The given JSON string failed to load previously because: $e
      Corrected JSON string:
      """

      # apply LLM specific instruct format
      externalService = config[:externalservice][:text2textinstruct]
      llminfo = externalService[:llminfo]
      prompt = 
      if llminfo[:name] == "llama3instruct"
        formatLLMtext_llama3instruct("system", _prompt)
      else
        error("llm model name is not defied yet $(@__LINE__)")
      end

      # send formatted input to user using GeneralUtils.sendReceiveMqttMsg
      msgMeta = GeneralUtils.generate_msgMeta(
        externalService[:mqtttopic],
        senderName= "jsoncorrection",
        senderId= string(uuid4()),
        receiverName= "text2textinstruct",
        mqttBroker= config[:mqttServerInfo][:broker],
        mqttBrokerPort= config[:mqttServerInfo][:port],
      )

      outgoingMsg = Dict(
        :msgMeta=> msgMeta,
        :payload=> Dict(
          :text=> prompt,
          :kwargs=> Dict(
            :max_tokens=> 512,
            :stop=> ["<|eot_id|>"],
            )
          )
        )
      result = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=120)
      incorrectjson = result[:response][:text]
    end
  end
end


























end # module llmfunction