module llmUtil

export formatLLMtext, formatLLMtext_llama3instruct, jsoncorrection, deFormatLLMtext, extractthink

using UUIDs, JSON, Dates
using GeneralUtils

# ---------------------------------------------- 100 --------------------------------------------- #
#[PENDING] update code to use JSON

""" Convert a single chat dictionary into LLM model instruct format.

# Llama 3 instruct format example
  <|begin_of_text|>
  <|start_header_id|>system<|end_header_id|>
  You are a helpful assistant.
  <|eot_id|>
  <|start_header_id|>user<|end_header_id|>
  Get me an icecream.
  <|eot_id|>
  <|start_header_id|>assistant<|end_header_id|>
  Go buy it yourself at 7-11.
  <|eot_id|>

# Arguments
  - `name::T`
    message owner name e.f. "system", "user" or "assistant"
  - `text::T`

# Return
  - `formattedtext::String`
    text formatted to model format

# Example
```jldoctest
julia> using Revise
julia> using YiemAgent
julia> d = Dict(:name=> "system",:text=> "You are a helpful, respectful and honest assistant.",)
julia> formattedtext = YiemAgent.formatLLMtext_llama3instruct(d[:name], d[:text])
"<|begin_of_text|>\n  <|start_header_id|>system<|end_header_id|>\n  You are a helpful, respectful and honest assistant.\n  <|eot_id|>\n"
```

Signature
"""
function formatLLMtext_llama3instruct(name::T, text::T;
  assistantStarter::Bool=false) where {T<:AbstractString}
  formattedtext = 
  if name == "system"
    """
    <|start_header_id|>$name<|end_header_id|>
    $text
    <|eot_id|>
    """
  else
    """
    <|start_header_id|>$name<|end_header_id|>
    $text
    <|eot_id|>
    """
  end

  if assistantStarter
    formattedtext *=
      """
      <|start_header_id|>assistant<|end_header_id|>
      """
  end

  return formattedtext
end


function formatLLMtext_qwen(name::T, text::T;
  assistantStarter::Bool=false) where {T<:AbstractString}
  formattedtext = 
  if name == "system"
    """
    <|im_start|>$name
    $text
    <|im_end|>
    """
  else
    """
    <|im_start|>$name
    $text
    <|im_end|>
    """
  end

  if assistantStarter
    formattedtext *=
      """
      <|im_start|>assistant
      """
  end

  return formattedtext
end


function formatLLMtext_qwen3(name::T, text::T;
  assistantStarter::Bool=false) where {T<:AbstractString}
  formattedtext = 
  if name == "system"
    """
    <|im_start|>$name
    $text
    <|im_end|>
    """
  else
    """
    <|im_start|>$name
    $text
    <|im_end|>
    """
  end

  if assistantStarter
    formattedtext *=
      """
      <|im_start|>assistant
      """
  end

  return formattedtext
end


function formatLLMtext_phi4(name::T, text::T;
  assistantStarter::Bool=false) where {T<:AbstractString}
  formattedtext = 
  if name == "system"
    """
    <|system|>
    $text
    <|end|>
    """
  else
    """
    <|assistant|>
    $text
    <|end|>
    """
  end

  if assistantStarter
    formattedtext *=
      """
      <|assistant|>
      """
  end

  return formattedtext
end


function formatLLMtext_granite3(name::T, text::T;
  assistantStarter::Bool=false) where {T<:AbstractString}
  formattedtext = 
  if name == "system"
    """
    <|start_of_role|>system<|end_of_role|>{$text}<|end_of_text|>
    """
  else
    """
    <|start_of_role|>$name<|end_of_role|>{$text}<|end_of_text|>
    """
  end

  if assistantStarter
    formattedtext *=
      """
      <|start_of_role|>assistant<|end_of_role|>{
      """
  end

  return formattedtext
end



