mcp_kit

defmodule MCPKit.Plug do
  @moduledoc """
  Plug entrypoint for the MCP Streamable HTTP transport.

  The plug is configured by `MCPKit.Router` and expects:
  - a host definition module implementing `MCPKit.Definition`
  - a persistent session store returned by that definition
  - a compile-time registry of declared tools for the mounted scope

  Current runtime scope is intentionally narrow:
  - `initialize`
  - `GET` session stream over SSE
  - `notifications/initialized`
  - `ping`
  - `completion/complete`
  - `prompts/list`
  - `prompts/get`
  - `resources/list`
  - `resources/read`
  - `resources/templates/list`
  - `tools/list`
  - `tools/call`
  - `notifications/cancelled`
  - `DELETE` session termination

  The plug now also uses request-time policy checks for tools, prompts, and
  resources via the host's configured `MCPKit.Policy` module.
  """
  @behaviour Plug

  require Logger

  import Plug.Conn

  alias MCPKit.Policy.AllowAll
  alias MCPKit.Protocol
  alias MCPKit.Response
  alias MCPKit.Runtime
  alias Plug.Conn.Unfetched

  @session_header "mcp-session-id"
  @protocol_header "mcp-protocol-version"
  @origin_header "origin"

  def init(opts) do
    definition = Keyword.fetch!(opts, :definition)
    prompts = Keyword.get(opts, :prompts, [])
    resources = Keyword.get(opts, :resources, [])
    tools = Keyword.get(opts, :tools, [])
    runtime = Keyword.get_lazy(opts, :runtime, fn -> Runtime.default_name(definition) end)
    policy = policy_module(definition)

    %{
      definition: definition,
      policy: policy,
      prompts: prompts,
      resources: resources,
      runtime: runtime,
      session_store: definition.session_store(),
      tools: tools
    }
  end

  def call(conn, config) do
    with :ok <- validate_origin(conn) do
      dispatch(conn, config)
    else
      {:error, :invalid_origin} ->
        log_transport_warning(conn, config, :invalid_origin)
        send_resp(conn, 403, "Forbidden")
    end
  end

  defp dispatch(conn, config) do
    case conn.method do
      "POST" -> handle_post(conn, config)
      "GET" -> handle_get(conn, config)
      "DELETE" -> handle_delete(conn, config)
      _ -> send_resp(conn, 405, "")
    end
  end

  defp handle_post(conn, config) do
    with :ok <- validate_accept_header(conn),
         {:ok, message} <- read_message(conn) do
      log_message_debug(conn, config, message)
      process_message(conn, message, config)
    else
      {:error, :invalid_accept_header} ->
        log_transport_warning(conn, config, :invalid_accept_header)
        send_resp(conn, 406, "")

      {:error, :parse_error} ->
        log_transport_warning(conn, config, :parse_error)

        conn
        |> put_resp_content_type("application/json")
        |> send_resp(400, Protocol.encode!(Protocol.parse_error()))

      {:error, :invalid_request} ->
        log_transport_warning(conn, config, :invalid_request)

        conn
        |> put_resp_content_type("application/json")
        |> send_resp(400, Protocol.encode!(Protocol.invalid_request()))

      {:error, reason} ->
        log_transport_warning(conn, config, reason)
        send_resp(conn, 400, "")
    end
  end

  defp handle_get(conn, config) do
    log_get_debug(conn, config)

    with :ok <- validate_sse_accept_header(conn),
         {:ok, session} <- fetch_mcp_session(conn, config),
         {:ok, generation} <-
           Runtime.attach_stream(config.runtime, Map.fetch!(session, :id), self()),
         {:ok, chunked_conn} <- open_sse_stream(conn) do
      Logger.debug(
        "MCP SSE stream opened session_id=#{Map.fetch!(session, :id)} generation=#{generation} runtime=#{inspect(config.runtime)}"
      )

      sse_loop(chunked_conn, config, Map.fetch!(session, :id), generation)
    else
      {:error, :invalid_accept_header} ->
        log_transport_warning(conn, config, :invalid_accept_header)
        send_resp(conn, 406, "")

      {:error, :session_required} ->
        log_transport_warning(conn, config, :session_required)
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        log_transport_warning(conn, config, :not_found)
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        log_transport_warning(conn, config, :unsupported_protocol_version)
        send_resp(conn, 400, "")

      {:error, :runtime_unavailable} ->
        log_transport_error(conn, config, :runtime_unavailable)
        send_resp(conn, 500, "")

      {:error, reason} ->
        log_transport_error(conn, config, reason)
        send_resp(conn, 500, "")
    end
  end

  defp read_message(%{body_params: %Unfetched{aspect: :body_params}} = conn) do
    with {:ok, raw_body, _conn} <- read_body(conn),
         {:ok, message} <- Protocol.decode(raw_body) do
      {:ok, message}
    end
  end

  defp read_message(%{body_params: body}) when is_map(body), do: {:ok, body}
  defp read_message(_conn), do: {:error, :invalid_request}

  defp process_message(conn, message, config) do
    cond do
      Protocol.request?(message) ->
        handle_request(conn, message, config)

      Protocol.notification?(message) ->
        handle_notification(conn, message, config)

      Protocol.response?(message) ->
        handle_response(conn, message, config)

      true ->
        conn
        |> put_resp_content_type("application/json")
        |> send_resp(400, Protocol.encode!(Protocol.invalid_request()))
    end
  end

  defp handle_request(conn, %{"method" => "initialize", "params" => params, "id" => id}, config) do
    with :ok <- reject_existing_session_header(conn),
         {:ok, protocol_version, requested_protocol_version, client_info, client_capabilities} <-
           validate_initialize_params(params, config.definition),
         {:ok, session} <-
           config.session_store.create_session(%{
             protocol_version: protocol_version,
             client_info: client_info,
             client_capabilities: client_capabilities
           }) do
      result = %{
        "protocolVersion" => protocol_version,
        "capabilities" => capabilities(config),
        "serverInfo" => config.definition.server_info()
      }

      result = maybe_put_instructions(result, config.definition)

      Logger.debug(
        "MCP initialize succeeded request_id=#{inspect(id)} requested_protocol_version=#{requested_protocol_version} negotiated_protocol_version=#{protocol_version} session_id=#{Map.fetch!(session, :id)} client_info=#{inspect(client_info)}"
      )

      conn
      |> put_resp_header(@session_header, Map.fetch!(session, :id))
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, result)))
    else
      {:error, :session_header_not_allowed} ->
        send_jsonrpc_error(
          conn,
          400,
          Protocol.invalid_request(id, "Initialize requests must not include a session ID")
        )

      {:error, message} when is_binary(message) ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_params(id, message))

      {:error, _reason} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, "Failed to create session"))
    end
  end

  defp handle_request(conn, %{"method" => "ping", "id" => id} = message, config) do
    with {:ok, _session_or_nil} <- fetch_optional_session(conn, message, config) do
      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, %{})))
    else
      {:error, :session_required} -> send_resp(conn, 400, "")
      {:error, :not_found} -> send_resp(conn, 404, "")
      {:error, :unsupported_protocol_version} -> send_resp(conn, 400, "")
    end
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "completion/complete", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config),
         {:ok, values} <- complete_argument(conn, session, params, id, config) do
      conn
      |> put_resp_content_type("application/json")
      |> send_resp(
        200,
        Protocol.encode!(
          Protocol.success(id, %{
            "completion" => %{"values" => values, "total" => length(values), "hasMore" => false}
          })
        )
      )
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))

      {:error, {:invalid_params, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_params(id, message))

      {:error, {:forbidden, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.forbidden(id, message))

      {:error, {:internal_error, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, message))
    end
  end

  defp handle_request(conn, %{"method" => "completion/complete"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "resources/list", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config) do
      resources = list_resources(config, conn, session, params, id)

      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, %{"resources" => resources})))
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))
    end
  end

  defp handle_request(conn, %{"method" => "resources/list"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "resources/templates/list", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config) do
      templates = list_resource_templates(config, conn, session, params, id)

      conn
      |> put_resp_content_type("application/json")
      |> send_resp(
        200,
        Protocol.encode!(Protocol.success(id, %{"resourceTemplates" => templates}))
      )
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))
    end
  end

  defp handle_request(conn, %{"method" => "resources/templates/list"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "resources/read", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config),
         {:ok, resource_module, uri} <- fetch_resource_and_uri(conn, session, params, id, config),
         {:reply, contents, _context} when is_list(contents) <-
           resource_module.read(uri, %{session: session}) do
      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, %{"contents" => contents})))
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))

      {:error, {:invalid_params, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_params(id, message))

      {:error, {:forbidden, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.forbidden(id, message))

      {:error, {:internal_error, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, message))

      {:reply, _contents, _context} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, "Resource read failed"))
    end
  end

  defp handle_request(conn, %{"method" => "resources/read"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "prompts/list", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config) do
      prompts = list_prompts(config, conn, session, params, id)

      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, %{"prompts" => prompts})))
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))
    end
  end

  defp handle_request(conn, %{"method" => "prompts/list"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "prompts/get", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config),
         {:ok, prompt, arguments} <- fetch_prompt_and_arguments(conn, session, params, id, config),
         {:ok, validated_arguments} <- prompt.validate_arguments(arguments),
         {:reply, messages, _context} when is_list(messages) <-
           prompt.render(validated_arguments, %{session: session}) do
      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, %{"messages" => messages})))
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))

      {:error, {:invalid_params, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_params(id, message))

      {:error, {:forbidden, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.forbidden(id, message))

      {:error, {:internal_error, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, message))

      {:reply, _messages, _context} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, "Prompt rendering failed"))
    end
  end

  defp handle_request(conn, %{"method" => "prompts/get"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "tools/list", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config),
         :ok <- validate_tools_list_params(params) do
      tools = list_tools(config, conn, session, params, id)

      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, %{"tools" => tools})))
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))

      {:error, message} when is_binary(message) ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_params(id, message))
    end
  end

  defp handle_request(conn, %{"method" => "tools/list"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(
         conn,
         %{"id" => id, "method" => "tools/call", "params" => params} = message,
         config
       )
       when is_map(params) do
    with {:ok, session} <- fetch_initialized_session(conn, message, config),
         {:ok, tool, arguments} <- fetch_tool_and_arguments(conn, session, params, id, config),
         {:ok, validated_arguments} <- tool.validate_arguments(arguments),
         {:reply, response, _context} <- tool.execute(validated_arguments, %{session: session}) do
      conn
      |> put_resp_content_type("application/json")
      |> send_resp(200, Protocol.encode!(Protocol.success(id, Response.to_result(response))))
    else
      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")

      {:error, :not_initialized} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_request(id, "Server not initialized"))

      {:error, {:invalid_params, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.invalid_params(id, message))

      {:error, {:forbidden, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.forbidden(id, message))

      {:error, {:internal_error, message}} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, message))

      {:error, _reason} ->
        send_jsonrpc_error(conn, 200, Protocol.internal_error(id, "Tool execution failed"))
    end
  end

  defp handle_request(conn, %{"method" => "tools/call"} = message, config) do
    handle_request(conn, Map.put(message, "params", %{}), config)
  end

  defp handle_request(conn, %{"id" => id, "method" => method} = message, config) do
    with {:ok, _session} <- fetch_optional_session(conn, message, config) do
      send_jsonrpc_error(conn, 200, Protocol.method_not_found(id, "Unknown method: #{method}"))
    else
      {:error, :session_required} -> send_resp(conn, 400, "")
      {:error, :not_found} -> send_resp(conn, 404, "")
      {:error, :unsupported_protocol_version} -> send_resp(conn, 400, "")
    end
  end

  defp handle_notification(conn, %{"method" => "notifications/initialized"}, config) do
    case fetch_mcp_session(conn, config) do
      {:ok, session} ->
        case config.session_store.mark_initialized(session) do
          {:ok, _session} ->
            Logger.debug(
              "MCP session marked initialized session_id=#{Map.fetch!(session, :id)} protocol_version=#{Map.fetch!(session, :protocol_version)}"
            )

            send_resp(conn, 202, "")

          {:error, _reason} ->
            send_resp(conn, 500, "")
        end

      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")
    end
  end

  defp handle_notification(conn, %{"method" => "notifications/cancelled"} = message, config) do
    case fetch_initialized_session(conn, message, config) do
      {:ok, _session} -> send_resp(conn, 202, "")
      {:error, :session_required} -> send_resp(conn, 400, "")
      {:error, :not_found} -> send_resp(conn, 404, "")
      {:error, :unsupported_protocol_version} -> send_resp(conn, 400, "")
      {:error, :not_initialized} -> send_resp(conn, 202, "")
    end
  end

  defp handle_notification(conn, message, config) do
    case fetch_optional_session(conn, message, config) do
      {:ok, _session_or_nil} -> send_resp(conn, 202, "")
      {:error, :session_required} -> send_resp(conn, 400, "")
      {:error, :not_found} -> send_resp(conn, 404, "")
      {:error, :unsupported_protocol_version} -> send_resp(conn, 400, "")
    end
  end

  defp handle_response(conn, message, config) do
    case fetch_optional_session(conn, message, config) do
      {:ok, session} ->
        Runtime.handle_response(config.runtime, Map.fetch!(session, :id), message)
        send_resp(conn, 202, "")

      {:error, :session_required} ->
        send_resp(conn, 400, "")

      {:error, :not_found} ->
        send_resp(conn, 404, "")

      {:error, :unsupported_protocol_version} ->
        send_resp(conn, 400, "")
    end
  end

  defp handle_delete(conn, config) do
    case get_req_header(conn, @session_header) do
      [session_id] when is_binary(session_id) and session_id != "" ->
        Runtime.terminate_session(config.runtime, session_id)

        case config.session_store.delete_session(session_id) do
          {:ok, _session} -> send_resp(conn, 200, "")
          {:error, :not_found} -> send_resp(conn, 404, "")
          {:error, _reason} -> send_resp(conn, 500, "")
        end

      _ ->
        send_resp(conn, 400, "")
    end
  end

  defp capabilities(config) do
    %{}
    |> maybe_put_capability("completions", config.prompts != [] or config.resources != [])
    |> maybe_put_capability("prompts", config.prompts != [])
    |> maybe_put_capability("resources", config.resources != [])
    |> maybe_put_capability("tools", config.tools != [])
  end

  defp list_tools(config, conn, session, params, request_id) do
    Enum.flat_map(config.tools, fn %{name: name, module: tool} ->
      action = {:tools, :list, name}
      context = policy_context(conn, session, params, %{}, request_id, config)

      case authorize(config, action, context) do
        :allow ->
          [
            %{
              "name" => name,
              "description" => tool.description(),
              "inputSchema" => tool.input_schema()
            }
            |> Map.reject(fn {_key, value} -> is_nil(value) end)
          ]

        {:deny, _reason} ->
          []
      end
    end)
  end

  defp list_prompts(config, conn, session, params, request_id) do
    Enum.flat_map(config.prompts, fn %{name: name, module: prompt} ->
      action = {:prompts, :list, name}
      context = policy_context(conn, session, params, %{}, request_id, config)

      case authorize(config, action, context) do
        :allow ->
          [
            %{
              "name" => name,
              "description" => prompt.description(),
              "arguments" => prompt.arguments()
            }
            |> Map.reject(fn {_key, value} -> is_nil(value) end)
          ]

        {:deny, _reason} ->
          []
      end
    end)
  end

  defp list_resources(config, conn, session, params, request_id) do
    Enum.flat_map(config.resources, fn %{name: name, module: resource} ->
      action = {:resources, :list, name}
      context = policy_context(conn, session, params, %{}, request_id, config)

      case authorize(config, action, context) do
        :allow ->
          resource.resources()

        {:deny, _reason} ->
          []
      end
    end)
  end

  defp list_resource_templates(config, conn, session, params, request_id) do
    Enum.flat_map(config.resources, fn %{name: name, module: resource} ->
      action = {:resources, :templates_list, name}
      context = policy_context(conn, session, params, %{}, request_id, config)

      case authorize(config, action, context) do
        :allow ->
          resource.templates()

        {:deny, _reason} ->
          []
      end
    end)
  end

  defp fetch_tool_and_arguments(conn, session, %{"name" => name} = params, request_id, config)
       when is_binary(name) do
    arguments = Map.get(params, "arguments", %{})

    cond do
      not is_map(arguments) ->
        {:error, {:invalid_params, "Tool arguments must be an object"}}

      tool = Enum.find(config.tools, &(&1.name == name)) ->
        action = {:tools, :call, name}
        context = policy_context(conn, session, params, arguments, request_id, config)

        case authorize(config, action, context) do
          :allow -> {:ok, tool.module, arguments}
          {:deny, :not_found} -> {:error, {:invalid_params, "Unknown tool: #{name}"}}
          {:deny, :forbidden} -> {:error, {:forbidden, "Forbidden"}}
        end

      true ->
        {:error, {:invalid_params, "Unknown tool: #{name}"}}
    end
  end

  defp fetch_tool_and_arguments(_conn, _session, _params, _request_id, _config) do
    {:error, {:invalid_params, "Tool name is required"}}
  end

  defp fetch_prompt_and_arguments(conn, session, %{"name" => name} = params, request_id, config)
       when is_binary(name) do
    arguments = Map.get(params, "arguments", %{})

    cond do
      not is_map(arguments) ->
        {:error, {:invalid_params, "Prompt arguments must be an object"}}

      prompt = Enum.find(config.prompts, &(&1.name == name)) ->
        action = {:prompts, :get, name}
        context = policy_context(conn, session, params, arguments, request_id, config)

        case authorize(config, action, context) do
          :allow -> {:ok, prompt.module, arguments}
          {:deny, :not_found} -> {:error, {:invalid_params, "Unknown prompt: #{name}"}}
          {:deny, :forbidden} -> {:error, {:forbidden, "Forbidden"}}
        end

      true ->
        {:error, {:invalid_params, "Unknown prompt: #{name}"}}
    end
  end

  defp fetch_prompt_and_arguments(_conn, _session, _params, _request_id, _config) do
    {:error, {:invalid_params, "Prompt name is required"}}
  end

  defp fetch_resource_and_uri(conn, session, %{"uri" => uri} = params, request_id, config)
       when is_binary(uri) do
    case Enum.find(config.resources, fn %{module: resource} ->
           resource_handles_uri?(resource, uri)
         end) do
      nil ->
        {:error, {:invalid_params, "Unknown resource: #{uri}"}}

      %{module: resource} ->
        action = {:resources, :read, uri}
        context = policy_context(conn, session, params, %{}, request_id, config)

        case authorize(config, action, context) do
          :allow -> {:ok, resource, uri}
          {:deny, :not_found} -> {:error, {:invalid_params, "Unknown resource: #{uri}"}}
          {:deny, :forbidden} -> {:error, {:forbidden, "Forbidden"}}
        end
    end
  end

  defp fetch_resource_and_uri(_conn, _session, _params, _request_id, _config) do
    {:error, {:invalid_params, "Resource uri is required"}}
  end

  defp complete_argument(conn, session, params, request_id, config) do
    with {:ok, ref, argument_name, argument_value, arguments} <-
           validate_completion_params(params) do
      case ref do
        %{"type" => "ref/prompt", "name" => name} ->
          complete_prompt_argument(
            conn,
            session,
            name,
            argument_name,
            argument_value,
            arguments,
            request_id,
            config
          )

        %{"type" => "ref/resource", "uri" => uri} ->
          complete_resource_argument(
            conn,
            session,
            uri,
            argument_name,
            argument_value,
            arguments,
            request_id,
            config
          )
      end
    end
  end

  defp complete_prompt_argument(
         conn,
         session,
         prompt_name,
         argument_name,
         argument_value,
         arguments,
         request_id,
         config
       ) do
    case Enum.find(config.prompts, &(&1.name == prompt_name)) do
      nil ->
        {:error, {:invalid_params, "Unknown prompt: #{prompt_name}"}}

      %{module: prompt} ->
        context =
          policy_context(conn, session, %{"name" => prompt_name}, arguments, request_id, config)

        case authorize(
               config,
               {:completion, {:prompt_argument, prompt_name, argument_name}},
               context
             ) do
          :allow ->
            invoke_completion(
              prompt,
              argument_name,
              Map.put(arguments, argument_name, argument_value),
              %{session: session}
            )

          {:deny, :not_found} ->
            {:error, {:invalid_params, "Unknown prompt: #{prompt_name}"}}

          {:deny, :forbidden} ->
            {:error, {:forbidden, "Forbidden"}}
        end
    end
  end

  defp complete_resource_argument(
         conn,
         session,
         uri,
         argument_name,
         argument_value,
         arguments,
         request_id,
         config
       ) do
    case Enum.find(config.resources, fn %{module: resource} ->
           resource_template?(resource, uri)
         end) do
      nil ->
        {:error, {:invalid_params, "Unknown resource template: #{uri}"}}

      %{name: name, module: resource} ->
        context = policy_context(conn, session, %{"uri" => uri}, arguments, request_id, config)

        case authorize(config, {:completion, {:resource_template, name, argument_name}}, context) do
          :allow ->
            invoke_completion(
              resource,
              argument_name,
              Map.put(arguments, argument_name, argument_value),
              %{session: session}
            )

          {:deny, :not_found} ->
            {:error, {:invalid_params, "Unknown resource template: #{uri}"}}

          {:deny, :forbidden} ->
            {:error, {:forbidden, "Forbidden"}}
        end
    end
  end

  defp invoke_completion(module, argument_name, arguments, context) do
    case module.complete(argument_name, arguments, context) do
      {:ok, values} when is_list(values) -> {:ok, Enum.take(values, 100)}
      :unsupported -> {:ok, []}
      _other -> {:error, {:internal_error, "Completion failed"}}
    end
  end

  defp validate_completion_params(
         %{
           "ref" => ref,
           "argument" => %{"name" => argument_name, "value" => argument_value}
         } = params
       )
       when is_map(ref) and is_binary(argument_name) and is_binary(argument_value) do
    arguments =
      case get_in(params, ["context", "arguments"]) do
        nil -> %{}
        args when is_map(args) -> args
        _other -> :invalid
      end

    case {ref, arguments} do
      {%{"type" => "ref/prompt", "name" => name}, arguments}
      when is_binary(name) and is_map(arguments) ->
        {:ok, ref, argument_name, argument_value, arguments}

      {%{"type" => "ref/resource", "uri" => uri}, arguments}
      when is_binary(uri) and is_map(arguments) ->
        {:ok, ref, argument_name, argument_value, arguments}

      {_ref, :invalid} ->
        {:error, {:invalid_params, "Completion context arguments must be an object"}}

      _other ->
        {:error, {:invalid_params, "Completion ref must target a prompt name or resource uri"}}
    end
  end

  defp validate_completion_params(_params) do
    {:error, {:invalid_params, "Completion params must include ref and argument"}}
  end

  defp fetch_initialized_session(conn, _message, config) do
    with {:ok, session} <- fetch_mcp_session(conn, config),
         true <- Map.fetch!(session, :initialized) || {:error, :not_initialized} do
      {:ok, session}
    end
  end

  defp fetch_optional_session(_conn, %{"method" => "initialize"}, _config), do: {:ok, nil}

  defp fetch_optional_session(conn, _message, config) do
    case get_req_header(conn, @session_header) do
      [] -> {:error, :session_required}
      [_session_id] -> fetch_mcp_session(conn, config)
    end
  end

  defp fetch_mcp_session(conn, config) do
    case get_req_header(conn, @session_header) do
      [session_id] when is_binary(session_id) and session_id != "" ->
        Logger.debug(
          "MCP session lookup start session_id=#{session_id} protocol_header=#{inspect(protocol_header(conn))} store=#{inspect(config.session_store)}"
        )

        with {:ok, session} <- config.session_store.fetch_session(session_id),
             :ok <- validate_protocol_version(conn, session),
             {:ok, session} <- config.session_store.touch_session(session) do
          Logger.debug(
            "MCP session lookup succeeded session_id=#{session_id} protocol_version=#{Map.fetch!(session, :protocol_version)} initialized=#{Map.fetch!(session, :initialized)}"
          )

          {:ok, session}
        else
          {:error, :not_found} = error ->
            Logger.debug("MCP session lookup failed session_id=#{session_id} reason=:not_found")
            error

          {:error, :unsupported_protocol_version} = error ->
            Logger.debug(
              "MCP session lookup failed session_id=#{session_id} reason=:unsupported_protocol_version protocol_header=#{inspect(protocol_header(conn))}"
            )

            error

          {:error, reason} ->
            Logger.debug(
              "MCP session lookup failed session_id=#{session_id} reason=#{inspect(reason)} mapped_reason=:not_found"
            )

            {:error, :not_found}
        end

      _ ->
        Logger.debug(
          "MCP session lookup failed session_id=nil reason=:session_required protocol_header=#{inspect(protocol_header(conn))}"
        )

        {:error, :session_required}
    end
  end

  defp validate_protocol_version(conn, session) do
    session_protocol_version = Map.fetch!(session, :protocol_version)

    case get_req_header(conn, @protocol_header) do
      [] ->
        Logger.debug(
          "MCP protocol header missing session_id=#{Map.fetch!(session, :id)} assumed_protocol_version=#{session_protocol_version}"
        )

        :ok

      [version] when version == session_protocol_version ->
        :ok

      [_version] ->
        {:error, :unsupported_protocol_version}
    end
  end

  defp validate_initialize_params(
         %{
           "protocolVersion" => protocol_version,
           "capabilities" => capabilities,
           "clientInfo" => client_info
         },
         _definition
       )
       when is_binary(protocol_version) and is_map(capabilities) and is_map(client_info) do
    negotiated_protocol_version = MCPKit.Definition.protocol_version()

    {:ok, negotiated_protocol_version, protocol_version, client_info, capabilities}
  end

  defp validate_initialize_params(_params, _definition) do
    {:error, "Initialize params must include protocolVersion, capabilities, and clientInfo"}
  end

  defp validate_tools_list_params(%{} = params) do
    case Map.get(params, "cursor") do
      nil -> :ok
      _cursor -> {:error, "Tool pagination cursors are not supported"}
    end
  end

  defp validate_accept_header(conn) do
    accept_header = conn |> get_req_header("accept") |> List.first("")

    if accept_header == "" or String.contains?(accept_header, "application/json") do
      :ok
    else
      {:error, :invalid_accept_header}
    end
  end

  defp validate_sse_accept_header(conn) do
    accept_header = conn |> get_req_header("accept") |> List.first("")

    if accept_header == "" or String.contains?(accept_header, "text/event-stream") do
      :ok
    else
      {:error, :invalid_accept_header}
    end
  end

  defp reject_existing_session_header(conn) do
    if get_req_header(conn, @session_header) == [] do
      :ok
    else
      {:error, :session_header_not_allowed}
    end
  end

  defp validate_origin(conn) do
    case get_req_header(conn, "origin") do
      [] -> :ok
      [origin] -> if allowed_origin?(conn, origin), do: :ok, else: {:error, :invalid_origin}
      _ -> {:error, :invalid_origin}
    end
  end

  defp allowed_origin?(conn, origin) do
    uri = URI.parse(origin)

    uri.scheme in ["http", "https"] and uri.host == conn.host and
      normalize_port(uri.port, uri.scheme) == conn.port
  end

  defp normalize_port(nil, "http"), do: 80
  defp normalize_port(nil, "https"), do: 443
  defp normalize_port(port, _scheme), do: port

  defp log_get_debug(conn, config) do
    Logger.debug(
      "MCP transport GET received path=#{conn.request_path} runtime=#{inspect(config.runtime)} accept=#{inspect(List.first(get_req_header(conn, "accept")))} session_id=#{inspect(session_header(conn))} protocol_header=#{inspect(protocol_header(conn))} origin=#{inspect(List.first(get_req_header(conn, @origin_header)))}"
    )
  end

  defp log_message_debug(conn, config, message) do
    Logger.debug(
      "MCP transport message received method=#{conn.method} path=#{conn.request_path} runtime=#{inspect(config.runtime)} rpc_method=#{inspect(Map.get(message, "method"))} rpc_id=#{inspect(Map.get(message, "id"))} message_type=#{message_type(message)} accept=#{inspect(List.first(get_req_header(conn, "accept")))} session_id=#{inspect(session_header(conn))} protocol_header=#{inspect(protocol_header(conn))} origin=#{inspect(List.first(get_req_header(conn, @origin_header)))}"
    )
  end

  defp log_transport_warning(conn, config, reason) do
    Logger.warning(
      log_message(conn, config, reason),
      log_metadata(conn, config, reason)
    )
  end

  defp log_transport_error(conn, config, reason) do
    Logger.error(
      log_message(conn, config, reason),
      log_metadata(conn, config, reason)
    )
  end

  defp log_message(conn, config, reason) do
    "MCP transport request failed reason=#{inspect(reason)} method=#{conn.method} path=#{conn.request_path} runtime=#{inspect(config.runtime)}"
  end

  defp log_metadata(conn, config, reason) do
    [
      reason: inspect(reason),
      http_method: conn.method,
      request_path: conn.request_path,
      runtime: inspect(config.runtime),
      session_id: session_header(conn),
      protocol_version: protocol_header(conn)
    ]
  end

  defp message_type(message) do
    cond do
      Protocol.request?(message) -> "request"
      Protocol.notification?(message) -> "notification"
      Protocol.response?(message) -> "response"
      true -> "unknown"
    end
  end

  defp session_header(conn), do: List.first(get_req_header(conn, @session_header))

  defp protocol_header(conn), do: List.first(get_req_header(conn, @protocol_header))

  defp maybe_put_instructions(result, definition) do
    if function_exported?(definition, :instructions, 0) do
      Map.put(result, "instructions", definition.instructions())
    else
      result
    end
  end

  defp open_sse_stream(conn) do
    {:ok,
     conn
     |> put_resp_header("cache-control", "no-cache")
     |> put_resp_header("x-accel-buffering", "no")
     |> put_resp_content_type("text/event-stream")
     |> send_chunked(200)}
  end

  defp sse_loop(conn, config, session_id, generation) do
    receive do
      {:mcp_runtime_sse, ^generation, payload} ->
        Logger.debug(
          "MCP SSE payload sent session_id=#{session_id} generation=#{generation} rpc_method=#{inspect(Map.get(payload, "method"))} rpc_id=#{inspect(Map.get(payload, "id"))}"
        )

        case chunk(conn, "data: #{Protocol.encode!(payload)}\n\n") do
          {:ok, conn} -> sse_loop(conn, config, session_id, generation)
          {:error, _reason} -> finalize_stream(conn, config, session_id, generation)
        end

      {:mcp_runtime_close, _reason} ->
        finalize_stream(conn, config, session_id, generation)

      _message ->
        sse_loop(conn, config, session_id, generation)
    after
      15_000 ->
        case chunk(conn, ": keep-alive\n\n") do
          {:ok, conn} -> sse_loop(conn, config, session_id, generation)
          {:error, _reason} -> finalize_stream(conn, config, session_id, generation)
        end
    end
  end

  defp finalize_stream(conn, config, session_id, generation) do
    Logger.debug(
      "MCP SSE stream finalized session_id=#{session_id} generation=#{generation} runtime=#{inspect(config.runtime)}"
    )

    Runtime.stream_detached(config.runtime, session_id, generation, self())
    conn
  end

  defp policy_context(conn, session, params, arguments, request_id, config) do
    %{
      arguments: arguments,
      client_capabilities: Map.fetch!(session, :client_capabilities),
      client_info: Map.fetch!(session, :client_info),
      conn: conn,
      definition: config.definition,
      params: params,
      request_id: request_id,
      session: session
    }
  end

  defp authorize(config, action, context) do
    config.policy.authorize(action, context)
  end

  defp resource_handles_uri?(resource, uri) do
    Enum.any?(resource.resources(), &(Map.get(&1, "uri") == uri)) or
      Enum.any?(resource.templates(), &uri_matches_template?(uri, Map.get(&1, "uriTemplate")))
  end

  defp resource_template?(resource, uri) do
    Enum.any?(resource.templates(), &(Map.get(&1, "uriTemplate") == uri))
  end

  defp uri_matches_template?(_uri, nil), do: false

  defp uri_matches_template?(uri, template) when is_binary(uri) and is_binary(template) do
    regex =
      template
      |> Regex.escape()
      |> String.replace(~r/\\\{[^}]+\\\}/, "[^/]+")
      |> then(&Regex.compile!("^" <> &1 <> "$"))

    Regex.match?(regex, uri)
  end

  defp maybe_put_capability(capabilities, _name, false), do: capabilities
  defp maybe_put_capability(capabilities, name, true), do: Map.put(capabilities, name, %{})

  defp policy_module(definition) do
    if function_exported?(definition, :policy, 0) do
      definition.policy()
    else
      AllowAll
    end
  end

  defp send_jsonrpc_error(conn, status, payload) do
    conn
    |> put_resp_content_type("application/json")
    |> send_resp(status, Protocol.encode!(payload))
  end
end