# TODO: take $chat_history from tidyprompt object & add when sending to provider

#' Send a prompt to a LLM provider
#'
#' This function is responsible for sending prompts to a LLM provider for evaluation.
#' The function will interact with the LLM provider until a successful response
#' is received or the maximum number of interactions is reached. The function will
#' apply extraction and validation functions to the LLM response, as specified
#' in the prompt wraps (see [prompt_wrap()]). If the maximum number of interactions
#'
#' @param prompt A string or a [tidyprompt-class] object
#'
#' @param llm_provider [llm_provider-class] object
#'  (default is [llm_provider_ollama()]).
#' This object and its settings will be used to evaluate the prompt.
#' You may also pass an `ellmer::chat()` object directly; in that case,
#' [send_prompt()] will clone it, clear any existing turns, and wrap the clean
#' clone with [llm_provider_ellmer()] before evaluation.
#' When using [llm_provider_ellmer()] or a raw `ellmer::chat()`, the working
#' `ellmer` turns are rebuilt from tidyprompt's own chat history for each call.
#' Note that the 'verbose' and 'stream' settings in the LLM provider will be
#'  overruled by the 'verbose' and 'stream' arguments in this function
#'  when those are not NULL.
#' Furthermore, advanced [tidyprompt-class] objects may carry '$parameter_fn'
#'  functions which can set parameters in the llm_provider object
#'  (see [prompt_wrap()] and [llm_provider-class] for more ).
#'
#' @param max_interactions Maximum number of interactions allowed with the
#' LLM provider. Default is 10. If the maximum number of interactions is reached
#' without a successful response, 'NULL' is returned as the response (see return
#' value). The first interaction is the initial chat completion
#' @param clean_chat_history If the chat history should be cleaned after each
#' interaction. Cleaning the chat history means that only the
#' first and last message from the user, the last message from the assistant,
#' all messages from the system, and all tool results are kept in a 'clean'
#' chat history. This clean chat history is used when requesting a new chat completion.
#' Rows marked as non-replayable are excluded from new requests regardless of
#' this setting, so the returned transcript may contain more rows than the model
#' actually sees on a retry or follow-up call.
#' (i.e., if a LLM repeatedly fails to provide a correct response, only its last failed response
#' will included in the context window). This may increase the LLM performance
#' on the next interaction
#' @param verbose If the interaction with the LLM provider should be printed
#' to the console. This will overrule the 'verbose' setting in the LLM provider
#' @param stream If the interaction with the LLM provider should be streamed.
#' This setting will only be used if the LLM provider already has a
#' 'stream' parameter (which indicates there is support for streaming). Note
#' that when 'verbose' is set to FALSE, the 'stream' setting will be ignored
#' @param return_mode One of 'full' or 'only_response'. See return value
#' @return \itemize{
#'  \item If return mode 'only_response', the function will return only the LLM response
#' after extraction and validation functions have been applied (NULL is returned
#' when unsuccessful after the maximum number of interactions).
#'  \item If return mode 'full', the function will return a list with the following elements:
#'  \itemize{
#'    \item 'response' (the LLM response after extraction and validation functions have been applied;
#'  NULL is returned when unsuccessful after the maximum number of interactions),
#'    \item 'interactions' (the number of interactions with the LLM provider),
#'    \item 'chat_history' (a dataframe with the full chat history which led to the final response.
#'    This may include rows that are retained for inspection but not re-sent to the model;
#'    such rows are marked with column `hidden_from_llm = TRUE`),
#'    \item 'chat_history_clean' (a dataframe with the cleaned chat history which led to
#' the final response; here, only the first and last message from the user, the
#' last message from the assistant, and all messages from the system are kept,
#' after excluding any non-replayable rows),
#'    \item 'start_time' (the time when the function was called),
#'    \item 'end_time' (the time when the function ended),
#'    \item 'duration_seconds' (the duration of the function in seconds),
#'    \item 'http' (a list with all HTTP requests and responses made during the interactions;
#'    as returned by `llm_provider$complete_chat()`),
#'    \item 'ellmer_chat' (if [llm_provider_ellmer()] or a raw `ellmer::chat()`
#'    object was used, this will be
#'    the updated 'ellmer' chat object, containing for instance
#'    the turns and possible tool calls. (As this function
#'    uses a clone of the provided LLM provider, the 'ellmer' chat object in the
#'    LLM provider will not be updated; use this return value when you need the
#'    updated native `ellmer` turns from the current evaluation. Note that turns in the 'ellmer'
#'    chat object may not contain the full chat history when `clean_chat_history = TRUE`
#'    was used.)
#'  }
#' }
#'
#'
#' @export
#' @example inst/examples/send_prompt.R
#'
#' @seealso [tidyprompt-class], [prompt_wrap()], [llm_provider-class], [llm_provider_ollama()],
#' [llm_provider_openai()]
#'
#' @family prompt_evaluation
send_prompt <- function(
  prompt,
  llm_provider = llm_provider_ollama(),
  max_interactions = 10,
  clean_chat_history = FALSE,
  verbose = NULL,
  stream = NULL,
  return_mode = c("only_response", "full")
) {
  ## 1 Validate arguments

  # Basic validation
  prompt <- tidyprompt(prompt)
  return_mode <- match.arg(return_mode)
  llm_provider <- as_send_prompt_llm_provider(llm_provider, verbose, stream)
  stopifnot(
    inherits(llm_provider, "LlmProvider"),
    max_interactions > 0,
    max_interactions == floor(max_interactions),
    is.logical(clean_chat_history),
    is.null(verbose) | is.logical(verbose),
    is.null(stream) | is.logical(stream)
  )

  # Add provider-specific pre/post prompt wraps
  if (isTRUE(is.function(llm_provider$apply_prompt_wraps))) {
    prompt <- llm_provider$apply_prompt_wraps(prompt)
  }

  # Verify and configure llm_provider
  llm_provider <- llm_provider$clone()
  # Deep-clone the ellmer chat object so mutations (register_tool, set_turns)
  # don't affect the caller's original provider. When send_prompt() received a
  # raw ellmer chat directly, reset the working clone before evaluation.
  if (!is.null(llm_provider[["ellmer_chat"]])) {
    if (isTRUE(llm_provider$parameters$.reset_ellmer_chat)) {
      llm_provider$ellmer_chat <- ellmer_chat_clone_reset(
        llm_provider$ellmer_chat,
        context = paste0(
          "When passing an ellmer chat directly to `send_prompt()`, `llm_provider`"
        )
      )
      llm_provider$parameters$.reset_ellmer_chat <- NULL
    } else if (is.function(llm_provider$ellmer_chat$clone)) {
      llm_provider$ellmer_chat <- llm_provider$ellmer_chat$clone()
    }
  }
  if (!is.null(verbose)) {
    llm_provider$verbose <- verbose
  }
  if (
    !is.null(stream) && !is.null(llm_provider$parameters$stream) # This means the provider supports streaming
  ) {
    llm_provider$parameters$stream <- stream
  }
  # Apply parameter_fn's to the llm_provider
  for (prompt_wrap in get_prompt_wraps(prompt)) {
    if (!is.null(prompt_wrap$parameter_fn)) {
      parameter_fn <- prompt_wrap$parameter_fn
      llm_provider$set_parameters(parameter_fn(llm_provider))
    }
  }
  # Add handler_fn's to the llm_provider
  for (prompt_wrap in get_prompt_wraps(prompt)) {
    if (!is.null(prompt_wrap$handler_fn)) {
      llm_provider$add_handler_fn(prompt_wrap$handler_fn)
    }
  }

  # Initialize variables which keep track of the process
  if (return_mode == "full") {
    start_time <- Sys.time()
  }
  http <- list(requests = list(), responses = list())
  # Object which keeps ellmer chat object, turns, structured output;
  #   when using an ellmer LLM provider:
  ellmer_chat <- NULL

  ## 2 Chat history, send_chat, handler_fns

  chat_history <- prompt$get_chat_history(llm_provider)

  # Internal function to send chat messages
  send_chat <- function(
    message,
    role = "user",
    tool_result = FALSE
  ) {
    merge_completed_chat_history <- function(
      original_chat_history,
      completed_chat_history,
      request_rows
    ) {
      updated_chat_history <- dplyr::bind_rows(
        original_chat_history,
        completed_chat_history[0, , drop = FALSE]
      )
      updated_chat_history <- updated_chat_history[
        seq_len(nrow(original_chat_history)),
        ,
        drop = FALSE
      ]

      request_rows <- as.integer(request_rows %||% integer())
      request_rows <- request_rows[
        !is.na(request_rows) &
          request_rows >= 1L &
          request_rows <= nrow(updated_chat_history)
      ]

      sent_rows <- min(length(request_rows), nrow(completed_chat_history))
      if (sent_rows > 0L) {
        updated_chat_history[
          request_rows[seq_len(sent_rows)],
          names(completed_chat_history)
        ] <- completed_chat_history[seq_len(sent_rows), , drop = FALSE]
      }

      if (nrow(completed_chat_history) > sent_rows) {
        updated_chat_history <- dplyr::bind_rows(
          updated_chat_history,
          completed_chat_history[
            (sent_rows + 1L):nrow(completed_chat_history),
            ,
            drop = FALSE
          ]
        )
      }

      updated_chat_history
    }

    if (!is.null(message)) {
      message <- as.character(message)
      chat_history <<- chat_history |>
        add_msg_to_chat_history(message, role, tool_result)
    }

    if (clean_chat_history) {
      cleaned_chat_history <- clean_chat_history(chat_history)
      response <- llm_provider$complete_chat(
        list(chat_history = cleaned_chat_history)
      )
      chat_history <<- merge_completed_chat_history(
        chat_history,
        response$completed,
        attr(cleaned_chat_history, "source_rows") %||%
          as.integer(rownames(cleaned_chat_history))
      )
    } else {
      response <- llm_provider$complete_chat(list(chat_history = chat_history))
      chat_history <<- merge_completed_chat_history(
        chat_history,
        response$completed,
        seq_len(nrow(chat_history))
      )
    }

    for (http_response in response$http$response) {
      http$responses[[length(http$responses) + 1]] <<- http_response
    }
    for (http_request in response$http$request) {
      http$requests[[length(http$requests) + 1]] <<- http_request
    }

    if (!is.null(response$ellmer_chat)) {
      ellmer_chat <<- response$ellmer_chat
    }

    utils::tail(chat_history$content, 1)
  }

  ## 3 Retrieve initial response

  response <- send_chat(NULL)
  # (NULL as initial message is already included in the chat_history)

  ## 4 Apply extractions and validations

  prompt_wraps <- get_prompt_wraps(prompt, order = "evaluation")
  # (Tools, then modes, then unspecified prompt_wraps)

  interactions <- 1
  success <- FALSE
  while (interactions < max_interactions && !success) {
    interactions <- interactions + 1

    if (length(prompt_wraps) == 0) {
      success <- TRUE
    }

    # Initialize variables for the loop
    any_prompt_wrap_not_done <- FALSE
    llm_break <- FALSE

    for (pw_index in seq_along(prompt_wraps)) {
      prompt_wrap <- prompt_wraps[[pw_index]]

      # Apply extraction function
      if (!is.null(prompt_wrap$extraction_fn)) {
        extraction_function <- prompt_wrap$extraction_fn

        # Set the environment of the extraction function, if an
        #   environment was given as attribute. At the moment, this is used
        #   for passing 'tool_functions' from add_tools() to execution here
        if (!is.null(attr(extraction_function, "environment"))) {
          environment <- attr(extraction_function, "environment")
          environment(extraction_function) <- environment
        }
        extraction_result <-
          extraction_function(response, llm_provider, http)

        # If it inherits llm_feedback,
        #   send the feedback to the LLM & get new response
        if (
          inherits(extraction_result, "llm_feedback") ||
            inherits(extraction_result, "llm_feedback_tool_result")
        ) {
          if (inherits(extraction_result, "llm_feedback_tool_result")) {
            # This ensures tool results are not filtered out when cleaning
            #   the context window in send_prompt()
            response <- send_chat(extraction_result, tool_result = TRUE)
          } else {
            response <- send_chat(extraction_result, tool_result = FALSE)
          }
          any_prompt_wrap_not_done <- TRUE
          break
        }

        if (inherits(extraction_result, "llm_break_soft")) {
          interactions <- max_interactions
        }

        if (inherits(extraction_result, "llm_break")) {
          # Still apply remaining prompt wraps of type 'check'
          #   (this may block the break when feedback is returned)
          if (pw_index < length(prompt_wraps)) {
            prompt_wraps_remaining <- prompt_wraps[
              (pw_index + 1):length(prompt_wraps)
            ]
            for (prompt_wrap_remaining in prompt_wraps_remaining) {
              if (prompt_wrap_remaining$type == "check") {
                check_result <-
                  prompt_wrap_remaining$validation_fn(
                    response,
                    llm_provider,
                    http
                  )
                if (inherits(check_result, "llm_feedback")) {
                  response <- send_chat(check_result, tool_result = TRUE)
                  any_prompt_wrap_not_done <- TRUE
                  break
                }
              }
            }
          }

          if (!extraction_result$success) {
            any_prompt_wrap_not_done <- TRUE # Will result in no success
          } else {
            any_prompt_wrap_not_done <- FALSE # Will result in success
          }
          response <- extraction_result$object_to_return
          llm_break <- TRUE
          break
        }

        # If no llm_feedback or break, extraction was succesful
        response <- extraction_result
      }

      # Apply validation function
      if (!is.null(prompt_wrap$validation_fn)) {
        validation_result <-
          prompt_wrap$validation_fn(response, llm_provider, http)

        # If it inherits llm_feedback, send the feedback to the LLM & get new response
        if (inherits(validation_result, "llm_feedback")) {
          # Set as 'tool_result' to not clear user feedback when type = 'check'
          tool_result <- FALSE
          if (prompt_wrap$type == "check") {
            tool_result <- TRUE
          }

          response <- send_chat(validation_result, tool_result = tool_result)
          any_prompt_wrap_not_done <- TRUE
          break
        }

        if (inherits(validation_result, "llm_break_soft")) {
          interactions <- max_interactions
        }

        if (inherits(validation_result, "llm_break")) {
          # Still apply remaining prompt wraps of type 'check'
          #   (this may block the break when feedback is returned)
          if (pw_index < length(prompt_wraps)) {
            prompt_wraps_remaining <- prompt_wraps[
              (pw_index + 1):length(prompt_wraps)
            ]
            for (prompt_wrap_remaining in prompt_wraps_remaining) {
              if (prompt_wrap_remaining$type == "check") {
                check_result <-
                  prompt_wrap_remaining$validation_fn(
                    response,
                    llm_provider,
                    http
                  )
                if (inherits(check_result, "llm_feedback")) {
                  response <- send_chat(check_result, tool_result = TRUE)
                  any_prompt_wrap_not_done <- TRUE
                  break
                }
              }
            }
          }

          if (!validation_result$success) {
            any_prompt_wrap_not_done <- TRUE # Will result in no success
          } else {
            any_prompt_wrap_not_done <- FALSE # Will result in success
          }
          response <- validation_result$object_to_return
          llm_break <- TRUE
          break
        }
      }
    }

    if (!any_prompt_wrap_not_done) {
      success <- TRUE
    }

    if (llm_break) break
  }

  ## 5 Final evaluation

  if (!success) {
    warning(
      paste0(
        "Failed to reach a valid answer after ",
        interactions,
        " interactions"
      )
    )
    response <- NULL
  }

  if (return_mode == "only_response") {
    return(response)
  }

  if (return_mode == "full") {
    return_list <- list()

    chat_history <- normalize_chat_history_metadata(chat_history)

    return_list$response <- response
    return_list$interactions <- interactions
    return_list$chat_history <- chat_history

    if (clean_chat_history) {
      return_list$chat_history_clean <- clean_chat_history(chat_history)
    }

    return_list$start_time <- start_time
    return_list$end_time <- Sys.time()

    return_list$duration_seconds <-
      as.numeric(
        difftime(
          return_list$end_time,
          return_list$start_time,
          units = "secs"
        )
      )

    return_list$http <- http
    return_list$ellmer_chat <- ellmer_chat

    return(return_list)
  }
}