""" Convert a vector of chat message dictionaries into LLM model instruct format.

# Arguments
  - `messages::Vector{Dict{Symbol, T}}`
    A vector of dictionaries where each dictionary contains the keys `:name` (the name of the message owner) and `:text` (the text of the message).
  - `formatname::T`
    The name of the format to be used for converting the chat messages.
# Return
  - `formattedtext::String`
    text formatted to model format

# Example
```jldoctest
julia> using Revise
julia> using YiemAgent
julia> chatmessage = [
          Dict(:name=> "system",:text=> "You are a helpful, respectful and honest assistant.",),
          Dict(:name=> "user",:text=> "list me all planets in our solar system.",),
          Dict(:name=> "assistant",:text=> "I'm sorry. I don't know. You tell me.",),
        ]
julia> formattedtext = YiemAgent.formatLLMtext(chatmessage, "llama3instruct")
"<|begin_of_text|>\n  <|start_header_id|>system<|end_header_id|>\n  You are a helpful, respectful and honest assistant.\n  <|eot_id|>\n  <|start_header_id|>user<|end_header_id|>\n  list me all planets in our solar system.\n  <|eot_id|>\n  <|start_header_id|>assistant<|end_header_id|>\n  I'm sorry. I don't know. You tell me.\n  <|eot_id|>\n"
```
"""
function formatLLMtext(messages::Vector{Dict{Symbol, T}}, formatname::String
  )::String where {T<:AbstractString}
  f = 
    if formatname == "llama3instruct"
      formatLLMtext_llama3instruct
    elseif formatname == "mistral"
      # not define yet
    elseif formatname == "phi3instruct"
      # not define yet
    elseif formatname == "qwen"
      formatLLMtext_qwen
    elseif formatname == "qwen3"
      formatLLMtext_qwen3
    elseif formatname == "phi4"
      formatLLMtext_phi4
    elseif formatname == "granite3"
      formatLLMtext_granite3
    else
      error("$formatname template not define yet")
    end

  str = ""
  for (i, t) in enumerate(messages)
    if i < length(messages)
      str *= f(t[:name], t[:text])
    else
      str *= f(t[:name], t[:text]; assistantStarter=true)
    end
  end

  return str
end

""" Revert LLM-format response back into regular text.

# Arguments
  - `text::String`
    The LLM formatted string to be converted.

# Return
  - `normalText::String`
    The original plain text extracted from the given LLM-formatted string.

# Example
```jldoctest
julia> using Revise
julia> using YiemAgent
julia> response = "<|begin_of_text|>This is a sample system instruction.<|eot_id|>"
julia> normalText = YiemAgent.deFormatLLMtext(response, "granite3")
"This is a sample system instruction."
```
"""
function deFormatLLMtext(text::String, formatname::String; includethink::Bool=false
  )::String
  f = 
    if formatname == "granite3"
      deFormatLLMtext_granite3
    elseif formatname == "qwen3"
      deFormatLLMtext_qwen3
    else
      error("$formatname template not define yet")
    end

  r = f(text)
  result = r === nothing ? text : r
  return result
end


""" Revert LLM-format response back into regular text for Granite 3 format.

# Arguments
  - `text::String`
    The LLM formatted string to be converted.

# Return
  - `normalText::Union{Nothing, String}`
    The original plain text extracted from the given LLM-formatted string.
    Returns nothing if the text is not in Granite 3 format.

# Example
```jldoctest
julia> using Revise
julia> using YiemAgent
julia> response = "{This is a sample LLM response.}"
julia> normalText = YiemAgent.deFormatLLMtext(response, "granite3")
"This is a sample LLM response."
"""
function deFormatLLMtext_granite3(text::String)::Union{Nothing, String}
  # check if '{' and '}' are in the text because it's a special format for the LLM response
  if contains(text, "<|im_start|>assistant")
    # get the text between '{' and '}'
    text_between_braces = GeneralUtils.extractTextBetweenCharacter(text, '{', '}')[1]
    return text_between_braces
  elseif text[end] == '}'
    text = "{$text"
    text_between_braces = GeneralUtils.extractTextBetweenCharacter(text, '{', '}')[1]
  else
    return nothing
  end
end


function deFormatLLMtext_qwen3(text::String)::Union{Nothing, String}
  return text
end

# function deFormatLLMtext_qwen3(text::String; includethink::Bool=false)::Union{Nothing, String}
#   think = nothing
#   str = nothing

#   if occursin("<think>", text)
#     r = GeneralUtils.extractTextBetweenString(text, "<think>", "</think>")
#     if r[:success]
#       think = r[:text]
#     end
#     str = string(split(text, "</think>")[2])
#   end

#   if includethink == true && occursin("<think>", text)
#     result = "ModelThought: $think $str"
#     return result
#   elseif includethink == false && occursin("<think>", text)
#     result = str
#     return result
#   else
#     return text
#   end
# 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= uuid4snakecase(),
        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


function extractthink(text::String)
  think = nothing
  str = nothing
  if occursin("<think>", text)
    r = GeneralUtils.extractTextBetweenString(text, "<think>", "</think>")
    if r[:success]
      think = r[:text]
    end
    str = string(split(text, "</think>")[2])
  else
    str = text
  end
  return think, str
end

































end # module llmUtil