defmodule LangChain.Chains.LLMChain do
  @doc """
  Define an LLMChain. This is the heart of the LangChain library.

  The chain deals with functions, a function map, delta tracking, last_message
  tracking, conversation messages, and verbose logging. This helps by separating
  these responsibilities from the LLM making it easier to support additional
  LLMs because the focus is on communication and formats instead of all the
  extra logic.



  """
  use Ecto.Schema
  import Ecto.Changeset
  require Logger
  alias LangChain.PromptTemplate
  alias __MODULE__
  alias LangChain.Message
  alias LangChain.MessageDelta
  alias LangChain.Function
  alias LangChain.LangChainError

  @primary_key false
  embedded_schema do
    field :llm, :any, virtual: true
    field :verbose, :boolean, default: false
    field :functions, {:array, :any}, default: [], virtual: true
    # set and managed privately through functions
    field :function_map, :map, default: %{}, virtual: true

    # List of `Message` structs for creating the conversation with the LLM.
    field :messages, {:array, :any}, default: [], virtual: true

    # Custom context data made available to functions when executed.
    # Could include information like account ID, user data, etc.
    field :custom_context, :any, virtual: true

    # Track the current merged `%MessageDelta{}` struct received when streamed.
    # Set to `nil` when there is no current delta being tracked. This happens
    # when the final delta is received that completes the message. At that point,
    # the delta is converted to a message and the delta is set to nil.
    field :delta, :any, virtual: true
    # Track the last `%Message{}` received in the chain.
    field :last_message, :any, virtual: true
    # Track if the state of the chain expects a response from the LLM. This
    # happens after sending a user message or when a function_call is received,
    # we've provided a function response and the LLM needs to respond.
    field :needs_response, :boolean, default: false

    # A callback function to execute when messages are added. Don't allow caller
    # to setup in `.new` function. Want to set it from the `.run` function to
    # avoid multiple chain instances (across processes) from both firing
    # callbacks.
    field :callback_fn, :any, virtual: true
  end

  @type t :: %LLMChain{}

  @create_fields [:llm, :functions, :custom_context, :verbose]
  @required_fields [:llm]

  @doc """
  Start a new LLMChain configuration.
  """
  @spec new(attrs :: map()) :: {:ok, t} | {:error, Ecto.Changeset.t()}
  def new(attrs \\ %{}) do
    %LLMChain{}
    |> cast(attrs, @create_fields)
    |> common_validation()
    |> apply_action(:insert)
  end

  @doc """
  Start a new LLMChain configuration and return it or raise an error if invalid.
  """
  @spec new!(attrs :: map()) :: t() | no_return()
  def new!(attrs \\ %{}) do
    case new(attrs) do
      {:ok, chain} ->
        chain

      {:error, changeset} ->
        raise LangChainError, changeset
    end
  end

  def common_validation(changeset) do
    changeset
    |> validate_required(@required_fields)
    |> validate_llm_is_struct()
    |> build_functions_map_from_functions()
  end

  defp validate_llm_is_struct(changeset) do
    case get_change(changeset, :llm) do
      nil -> changeset
      llm when is_struct(llm) -> changeset
      _other -> add_error(changeset, :llm, "LLM must be a struct")
    end
  end

  @doc false
  def build_functions_map_from_functions(changeset) do
    functions = get_field(changeset, :functions, [])

    # get a list of all the functions indexed into a map by name
    fun_map =
      Enum.reduce(functions, %{}, fn f, acc ->
        Map.put(acc, f.name, f)
      end)

    put_change(changeset, :function_map, fun_map)
  end

  @doc """
  Add more functions to an LLMChain.
  """
  @spec add_functions(t(), Function.t() | [Function.t()]) :: t() | no_return()
  def add_functions(%LLMChain{} = chain, %Function{} = function) do
    add_functions(chain, [function])
  end

  def add_functions(%LLMChain{functions: existing} = chain, functions) when is_list(functions) do
    updated = existing ++ functions

    chain
    |> change()
    |> cast(%{functions: updated}, [:functions])
    |> build_functions_map_from_functions()
    |> apply_action!(:update)
  end

  @doc """
  Run the chain on the LLM using messages and any registered functions. This
  formats the request for a ChatLLMChain where messages are passed to the API.

  When successful, it returns `{:ok, updated_chain, message_or_messages}`

  ## Options

  - `:while_needs_response` - repeatedly evaluates functions and submits to the
    LLM so long as we still expect to get a response.
  - `:callback_fn` - the callback function to execute as messages are received.

  """
  @spec run(t(), Keyword.t()) :: {:ok, t(), Message.t() | [Message.t()]} | {:error, String.t()}
  def run(chain, opts \\ [])

  def run(%LLMChain{} = chain, opts) do
    # set the callback function on the chain
    chain = %LLMChain{chain | callback_fn: Keyword.get(opts, :callback_fn)}

    if chain.verbose, do: IO.inspect(chain.llm, label: "LLM")

    if chain.verbose, do: IO.inspect(chain.messages, label: "MESSAGES")

    functions = chain.functions
    if chain.verbose, do: IO.inspect(functions, label: "FUNCTIONS")

    if Keyword.get(opts, :while_needs_response, false) do
      run_while_needs_response(chain)
    else
      # run the chain and format the return
      case do_run(chain) do
        {:ok, chain} ->
          {:ok, chain, chain.last_message}

        {:error, _reason} = error ->
          error
      end
    end
  end

  # Repeatedly run the chain while `needs_response` is true. This will execute
  # functions and re-submit the function result to the LLM giving the LLM an
  # opportunity to execute more functions or return a response.
  @spec run_while_needs_response(t()) :: {:ok, t(), Message.t()} | {:error, String.t()}
  defp run_while_needs_response(%LLMChain{needs_response: false} = chain) do
    {:ok, chain, chain.last_message}
  end

  defp run_while_needs_response(%LLMChain{needs_response: true} = chain) do
    chain
    |> execute_function()
    |> do_run()
    |> case do
      {:ok, updated_chain} ->
        run_while_needs_response(updated_chain)

      {:error, reason} ->
        {:error, reason}
    end
  end

  # internal reusable function for running the chain
  @spec do_run(t()) :: {:ok, t()} | {:error, String.t()}
  defp do_run(%LLMChain{} = chain) do
    # submit to LLM. The "llm" is a struct. Match to get the name of the module
    # then execute the `.call` function on that module.
    %module{} = chain.llm

    # handle and output response
    case module.call(chain.llm, chain.messages, chain.functions, chain.callback_fn) do
      {:ok, [%Message{} = message]} ->
        if chain.verbose, do: IO.inspect(message, label: "SINGLE MESSAGE RESPONSE")
        {:ok, add_message(chain, message)}

      {:ok, [%Message{} = message, _others] = messages} ->
        if chain.verbose, do: IO.inspect(messages, label: "MULTIPLE MESSAGE RESPONSE")
        # return the list of message responses. Happens when multiple
        # "choices" are returned from LLM by request.
        {:ok, add_message(chain, message)}

      {:ok, [[%MessageDelta{} | _] | _] = deltas} ->
        if chain.verbose, do: IO.inspect(deltas, label: "DELTA MESSAGE LIST RESPONSE")
        {:ok, apply_deltas(chain, deltas)}

      {:error, reason} ->
        if chain.verbose, do: IO.inspect(reason, label: "ERROR")
        Logger.error("Error during chat call. Reason: #{inspect(reason)}")
        {:error, reason}
    end
  end

  @doc """
  Update the LLMChain's `custom_context` map. Passing in a `context_update` map
  will by default merge the map into the existing `custom_context`.

  Use the `:as` option to:
  - `:merge` - Merge update changes in. Default.
  - `:replace` - Replace the context with the `context_update`.
  """
  @spec update_custom_context(t(), context_update :: %{atom() => any()}, opts :: Keyword.t()) ::
          t() | no_return()
  def update_custom_context(chain, context_update, opts \\ [])

  def update_custom_context(
        %LLMChain{custom_context: %{} = context} = chain,
        %{} = context_update,
        opts
      ) do
    new_context =
      case Keyword.get(opts, :as) || :merge do
        :merge ->
          Map.merge(context, context_update)

        :replace ->
          context_update

        other ->
          raise LangChain.LangChainError,
                "Invalid update_custom_context :as option of #{inspect(other)}"
      end

    %LLMChain{chain | custom_context: new_context}
  end

  def update_custom_context(
        %LLMChain{custom_context: nil} = chain,
        %{} = context_update,
        _opts
      ) do
    # can't merge a map with `nil`. Replace it.
    %LLMChain{chain | custom_context: context_update}
  end

  @doc """
  Apply a received MessageDelta struct to the chain. The LLMChain tracks the
  current merged MessageDelta state. When the final delta is received that
  completes the message, the LLMChain is updated to clear the `delta` and the
  `last_message` and list of messages are updated.
  """
  @spec apply_delta(t(), MessageDelta.t()) :: t()
  def apply_delta(%LLMChain{delta: nil} = chain, %MessageDelta{} = new_delta) do
    %LLMChain{chain | delta: new_delta}
  end

  def apply_delta(%LLMChain{delta: %MessageDelta{} = delta} = chain, %MessageDelta{} = new_delta) do
    merged = MessageDelta.merge_delta(delta, new_delta)

    # if the merged delta is now complete, updates as a message.
    if merged.status in [:complete, :length] do
      case MessageDelta.to_message(merged) do
        {:ok, %Message{} = message} ->
          fire_callback(chain, message)
          add_message(%LLMChain{chain | delta: nil}, message)

        {:error, reason} ->
          # should not have failed, but it did. Log the error and return
          # the chain unmodified.
          Logger.warning("Error applying delta message. Reason: #{inspect(reason)}")
          chain
      end
    else
      # the delta message is not yet complete. Update the delta with the merged
      # result.
      %LLMChain{chain | delta: merged}
    end
  end

  @doc """
  Apply a list of deltas to the chain.
  """
  @spec apply_deltas(t(), list()) :: t()
  def apply_deltas(%LLMChain{} = chain, deltas) when is_list(deltas) do
    deltas
    |> List.flatten()
    |> Enum.reduce(chain, fn d, acc -> apply_delta(acc, d) end)
  end

  @doc """
  Add a received Message struct to the chain. The LLMChain tracks the
  `last_message` received and the complete list of messages exchanged. Depending
  on the message role, the chain may be in a pending or incomplete state where
  a response from the LLM is anticipated.
  """
  @spec add_message(t(), Message.t()) :: t()
  def add_message(%LLMChain{} = chain, %Message{} = new_message) do
    needs_response =
      cond do
        new_message.role in [:user, :function_call, :function] -> true
        Message.is_function_call?(new_message) -> true
        new_message.role in [:system, :assistant] -> false
      end

    %LLMChain{
      chain
      | messages: chain.messages ++ [new_message],
        last_message: new_message,
        needs_response: needs_response
    }
  end

  @doc """
  Add a set of Message structs to the chain. This enables quickly building a chain
  for submitting to an LLM.
  """
  @spec add_messages(t(), [Message.t()]) :: t()
  def add_messages(%LLMChain{} = chain, messages) do
    Enum.reduce(messages, chain, fn msg, acc ->
      add_message(acc, msg)
    end)
  end

  @doc """
  Apply a set of PromptTemplates to the chain. The list of templates can also
  include Messages with no templates. Provide the inputs to apply to the
  templates for rendering as a message. The prepared messages are applied to the
  chain.
  """
  @spec apply_prompt_templates(t(), [Message.t() | PromptTemplate.t()], %{atom() => any()}) ::
          t() | no_return()
  def apply_prompt_templates(%LLMChain{} = chain, templates, %{} = inputs) do
    messages = PromptTemplate.to_messages!(templates, inputs)
    add_messages(chain, messages)
  end

  @doc """
  Convenience function for setting the prompt text for the LLMChain using
  prepared text.
  """
  @spec quick_prompt(t(), String.t()) :: t()
  def quick_prompt(%LLMChain{} = chain, text) do
    messages = [
      Message.new_system!(),
      Message.new_user!(text)
    ]

    add_messages(chain, messages)
  end

  @doc """
  If the `last_message` is a `%Message{role: :function_call}`, then the linked
  function is executed. If there is no `last_message` or the `last_message` is
  not a `:function_call`, the LLMChain is returned with no action performed.
  This makes it safe to call any time.

  The `context` is additional data that will be passed to the executed function.
  The value given here will override any `custom_context` set on the LLMChain.
  If not set, the global `custom_context` is used.

  https://platform.openai.com/docs/guides/gpt/function-calling
  """
  @spec execute_function(t(), context :: any()) :: t()
  def execute_function(chain, context \\ nil)
  def execute_function(%LLMChain{last_message: nil} = chain, _context), do: chain

  def execute_function(
        %LLMChain{last_message: %Message{} = message} = chain,
        context
      ) do
    if Message.is_function_call?(message) do
      # context to use
      use_context = context || chain.custom_context

      # find and execute the linked function
      case chain.function_map[message.function_name] do
        %Function{} = function ->
          if chain.verbose, do: IO.inspect(function.name, label: "EXECUTING FUNCTION")

          # execute the function
          result = Function.execute(function, message.arguments, use_context)
          if chain.verbose, do: IO.inspect(result, label: "FUNCTION RESULT")

          # add the :function response to the chain
          function_result = Message.new_function!(function.name, result)
          # fire the callback as this is newly generated message
          fire_callback(chain, function_result)
          LLMChain.add_message(chain, function_result)

        nil ->
          Logger.warning(
            "Received function_call for missing function #{inspect(message.function_name)}"
          )

          chain
      end
    else
      # Either not a function_call or an incomplete function_call, do nothing.
      chain
    end
  end

  # Fire the callback if set.
  defp fire_callback(%LLMChain{callback_fn: nil}, _data), do: :ok

  # OPTIONAL: Execute callback function
  defp fire_callback(%LLMChain{callback_fn: callback_fn}, data) when is_function(callback_fn) do
    case data do
      value when is_list(value) ->
        value
        |> List.flatten()
        |> Enum.each(fn item -> callback_fn.(item) end)

        :ok

      # not a list, pass the item as-is
      item ->
        callback_fn.(item)
        :ok
    end
  end

  @doc """
  Remove an incomplete MessageDelta from `delta` and add a Message with the
  desired status to the chain.
  """
  def cancel_delta(%LLMChain{delta: nil} = chain, _message_status), do: chain

  def cancel_delta(%LLMChain{delta: delta} = chain, message_status) do
    # remove the in-progress delta
    updated_chain = %LLMChain{chain | delta: nil}

    case MessageDelta.to_message(%MessageDelta{delta | status: :complete}) do
      {:ok, message} ->
        message = %Message{message | status: message_status}
        add_message(updated_chain, message)

      {:error, reason} ->
        Logger.error("Error attempting to cancel_delta. Reason: #{inspect(reason)}")
        chain
    end
  end
end