#' Create chat history dataframe
#'
#' @description Internal function used by [send_prompt()] to create a chat history
#' dataframe. This dataframe is used to keep track of the chat history during
#' the interactions with the LLM provider.
#
#' @param role Role of the message (e.g. "user", "assistant", "system")
#' @param content Content of the message
#' @param tool_result Logical indicating whether the message is a tool result.
#' This will ensure it is not removed when the chat history context window is
#' cleaned by `clean_chat_history()` (which is an internal function which will keep only the
#' first and last message from the user, the last message from the assistant,
#' all messages from the system, and all tool_result messages)
#'
#' @return A dataframe with the chat history
#'
#' @noRd
#' @keywords internal
create_chat_df <- function(
  role = character(),
  content = character(),
  tool_result = logical()
) {
  data.frame(role = role, content = content, tool_result = tool_result)
}

#' Clean chat history
#'
#' @description Internal function used by [send_prompt()] to clean the chat history
#' context window after each interaction with the LLM provider.
#'
#' This function will keep only the first and last message from the user, the last
#' message from the assistant, all messages from the system, and all tool_result messages.
#'
#' This may be useful to keep the context window clean and increase the LLM's performance.
#'
#' @param chat_history A dataframe with the full chat history
#'
#' @return A dataframe with the cleaned chat history
#'
#' @noRd
#' @keywords internal
clean_chat_history <- function(chat_history) {
  filtered <- chat_history_to_send(chat_history)
  # source_rows maps positions in `filtered` back to the original
  # chat_history rows (accounting for hidden / tool_call filtering).
  original_source_rows <- attr(filtered, "source_rows") %||%
    seq_len(nrow(filtered))

  # Keep only first and last message from user;
  # keep only last message from assistant;
  # keep all messages from system;
  # keep all tool_result messages
  user_rows <- which(filtered$role == "user")
  assistant_rows <- which(filtered$role == "assistant")
  system_rows <- which(filtered$role == "system")
  tool_result_rows <- which(
    filtered$tool_result | filtered$role == "tool"
  )

  keep_rows <- c(
    system_rows,
    user_rows[c(1, length(user_rows))],
    utils::tail(assistant_rows, 1),
    tool_result_rows
  )

  keep_rows <- sort(unique(keep_rows))

  # Subset the dataframe with these rows
  cleaned_chat_history <- filtered[keep_rows, ]
  # Compose the two mappings: keep_rows indexes into `filtered`, whose

  # positions map back to the original frame via original_source_rows.
  attr(cleaned_chat_history, "source_rows") <- original_source_rows[keep_rows]
  # (sort(unique()) is used to ensure that the rows are in order)

  return(cleaned_chat_history)
}
