defmodule Oli.GenAI.Dialogue.Server do
  use GenServer
  require Logger

  @moduledoc """
  This module implements a dialogue server that manages a multi-turn
  conversation (or dialogue) between a user and a language model provider
  using the GenServer behavior.

  This dialogue is stateful and allows for asynchronous message handling,
  including function calls and streaming responses. It is designed to
  handle user messages, process them through the LLM, and manage the
  dialogue state entirely.

  The server is started with a static configuration that includes
  the service config to use, existing (i.e. starting) messages, functions available for calls,
  and a reply-to PID for sending responses back to the client.

  The server supports engaging in a dialogue by sending messages,
  processing responses, and automatically handling function calls. It can also
  handle streaming responses from the LLM, allowing for real-time
  interaction with the user.

  The dialogue server is designed to be resilient, with the ability to
  fallback to a backup model (if one exists) should the primary model fail.
  It maintains a state that includes the current messages, configuration,
  and any pending function calls.

  From a client, as an example of how to use this server:

  ```elixir
  {:ok, server} = Oli.GenAI.Dialogue.Server.new(configuration)

  Oli.GenAI.Dialogue.Server.engage(server, %Oli.GenAI.Completions.Message{
    role: :user,
    content: "What is the capital of France?"
  })
  ```

  Then, the server will stream back tokens as they are generated by the LLM,
  and you can handle them in your client (probably a LiveView) by listening to the
  server's messages, like so:

  ```elixir
  def handle_info({:dialogue_server, {:tokens_received, content}}, socket) do
    {:noreply, assign(socket, :active_message, socket.assigns.active_message <> content)}
  end

  def handle_info({:dialogue_server, {:tokens_finished}}, socket) do
    {:noreply, assign(socket, :streaming, false)}
  end
  ```

  The complete set of `:dialogue_server` messages that can be sent to the client are:
  - `{:dialogue_server, {:tokens_received, content}}`: Sent when new tokens are received from the LLM.
  - `{:dialogue_server, {:tokens_finished}}`: Sent when the LLM has finished sending tokens (i.e., the assisstant message is complete).
  - `{:dialogue_server, {:function_called, name, arguments}}`: Sent when a function has been requested by the LLM and called by the dialogue server.
  - `{:dialogue_server, {:error, reason}}`: Sent when an error occurs during processing.

  """

  alias Oli.GenAI.Completions
  alias Oli.GenAI.Completions.{Message, Function}
  alias Oli.GenAI.Dialogue.{State, Configuration}

  @doc """
  Starts a dialogue server given a static configuration.
  """
  def new(%Configuration{} = static_configuration) do
    GenServer.start_link(__MODULE__, static_configuration)
  end

  def engage(server, %Message{} = message) when is_pid(server) do
    GenServer.cast(server, {:engage, message})
  end

  def init(%Configuration{} = static_configuration) do
    Logger.debug(
      "Starting dialogue server for client process #{inspect(static_configuration.reply_to_pid)}"
    )

    {:ok, State.new(static_configuration)}
  end

  def handle_cast({:engage, message}, %State{} = state) do
    # Extract the current dialogue state and configuration
    %{messages: messages, configuration: cfg} = state

    # spawn a Task so we don’t block the GenServer loop
    server = self()

    Task.start(fn ->
      build_notify_fns(server, state)
      |> do_engage(state, message)
    end)

    Logger.debug(
      "Engaging dialogue for client #{inspect(cfg.reply_to_pid)} with message: #{inspect(message)}"
    )

    {:noreply, %State{state | messages: messages ++ [message]}}
  end

  defp build_notify_fns(server_pid, %State{
         configuration: %Configuration{reply_to_pid: reply_to_pid}
       }) do
    {
      fn message ->
        send(server_pid, message)
      end,
      fn message ->
        send(reply_to_pid, message)
      end
    }
  end

  def do_engage(
        {notify_server_fn, notify_client_fn},
        %State{configuration: configuration, registered_model: registered_model} = state,
        %Message{} = message
      ) do
    Logger.debug("do_engage for client #{inspect(configuration.reply_to_pid)}")

    process_chunk = fn chunk ->
      case chunk do
        {:error} ->
          notify_server_fn.({:stream_chunk, {:error}})

          notify_client_fn.(
            {:dialogue_server, {:error, "An error occurred while processing the request"}}
          )

        {:function_call, function_call} ->
          notify_server_fn.({:stream_chunk, {:function_call, function_call}})

        {:function_call_finished} ->
          notify_server_fn.({:stream_chunk, {:function_call_finished}})

        {:tokens_received, content} ->
          notify_server_fn.({:stream_chunk, {:tokens_received, content}})
          notify_client_fn.({:dialogue_server, {:tokens_received, content}})

        {:tokens_finished} ->
          notify_server_fn.({:stream_chunk, {:tokens_finished}})
          notify_client_fn.({:dialogue_server, {:tokens_finished}})
      end
    end

    response_handler_fn = fn response ->
      case response do
        chunks when is_list(chunks) ->
          Enum.each(chunks, fn chunk -> process_chunk.(chunk) end)

        chunk ->
          process_chunk.(chunk)
      end
    end

    case Completions.stream(
           state.messages ++ [message],
           configuration.functions,
           registered_model,
           response_handler_fn
         ) do
      :ok ->
        {:noreply, state}

      {:error, error} ->
        Logger.info(
          "Encountered completions http error for client #{inspect(configuration.reply_to_pid)}: #{inspect(error)}"
        )

        if should_retry?(state) do
          Logger.info("Falling back to backup for client #{inspect(configuration.reply_to_pid)}")

          state = fallback(state)
          do_engage({notify_server_fn, notify_client_fn}, state, message)
        else
          Logger.warning(
            "Failed without backup for client #{inspect(configuration.reply_to_pid)}"
          )

          notify_client_fn.(
            {:dialogue_server, {:error, "An error occurred while processing the request"}}
          )

          {:noreply, state}
        end
    end
  end

  defp should_retry?(state) do
    state.configuration.service_config.backup_model &&
      state.configuration.service_config.backup_model != state.registered_model
  end

  defp fallback(state),
    do: Map.put(state, :registered_model, state.configuration.service_config.backup_model)

  defp send_to_listener(%State{configuration: %Configuration{reply_to_pid: pid}}, message) do
    send(pid, message)
  end

  # All stream chunks come back here

  def handle_info({:stream_chunk, {:error}}, state) do
    if should_retry?(state) do
      Logger.info(
        "Falling back to backup for client #{inspect(state.configuration.reply_to_pid)}"
      )

      {:noreply, fallback(state)}
    else
      {:noreply, state}
    end
  end

  def handle_info({:stream_chunk, {:tokens_received, content}}, state) do
    # Append tokens to the “assistant” draft
    messages = append_token(state.messages, content)
    {:noreply, %{state | messages: messages}}
  end

  def handle_info({:stream_chunk, {:function_call, content}}, state) do
    # we are processing a chunk of information regarding a function to call
    case content do
      %{"name" => name, "arguments" => args, "id" => id} ->
        {:noreply, %{state | function_args: args, function_name: name, function_id: id}}

      %{"arguments" => args} ->
        {:noreply, %{state | function_args: state.function_args <> args}}
    end
  end

  def handle_info({:stream_chunk, {:function_call_finished}}, state) do
    # The LLM has finished spitting out the function args, we need to execute it.
    %{function_name: name, function_args: args, function_id: id, configuration: configuration} =
      state

    %{functions: functions} = configuration

    case Jason.decode(args) do
      {:ok, arguments_as_map} ->
        case Function.call(functions, name, arguments_as_map) do
          {:ok, result} ->
            message = %Message{
              role: :function,
              content: result,
              name: name,
              id: id,
              input: arguments_as_map
            }

            Logger.info("Function #{name} executed successfully with result: #{result}")

            # Let the client know that a function was called, it may want to persist this
            # or do something with it
            send_to_listener(
              state,
              {:dialogue_server, {:function_called, name, arguments_as_map}}
            )

            state = %{state | function_args: nil, function_name: nil, function_message: message}
            {:noreply, state, {:continue, :execute_function}}

          {:error, reason} ->
            Logger.error("Failed to execute function #{name}: #{reason}")
            {:noreply, state}
        end

      {:error, _} ->
        Logger.error("Failed to decode function arguments: #{args}")
        {:noreply, state}
    end
  end

  def handle_info({:stream_chunk, {:tokens_finished}}, state) do
    # done streaming tokens; nothing more to do here
    {:noreply, state}
  end

  # This is the entry point for engaging the dialogue server from within the server itself
  # (e.g., when the server sends a message to itself after executing a function call)
  def handle_info({:engage, message}, %State{} = state) do
    %{messages: messages} = state

    server = self()

    Task.start(fn ->
      build_notify_fns(server, state)
      |> do_engage(state, message)
    end)

    {:noreply, %State{state | messages: messages ++ [message]}}
  end

  def handle_continue(:execute_function, %{function_message: message} = state) do
    # We have executed the function successfully, now we want to
    # engage the dialogue with the result, by posting a message
    send(self(), {:engage, message})
    {:noreply, state}
  end

  defp append_token(messages, token) do
    case Enum.reverse(messages) do
      [%{role: :assistant, content: old} = last | rest] ->
        updated = %{last | content: old <> token}
        Enum.reverse([updated | rest])

      list ->
        # If the last message is not from the assistant, we need to create a new one
        new_message = Message.new(:assistant, token)
        Enum.reverse([new_message | list])
    end
  end
end
