defmodule AI.Agent.Memory.AssociativeLearning do
  @moduledoc """
  Agent that scores memories for a conversation.

  Given a conversation and a list of `AI.Memory` structs, it asks the LLM to
  assign each memory a relevance score from 1–10 and returns a map of
  `memory_id => score` keyed by each memory's `id`.

  Response format:

      {:ok, conversation_id, %{"memory_id_1" => score_1, ...}}

  On transient decode/validation failures, the agent will retry up to
  `@retry_limit` times before returning an error.
  """

  @behaviour AI.Agent

  @model AI.Model.fast()
  @retry_limit 3

  alias Store.Project.Conversation
  alias AI.Memory

  @type conversation_id :: String.t()
  @type memory_score_map :: %{optional(String.t()) => pos_integer()}

  @doc """
  Entry point required by the `AI.Agent` behaviour.

  Expected args map:

      %{
        agent: %AI.Agent{},
        conversation: %Store.Project.Conversation{} | %{id: id, messages: messages},
        memories: [%AI.Memory{}, ...]
      }

  Returns `{:ok, conversation_id, scores}` on success, where `scores` is a map
  of memory IDs (as strings) to integer scores in the range 1–10.
  """
  @spec get_response(map()) :: {:ok, conversation_id, memory_score_map} | {:error, term()}
  @impl AI.Agent
  def get_response(%{agent: agent, conversation: conversation, memories: memories})
      when is_list(memories) do
    with {:ok, convo_id} <- conversation_id(conversation) do
      do_score_with_retries(agent, convo_id, conversation, memories, @retry_limit)
    end
  end

  def get_response(_), do: {:error, :invalid_arguments}

  # ---------------------------------------------------------------------------
  # Core scoring flow with simple retry
  # ---------------------------------------------------------------------------

  defp do_score_with_retries(_agent, convo_id, _conversation, _memories, 0) do
    {:error, {:max_retries_exceeded, convo_id}}
  end

  defp do_score_with_retries(agent, convo_id, conversation, memories, attempts_left) do
    case score_once(agent, conversation, memories) do
      {:ok, scores} ->
        {:ok, convo_id, scores}

      {:error, _reason} ->
        do_score_with_retries(agent, convo_id, conversation, memories, attempts_left - 1)
    end
  end

  defp score_once(agent, conversation, memories) do
    messages = build_messages(conversation, memories)

    case AI.Agent.get_completion(agent,
           model: @model,
           messages: messages,
           response_format: %{
             type: "json_schema",
             json_schema: response_schema()
           }
         ) do
      {:ok, %AI.Completion{response: response}} ->
        decode_scores(response)

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

  # ---------------------------------------------------------------------------
  # Prompt construction
  # ---------------------------------------------------------------------------

  @system_prompt """
  You are an associative memory selection and scoring helper inside a larger
  agent system.

  Your job is to read the current conversation and a list of candidate
  memories, then assign each memory a *relevance score* from 1 (barely
  relevant) to 10 (extremely central) **for this specific conversation**.

  Important rules:
  - You MUST return a JSON object where each key is a memory id (string) and
    each value is an integer score from 1 to 10.
  - Every provided memory id MUST appear in the object, even if you think it
    is only weakly relevant.
  - Use the full range 1–10 when appropriate; do not cluster all scores at one
    value.
  - Base your judgment only on the provided conversation messages and memory
    descriptions.
  - If a memory seems unrelated to the conversation, give it a low score like
    1 or 2, but still include it in the result.
  """

  defp build_messages(conversation, memories) do
    convo_block = format_conversation(conversation)
    memories_block = format_memories(memories)

    user_content = """
    Here is the current conversation:

    #{convo_block}

    -----

    Here are the candidate memories:

    #{memories_block}

    Return ONLY a JSON object mapping memory ids (as strings) to integer
    relevance scores from 1 to 10.
    """

    [
      AI.Util.system_msg(@system_prompt),
      AI.Util.user_msg(user_content)
    ]
  end

  defp format_conversation(%Conversation{} = convo) do
    {:ok, _ts, messages, _metadata} = Conversation.read(convo)
    format_messages(messages)
  end

  defp format_conversation(%{messages: messages}) when is_list(messages) do
    format_messages(messages)
  end

  defp format_conversation(_), do: "(no conversation messages available)"

  defp format_messages(messages) do
    messages
    |> Enum.with_index(1)
    |> Enum.map(fn {msg, idx} ->
      role = Map.get(msg, :role) || Map.get(msg, "role") || "unknown"
      content = Map.get(msg, :content) || Map.get(msg, "content") || ""
      "[#{idx}] #{role}: #{content}"
    end)
    |> Enum.join("\n")
  end

  defp format_memories(memories) do
    memories
    |> Enum.map(fn %Memory{id: id, label: label, scope: scope, response_template: template} = mem ->
      scope_str = to_string(scope)
      # Fallbacks in case some fields are nil
      label = label || id || "(no label)"
      template = template || "(no response template)"

      pattern_info =
        case Map.get(mem, :pattern_tokens) do
          %{} = tokens when map_size(tokens) > 0 ->
            sampled =
              tokens
              |> Enum.take(5)
              |> Enum.map_join(", ", fn {tok, weight} -> "#{tok}:#{weight}" end)

            "pattern_tokens: #{sampled}"

          _ ->
            "pattern_tokens: (none)"
        end

      "- id: #{id}\n  scope: #{scope_str}\n  label: #{label}\n  template: #{template}\n  #{pattern_info}"
    end)
    |> Enum.join("\n\n")
  end

  # ---------------------------------------------------------------------------
  # Response format and decoding
  # ---------------------------------------------------------------------------

  defp response_schema do
    %{
      name: "memory_relevance_scores",
      strict: true,
      schema: %{
        type: "object",
        description:
          "Map of memory ids to relevance scores (integer 1–10) for the current conversation.",
        additionalProperties: %{
          type: "integer",
          minimum: 1,
          maximum: 10
        }
      }
    }
  end

  defp decode_scores(json) when is_binary(json) do
    case Jason.decode(json) do
      {:ok, %{} = scores} -> validate_scores(scores)
      {:ok, _other} -> {:error, :invalid_response_shape}
      {:error, reason} -> {:error, {:decode_error, reason}}
    end
  end

  defp decode_scores(_), do: {:error, :invalid_response}

  defp validate_scores(scores) when is_map(scores) do
    with true <-
           Enum.all?(scores, fn {k, v} ->
             is_binary(k) and is_integer(v) and v >= 1 and v <= 10
           end) do
      {:ok, scores}
    else
      _ -> {:error, :invalid_scores}
    end
  end

  # ---------------------------------------------------------------------------
  # Helpers
  # ---------------------------------------------------------------------------

  defp conversation_id(%Conversation{} = convo) do
    {:ok, convo.id}
  end

  defp conversation_id(%{id: id}) when is_binary(id), do: {:ok, id}
  defp conversation_id(%{id: id}) when is_integer(id), do: {:ok, Integer.to_string(id)}
  defp conversation_id(_), do: {:error, :missing_conversation_id}
end