sagents

defmodule Sagents.Persistence.StateSerializer do
  @moduledoc """
  Handles serialization and deserialization of AgentServer state.

  Automatically handles:
  - Version migrations
  - Non-serializable data (PIDs, refs)
  - Consistent use of string keys (for JSONB compatibility)

  ## String Keys

  NOTE: This serializer uses string keys throughout to match PostgreSQL JSONB
  behavior. When JSONB maps are restored from the database, they will have
  string keys, not atoms. For consistency, the serializer exports to string keys
  and expects string keys on import.

  ## Agent Configuration vs Conversation State

  StateSerializer follows the principle: **Code lives in code. Data lives in
  data.**

  **What is serialized**:
  - ✅ Conversation state: messages, todos, metadata

  **What is NOT serialized**:
  - ❌ Agent configuration: middleware, tools, model
  - ❌ Runtime identifiers: agent_id

  Agent capabilities (middleware, tools, model) are code-defined by your
  application. When restoring a conversation, create the agent from your
  application code and restore only the conversation state.

  See `AgentServer` module documentation for complete restoration examples.

  ## Versioning

  The serialized state includes a version field that allows for future
  migrations of the state format. The current version is 2.

  ### v1 → v2

  The `Sagents.Middleware.SubAgent` `task` tool's required argument was renamed
  from `subagent_type` to `task_name`. Stored tool-call arguments on `assistant`
  messages are rewritten on read so historical conversations remain consistent
  with the current tool schema. Affects calls named `"task"` or
  `"get_task_instructions"`.
  """

  alias Sagents.{State, Agent, Todo}
  alias LangChain.Message
  alias LangChain.Message.{ContentPart, ToolCall, ToolResult}

  @current_version 2

  @doc """
  Get the current serialization format version.
  """
  def current_version, do: @current_version

  @doc """
  Serializes AgentServer state to a map with string keys.

  Returns a map with string keys suitable for storage in JSONB columns.

  Note: The `agent_id` is NOT included in the serialized state. The agent_id
  is a runtime identifier used for process registration and PubSub topics, not
  part of the conversation state. When restoring state, you must provide the
  agent_id to `deserialize_server_state/2` or `AgentServer.start_link_from_state/2`.
  """
  def serialize_server_state(%Agent{} = _agent, %State{} = state) do
    %{
      # Version stays at 1, but now only contains state
      "version" => @current_version,
      "state" => serialize_state(state),
      "serialized_at" => DateTime.utc_now() |> DateTime.to_iso8601()
    }
  end

  @doc """
  Deserializes a serialized state map back to a State struct.

  The agent configuration is NOT deserialized - agents must be created
  from your application code. This function only restores conversation
  state (messages, metadata).

  ## Parameters

  - `agent_id` - The agent_id to use (NOT serialized, provided by you)
  - `data` - The serialized state map (from JSONB or export)
  - `opts` - Options (currently unused, kept for compatibility)

  ## Returns

  - `{:ok, state}` - Successfully deserialized state
  - `{:error, reason}` - Deserialization failed
  """
  def deserialize_server_state(agent_id, data, _opts \\ [])
      when is_binary(agent_id) and is_map(data)
      when is_map(data) and is_binary(agent_id) do
    # Handle version migration if needed
    migrated_data = maybe_migrate(data)
    # Only deserialize state, not agent config
    deserialize_state(agent_id, migrated_data["state"])
  end

  @doc """
  Serializes a State struct to a map with string keys.
  """
  def serialize_state(%State{} = state) do
    %{
      "messages" => Enum.map(state.messages, &serialize_message/1),
      "todos" => Enum.map(state.todos, &Todo.to_map/1),
      "metadata" => serialize_map_to_string_keys(state.metadata)
    }
  end

  @doc """
  Deserializes a map with string keys into a State struct.

  The `agent_id` must be provided as it is not part of the serialized data
  (it's a runtime identifier for process registration and PubSub).

  Returns `{:ok, state}` on success or `{:error, reason}` on failure.
  """
  def deserialize_state(agent_id, data) when is_binary(agent_id) and is_map(data) do
    messages =
      case data["messages"] do
        messages when is_list(messages) ->
          Enum.map(messages, &deserialize_message/1)

        _other ->
          []
      end

    todos =
      case data["todos"] do
        todos when is_list(todos) ->
          {:ok, parsed} = Todo.list_from_maps(todos)
          parsed

        _other ->
          []
      end

    metadata =
      case data["metadata"] do
        metadata when is_map(metadata) -> metadata
        _other -> %{}
      end

    case State.new(%{agent_id: agent_id, messages: messages, todos: todos, metadata: metadata}) do
      {:ok, state} -> {:ok, state}
      {:error, changeset} -> {:error, {:invalid_state, changeset}}
    end
  end

  # Private Functions

  defp serialize_message(%Message{} = message) do
    base = %{
      "role" => to_string(message.role),
      "content" => serialize_content(message.content),
      "status" => if(message.status, do: to_string(message.status), else: nil)
    }

    # Add optional fields if present
    base = maybe_add_string_field(base, "name", message.name)
    base = maybe_add_string_field(base, "index", message.index)

    # Add tool_calls if present
    base =
      if message.tool_calls && message.tool_calls != [] do
        Map.put(base, "tool_calls", Enum.map(message.tool_calls, &serialize_tool_call/1))
      else
        base
      end

    # Add tool_results if present
    base =
      if message.tool_results && message.tool_results != [] do
        Map.put(base, "tool_results", Enum.map(message.tool_results, &serialize_tool_result/1))
      else
        base
      end

    base
  end

  defp deserialize_message(data) when is_map(data) do
    # Convert string keys to atom keys for Message.new
    attrs = %{
      role: String.to_existing_atom(data["role"] || "user"),
      content: deserialize_content(data["content"]),
      status: if(data["status"], do: String.to_existing_atom(data["status"]), else: nil)
    }

    # Add optional fields
    attrs = maybe_add_field(attrs, :name, data["name"])
    attrs = maybe_add_field(attrs, :index, data["index"])

    # Add tool_calls if present
    attrs =
      if data["tool_calls"] do
        Map.put(attrs, :tool_calls, Enum.map(data["tool_calls"], &deserialize_tool_call/1))
      else
        attrs
      end

    # Add tool_results if present
    attrs =
      if data["tool_results"] do
        Map.put(attrs, :tool_results, Enum.map(data["tool_results"], &deserialize_tool_result/1))
      else
        attrs
      end

    # Use Message.new! to create the message
    case Message.new(attrs) do
      {:ok, message} -> message
      {:error, _changeset} -> raise "Failed to deserialize message: #{inspect(data)}"
    end
  end

  defp serialize_content(content) when is_binary(content), do: content

  defp serialize_content(content) when is_list(content) do
    Enum.map(content, &serialize_content_part/1)
  end

  defp serialize_content(nil), do: nil

  defp deserialize_content(content) when is_binary(content), do: content

  defp deserialize_content(content) when is_list(content) do
    Enum.map(content, &deserialize_content_part/1)
  end

  defp deserialize_content(nil), do: nil

  defp serialize_content_part(%ContentPart{} = part) do
    %{
      "type" => to_string(part.type),
      "content" => part.content,
      "options" => serialize_options(part.options)
    }
    |> Enum.reject(fn {_k, v} -> is_nil(v) or v == [] or v == %{} end)
    |> Map.new()
  end

  # Convert options keyword list to a map with string keys for JSON serialization
  defp serialize_options(opts) when is_list(opts) and opts != [] do
    # Check if it's a keyword list (list of 2-tuples with atom keys)
    if Keyword.keyword?(opts) do
      opts
      |> Enum.map(fn {k, v} -> {to_string(k), v} end)
      |> Map.new()
    else
      opts
    end
  end

  defp serialize_options(opts) when is_map(opts), do: opts
  defp serialize_options(_other), do: %{}

  defp deserialize_content_part(part) when is_map(part) do
    case ContentPart.new(%{
           type: String.to_existing_atom(part["type"] || "text"),
           content: part["content"],
           options: deserialize_options(part["options"])
         }) do
      {:ok, content_part} -> content_part
      {:error, _reason} -> raise "Failed to deserialize content part: #{inspect(part)}"
    end
  end

  # Convert options map with string keys back to keyword list with atom keys
  defp deserialize_options(opts) when is_map(opts) and map_size(opts) > 0 do
    opts
    |> Enum.map(fn {k, v} -> {String.to_atom(k), v} end)
    |> Keyword.new()
  end

  defp deserialize_options(opts) when is_list(opts), do: opts
  defp deserialize_options(_other), do: []

  defp serialize_tool_call(%ToolCall{} = tool_call) do
    # Arguments can be either a map (when complete) or a string (during streaming)
    arguments =
      case tool_call.arguments do
        nil -> %{}
        arg when is_map(arg) -> serialize_map_to_string_keys(arg)
        arg when is_binary(arg) -> arg
        arg -> inspect(arg)
      end

    base = %{
      "call_id" => tool_call.call_id,
      "type" => to_string(tool_call.type),
      "name" => tool_call.name,
      "arguments" => arguments,
      "status" => to_string(tool_call.status)
    }

    # Add display_text if present (direct field on ToolCall)
    base =
      if tool_call.display_text do
        Map.put(base, "display_text", tool_call.display_text)
      else
        base
      end

    # Add metadata if present (for model-specific attributes)
    base =
      if tool_call.metadata && map_size(tool_call.metadata) > 0 do
        Map.put(base, "metadata", serialize_map_to_string_keys(tool_call.metadata))
      else
        base
      end

    base
    |> Enum.reject(fn {_k, v} -> is_nil(v) end)
    |> Map.new()
  end

  defp deserialize_tool_call(data) when is_map(data) do
    attrs = %{
      call_id: data["call_id"],
      type: String.to_existing_atom(data["type"] || "function"),
      name: data["name"],
      arguments: data["arguments"] || %{},
      status: String.to_existing_atom(data["status"] || "incomplete")
    }

    # Add display_text if present (direct field on ToolCall)
    attrs =
      if data["display_text"] do
        Map.put(attrs, :display_text, data["display_text"])
      else
        attrs
      end

    # Add metadata if present (for model-specific attributes)
    attrs =
      if data["metadata"] do
        Map.put(attrs, :metadata, data["metadata"])
      else
        attrs
      end

    case ToolCall.new(attrs) do
      {:ok, tool_call} -> tool_call
      {:error, _reason} -> raise "Failed to deserialize tool call: #{inspect(data)}"
    end
  end

  defp serialize_tool_result(%ToolResult{} = tool_result) do
    base = %{
      "type" => to_string(tool_result.type),
      "tool_call_id" => tool_result.tool_call_id,
      "name" => tool_result.name,
      "content" => serialize_content(tool_result.content),
      "is_error" => tool_result.is_error,
      "is_interrupt" => tool_result.is_interrupt
    }

    # `interrupt_data` is a virtual Ecto field, so it would otherwise be lost
    # on round-trip. Encode it opaquely as Base64-wrapped term_to_binary —
    # JSONB just sees a string, and the load path can decide whether to
    # restore or demote based on the agent's middleware.
    base =
      if tool_result.is_interrupt && not is_nil(tool_result.interrupt_data) do
        Map.put(base, "interrupt_data_b64", encode_interrupt_data(tool_result.interrupt_data))
      else
        base
      end

    base
    |> Enum.reject(fn {_k, v} -> is_nil(v) end)
    |> Map.new()
  end

  defp deserialize_tool_result(data) when is_map(data) do
    interrupt_data = decode_interrupt_data(data["interrupt_data_b64"])

    case ToolResult.new(%{
           type: String.to_existing_atom(data["type"] || "function"),
           tool_call_id: data["tool_call_id"],
           name: data["name"],
           content: deserialize_content(data["content"]),
           is_error: data["is_error"] || false,
           is_interrupt: data["is_interrupt"] || false,
           interrupt_data: interrupt_data
         }) do
      {:ok, tool_result} -> tool_result
      {:error, _reason} -> raise "Failed to deserialize tool result: #{inspect(data)}"
    end
  end

  defp encode_interrupt_data(data) do
    data |> :erlang.term_to_binary() |> Base.encode64()
  end

  # Best-effort decode. Any failure mode (missing key, bad Base64, malformed
  # binary, unknown atom under [:safe]) collapses to `nil`, which lets
  # `clean_stale_interrupts/2` demote the result to an error.
  defp decode_interrupt_data(nil), do: nil

  defp decode_interrupt_data(b64) when is_binary(b64) do
    with {:ok, bin} <- Base.decode64(b64),
         {:ok, term} <- safe_binary_to_term(bin) do
      term
    else
      _other -> nil
    end
  end

  defp decode_interrupt_data(_data), do: nil

  defp safe_binary_to_term(bin) when is_binary(bin) do
    {:ok, :erlang.binary_to_term(bin, [:safe])}
  rescue
    _other -> :error
  catch
    _kind, _other -> :error
  end

  defp serialize_map_to_string_keys(map) when is_map(map) do
    Map.new(map, fn
      {k, v} when is_atom(k) -> {to_string(k), serialize_value(v)}
      {k, v} when is_binary(k) -> {k, serialize_value(v)}
    end)
  end

  defp serialize_value(value) when is_struct(value) do
    # For structs, we'll just use their inspect representation as a string
    # This is not ideal for deserialization, but complex structs in middleware opts
    # are not typically needed for state restoration
    inspect(value)
  end

  defp serialize_value(value) when is_map(value), do: serialize_map_to_string_keys(value)
  defp serialize_value(value) when is_list(value), do: Enum.map(value, &serialize_value/1)
  defp serialize_value(value) when is_atom(value) and not is_nil(value), do: to_string(value)

  # Functions cannot be serialized to JSON - skip them
  # They will need to be recreated during deserialization from config
  defp serialize_value(value) when is_function(value), do: nil

  defp serialize_value(value), do: value

  defp maybe_add_string_field(map, _key, nil), do: map
  defp maybe_add_string_field(map, key, value), do: Map.put(map, key, value)

  defp maybe_add_field(map, _key, nil), do: map
  defp maybe_add_field(map, key, value), do: Map.put(map, key, value)

  defp maybe_migrate(%{"version" => 2} = data), do: data

  defp maybe_migrate(%{"version" => 1} = data) do
    data
    |> migrate_v1_to_v2()
    |> Map.put("version", 2)
  end

  defp maybe_migrate(%{"version" => _other} = data) do
    # Future versions go here. Unknown versions pass through
    data
  end

  defp maybe_migrate(data) do
    # No version field — treat as the oldest supported version and run
    # the full migration chain.
    data |> Map.put("version", 1) |> maybe_migrate()
  end

  # v1 → v2: SubAgent middleware renamed the `task` / `get_task_instructions`
  # tool argument from "subagent_type" to "task_name". Rewrite stored tool
  # calls so replayed conversations match the current tool schema.
  defp migrate_v1_to_v2(%{"state" => %{"messages" => messages} = state} = data)
       when is_list(messages) do
    migrated = Enum.map(messages, &migrate_v1_subagent_tool_call/1)
    %{data | "state" => %{state | "messages" => migrated}}
  end

  defp migrate_v1_to_v2(data), do: data

  defp migrate_v1_subagent_tool_call(%{"tool_calls" => calls} = msg) when is_list(calls) do
    %{msg | "tool_calls" => Enum.map(calls, &rename_subagent_type_arg/1)}
  end

  defp migrate_v1_subagent_tool_call(msg), do: msg

  defp rename_subagent_type_arg(%{"name" => name, "arguments" => args} = tc)
       when name in ["task", "get_task_instructions"] and is_map(args) do
    case Map.pop(args, "subagent_type") do
      {nil, _rest} -> tc
      {value, rest} -> %{tc | "arguments" => Map.put(rest, "task_name", value)}
    end
  end

  defp rename_subagent_type_arg(tc), do: tc
end