llm

defmodule LLM.Adapter.Anthropic do
  @moduledoc """
  Anthropic Messages API adapter.

  Handles the wire format for the `/v1/messages` endpoint.

  ## Capabilities

  - Streaming via SSE with content block events
  - Tool/use calling with `input_json_delta` accumulation
  - Extended thinking with `thinking` content blocks and signatures
  - Model listing (returns hardcoded model list)

  ## Request format

  Messages use Anthropic's content block format:
  - Text: `%{"type" => "text", "text" => "..."}`
  - Tool use: `%{"type" => "tool_use", "id" => ..., "name" => ..., "input" => ...}`
  - Thinking: `%{"type" => "thinking", "thinking" => "...", "signature" => "..."}`

  System prompts are sent as a top-level `"system"` field.
  Developer messages are merged into the system prompt.

  ## SSE events

  - `event: content_block_start` — new content block (tool use or thinking)
  - `event: content_block_delta` — incremental text, thinking, or JSON
  - `event: content_block_stop` — block complete (tool calls emitted here)
  - `event: message_delta` — stop reason and usage
  - `event: message_stop` — end of stream
  """

  @behaviour LLM.Adapter

  @impl true
  def init_stream_state, do: %{blocks: %{}, usage: nil}

  @impl true
  @structured_output_tool "__structured_output__"

  def build_request(context, opts) do
    {developer_content, messages} = extract_developer_messages(context.messages)
    system_prompt = merge_system_prompts(context.system, developer_content)
    {tools, tool_choice} = structured_output_tools(context.tools, opts[:structured_output])

    %{
      "model" => Keyword.fetch!(opts, :model),
      "max_tokens" => opts[:max_tokens] || 4096,
      "messages" => Enum.map(messages, &encode_message/1)
    }
    |> maybe_put("system", system_prompt)
    |> maybe_put("temperature", opts[:temperature])
    |> maybe_put("tools", encode_tools(tools))
    |> maybe_put("tool_choice", tool_choice)
    |> maybe_put("thinking", normalize_thinking(opts[:thinking]))
  end

  @impl true
  def decode_response(%{"content" => content, "stop_reason" => stop_reason} = raw) do
    usage =
      case raw do
        %{"usage" => u} ->
          LLM.Usage.new(input_tokens: u["input_tokens"], output_tokens: u["output_tokens"])

        _ ->
          nil
      end

    {:ok,
     %LLM.Response{
       message: build_response_message(content),
       usage: usage,
       stop_reason: String.to_atom(stop_reason),
       raw: raw
     }}
  end

  def decode_response(other), do: {:error, {:invalid_response, other}}

  @impl true
  def stream_path, do: "/v1/messages"

  @impl true
  def stream_headers(_opts) do
    [
      {"content-type", "application/json"},
      {"anthropic-version", "2023-06-01"},
      {"anthropic-beta", "messages-2023-12-15"}
    ]
  end

  @impl true
  def auth_headers(%{api_key: api_key}, _opts), do: [{"x-api-key", api_key}]

  @impl true
  def decode_chunk("event: message_stop" <> _, state), do: {:done, state}

  def decode_chunk("event: message_start\n" <> data, state) do
    case decode_data(data) do
      {:ok, %{"usage" => usage}} ->
        start_usage =
          LLM.Usage.new(
            input_tokens: usage["input_tokens"] || 0,
            cache_read_tokens: usage["cache_read_input_tokens"] || 0,
            cache_write_tokens: usage["cache_creation_input_tokens"] || 0
          )

        {[], %{state | usage: start_usage}}

      _ ->
        {[], state}
    end
  end

  def decode_chunk("event: content_block_start\n" <> data, state) do
    case decode_data(data) do
      {:ok, %{"index" => index, "content_block" => %{"type" => "tool_use"} = block}} ->
        block_state = %{type: "tool_use", id: block["id"], name: block["name"], input_json: ""}
        {[], put_in(state, [:blocks, index], block_state)}

      {:ok, %{"index" => index, "content_block" => %{"type" => "thinking"} = block}} ->
        {[],
         put_in(state, [:blocks, index], %{
           type: "thinking",
           text: block["thinking"] || "",
           signature: block["signature"]
         })}

      _ ->
        {[], state}
    end
  end

  def decode_chunk("event: content_block_delta\n" <> data, state) do
    case decode_data(data) do
      {:ok, %{"index" => _index, "delta" => %{"type" => "text_delta", "text" => text}}} ->
        {[%LLM.Stream.Chunk{text: text}], state}

      {:ok, %{"index" => index, "delta" => %{"type" => "thinking_delta", "thinking" => text}}} ->
        next_state =
          update_in(state, [:blocks, index, :text], fn existing -> (existing || "") <> text end)

        {[%LLM.Stream.Thinking{text: text}], next_state}

      {:ok,
       %{"index" => index, "delta" => %{"type" => "input_json_delta", "partial_json" => json}}} ->
        next_state =
          update_in(state, [:blocks, index, :input_json], fn existing ->
            (existing || "") <> json
          end)

        {[], next_state}

      _ ->
        {[], state}
    end
  end

  def decode_chunk("event: content_block_stop\n" <> data, state) do
    case decode_data(data) do
      {:ok, %{"index" => index}} ->
        emit_completed_block(index, state)

      _ ->
        {[], state}
    end
  end

  def decode_chunk("event: message_delta\n" <> data, state) do
    case decode_data(data) do
      {:ok, %{"delta" => %{"stop_reason" => reason}, "usage" => delta_usage}} ->
        output_tokens = delta_usage["output_tokens"] || 0
        cache_read_delta = delta_usage["cache_read_input_tokens"] || 0

        start_usage = state.usage || %LLM.Usage{}

        combined_usage = %{
          start_usage
          | output_tokens: start_usage.output_tokens + output_tokens,
            cache_read_tokens: start_usage.cache_read_tokens + cache_read_delta
        }

        combined_usage = %{
          combined_usage
          | total_tokens:
              combined_usage.input_tokens + combined_usage.output_tokens +
                combined_usage.cache_read_tokens + combined_usage.cache_write_tokens
        }

        {[
           %LLM.Stream.Stop{
             reason: String.to_atom(reason),
             usage: combined_usage
           }
         ], %{state | usage: nil}}

      {:ok, %{"delta" => %{"stop_reason" => reason}}} ->
        {[%LLM.Stream.Stop{reason: String.to_atom(reason)}], state}

      _ ->
        {[], state}
    end
  end

  def decode_chunk("event: ping\n" <> _, state), do: {[], state}

  def decode_chunk("event: error\n" <> data, state) do
    case Jason.decode(data) do
      {:ok, %{"error" => %{"message" => msg}}} -> {[%LLM.Stream.Error{message: msg}], state}
      _ -> {[], state}
    end
  end

  def decode_chunk(_, state), do: {[], state}

  @impl true
  def normalize_thinking(%{"type" => "enabled", "budget_tokens" => _} = t), do: t

  def normalize_thinking(%{type: "enabled", budget_tokens: _} = t),
    do: %{"type" => t.type, "budget_tokens" => t.budget_tokens}

  def normalize_thinking(:low), do: %{"type" => "enabled", "budget_tokens" => 2_000}
  def normalize_thinking(:medium), do: %{"type" => "enabled", "budget_tokens" => 10_000}
  def normalize_thinking(:high), do: %{"type" => "enabled", "budget_tokens" => 32_000}
  def normalize_thinking(:xhigh), do: %{"type" => "enabled", "budget_tokens" => 64_000}
  def normalize_thinking(:max), do: %{"type" => "enabled", "budget_tokens" => 100_000}
  def normalize_thinking(nil), do: nil
  def normalize_thinking(false), do: nil
  def normalize_thinking(val), do: val

  defp encode_message(%LLM.Message{role: :tool, content: content, tool_call_id: tool_call_id}) do
    %{
      "role" => "user",
      "content" => [
        %{"type" => "tool_result", "tool_use_id" => tool_call_id, "content" => to_string(content)}
      ]
    }
  end

  defp encode_message(%LLM.Message{role: :assistant} = msg) do
    thinking_blocks =
      case msg.thinking do
        nil ->
          []

        text when is_binary(text) ->
          [%{"type" => "thinking", "thinking" => text}]

        %{text: text, signature: sig} ->
          [%{"type" => "thinking", "thinking" => text, "signature" => sig}]
      end

    text_blocks =
      case msg.content do
        nil -> []
        "" -> []
        text -> [%{"type" => "text", "text" => text}]
      end

    tool_blocks =
      (msg.tools || [])
      |> Enum.map(fn %{id: id, name: name, args: args} ->
        %{"type" => "tool_use", "id" => id, "name" => name, "input" => args}
      end)

    %{"role" => "assistant", "content" => thinking_blocks ++ text_blocks ++ tool_blocks}
  end

  defp encode_message(%LLM.Message{role: role, content: content}) when is_binary(content) do
    %{"role" => to_string(role), "content" => content}
  end

  defp encode_message(%LLM.Message{role: :user, content: content}) when is_list(content) do
    blocks =
      Enum.map(content, fn
        {:text, text} ->
          %{"type" => "text", "text" => text}

        {:tool_result, id, _name, result} ->
          %{"type" => "tool_result", "tool_use_id" => id, "content" => to_string(result)}
      end)

    %{"role" => "user", "content" => blocks}
  end

  defp structured_output_tools(tools, nil), do: {tools, nil}

  defp structured_output_tools(tools, spec) do
    {_name, schema} = extract_schema(spec)

    so_tool = %LLM.Tool{
      name: @structured_output_tool,
      description: "Output the result as structured data.",
      input_schema: schema
    }

    {tools ++ [so_tool], %{"type" => "tool", "name" => @structured_output_tool}}
  end

  defp extract_schema(%{name: name, schema: schema}), do: {to_string(name), schema}
  defp extract_schema(%{"name" => name, "schema" => schema}), do: {name, schema}
  defp extract_schema(schema), do: {"output", schema}

  defp encode_tools([]), do: nil

  defp encode_tools(tools) do
    Enum.map(tools, fn
      %LLM.Tool{} = tool ->
        %{
          "name" => tool.name,
          "description" => tool.description,
          "input_schema" => tool.input_schema
        }

      module when is_atom(module) ->
        module |> LLM.Tool.from_module() |> then(&List.first(encode_tools([&1])))
    end)
  end

  defp build_response_message(content) do
    parts =
      Enum.flat_map(content, fn
        %{"type" => "thinking", "thinking" => text, "signature" => signature}
        when is_binary(text) and text != "" ->
          [{:thinking, %{text: text, signature: signature}}]

        %{"type" => "thinking", "thinking" => text} when is_binary(text) and text != "" ->
          [{:thinking, text}]

        %{"type" => "text", "text" => text} when is_binary(text) and text != "" ->
          [{:text, text}]

        %{"type" => "tool_use", "id" => id, "name" => name, "input" => input} ->
          [{:tool_use, id, name, input}]

        _ ->
          []
      end)

    thinking =
      Enum.find_value(parts, fn
        {:thinking, v} -> {:found, v}
        _ -> nil
      end)
      |> then(fn
        nil -> nil
        {:found, v} -> v
      end)

    text =
      parts
      |> Enum.flat_map(fn
        {:text, t} -> [t]
        _ -> []
      end)
      |> IO.iodata_to_binary()

    tools =
      parts
      |> Enum.flat_map(fn
        {:tool_use, id, name, args} -> [%{id: id, name: name, args: args}]
        _ -> []
      end)

    %LLM.Message{
      role: :assistant,
      content: if(text == "", do: nil, else: text),
      thinking: thinking,
      tools: if(tools == [], do: nil, else: tools)
    }
  end

  defp emit_completed_block(index, state) do
    case get_in(state, [:blocks, index]) do
      %{type: "tool_use", id: id, name: name, input_json: input_json} ->
        arguments =
          case Jason.decode(input_json) do
            {:ok, parsed} -> parsed
            _ -> %{}
          end

        chunk = %LLM.Stream.ToolCall{
          id: id,
          name: name,
          arguments: arguments,
          index: index,
          complete: true
        }

        {[chunk], update_in(state, [:blocks], &Map.delete(&1, index))}

      _ ->
        {[], state}
    end
  end

  defp decode_data(data) do
    case Regex.run(~r/^data: (.+)$/s, data, capture: :all_but_first) do
      [json] -> Jason.decode(json)
      _ -> {:error, :not_found}
    end
  end

  defp extract_developer_messages(messages) do
    {developer_msgs, other_msgs} = Enum.split_with(messages, &(&1.role == :developer))
    developer_content = developer_msgs |> Enum.map(& &1.content) |> Enum.join("\n\n")
    {developer_content, other_msgs}
  end

  defp merge_system_prompts(nil, developer_content) when developer_content == "", do: nil
  defp merge_system_prompts(nil, developer_content), do: developer_content
  defp merge_system_prompts(system, ""), do: system
  defp merge_system_prompts(system, developer_content), do: "#{system}\n\n#{developer_content}"

  defp maybe_put(map, _key, nil), do: map
  defp maybe_put(map, _key, []), do: map
  defp maybe_put(map, key, value), do: Map.put(map, key, value)

  @impl true
  def list_models(_config) do
    models = [
      %{
        id: "claude-sonnet-4-5-20250514",
        name: "Claude Sonnet 4.5",
        description: "Most intelligent model with extended thinking",
        context_window: 200_000,
        max_output: 128_000,
        capabilities: [:text, :vision, :tools, :thinking]
      },
      %{
        id: "claude-sonnet-4-20250514",
        name: "Claude Sonnet 4",
        description: "Balanced model for most tasks",
        context_window: 200_000,
        max_output: 64_000,
        capabilities: [:text, :vision, :tools]
      },
      %{
        id: "claude-3-5-haiku-20241022",
        name: "Claude 3.5 Haiku",
        description: "Fast and efficient model",
        context_window: 200_000,
        max_output: 8_192,
        capabilities: [:text, :vision, :tools]
      },
      %{
        id: "claude-3-opus-20240229",
        name: "Claude 3 Opus",
        description: "Most powerful model for complex tasks",
        context_window: 200_000,
        max_output: 4_096,
        capabilities: [:text, :vision, :tools]
      }
    ]

    {:ok, models}
  end
end