fnord

defmodule MCP.OAuth2.Client do
  @moduledoc """
  Pure OAuth2 + PKCE client implementation for MCP servers.

  Unlike OIDC libraries (like oidcc), this works with OAuth2 Authorization Server
  discovery (RFC 8414) at `/.well-known/oauth-authorization-server`, not just
  OpenID Connect discovery at `/.well-known/openid-configuration`.

  Implements:
  - Authorization Code flow with PKCE (RFC 7636)
  - Token refresh (RFC 6749)
  - OAuth2 server metadata discovery (RFC 8414)
  - Resource Indicators (RFC 8707) via the optional `:resource` config key

  The MCP authorization spec requires clients to send the canonical MCP
  server URI as the `resource` parameter on both the authorization and
  token requests so the server can bind the grant to that resource.
  Servers that advertise a `resource` in their metadata (e.g. Linear)
  hard-fail flows that omit it. Callers pass the server's `base_url` as
  `:resource`; when absent the parameter is omitted for compatibility
  with servers that predate RFC 8707.

  Security:
  - PKCE is always required (S256 challenge method)
  - Tokens are never logged
  - Uses secure random generation for state and verifier
  """

  @type config :: %{
          required(:discovery_url) => String.t(),
          required(:client_id) => String.t(),
          optional(:client_secret) => String.t(),
          required(:redirect_uri) => String.t(),
          required(:scopes) => [String.t()],
          optional(:resource) => String.t()
        }

  @type tokens :: %{
          access_token: String.t(),
          token_type: String.t(),
          expires_at: non_neg_integer(),
          refresh_token: String.t() | nil,
          scope: String.t() | nil
        }

  @doc """
  Start OAuth2 authorization flow with PKCE.

  Fetches server metadata, generates PKCE parameters, and builds authorization URL.

  Returns: `{:ok, %{auth_url: String.t(), state: String.t(), code_verifier: String.t()}}`
  """
  @spec start_flow(config) ::
          {:ok, %{auth_url: String.t(), state: String.t(), code_verifier: String.t()}}
          | {:error, term()}
  def start_flow(cfg) do
    with {:ok, metadata} <- fetch_metadata(cfg.discovery_url),
         {:ok, state} <- generate_state(),
         {:ok, verifier, challenge} <- generate_pkce(),
         {:ok, auth_url} <- build_authorization_url(metadata, cfg, state, challenge) do
      {:ok, %{auth_url: auth_url, state: state, code_verifier: verifier}}
    end
  end

  @doc """
  Handle OAuth2 callback and exchange authorization code for tokens.

  Validates state, extracts code, exchanges for tokens with PKCE verifier.

  Returns: `{:ok, tokens}` with normalized token map
  """
  @spec handle_callback(config, map(), String.t(), String.t()) ::
          {:ok, tokens} | {:error, term()}
  def handle_callback(cfg, params, expected_state, code_verifier) do
    with {:ok, metadata} <- fetch_metadata(cfg.discovery_url),
         :ok <- verify_state(params, expected_state),
         {:ok, code} <- extract_code(params),
         {:ok, tokens} <- exchange_code(metadata, cfg, code, code_verifier) do
      {:ok, normalize_tokens(tokens)}
    end
  end

  @doc """
  Refresh an expired access token using the refresh token.

  Returns: `{:ok, tokens}` with new access token and possibly new refresh token
  """
  @spec refresh_token(config, String.t()) :: {:ok, tokens} | {:error, term()}
  def refresh_token(cfg, refresh_token) do
    with {:ok, metadata} <- fetch_metadata(cfg.discovery_url),
         {:ok, tokens} <- refresh_with_server(metadata, cfg, refresh_token) do
      {:ok, normalize_tokens(tokens)}
    end
  end

  # -- Metadata Discovery --

  defp fetch_metadata(discovery_url) do
    case Http.Client.impl().get(discovery_url, [], recv_timeout: 10_000, timeout: 10_000) do
      {:ok, %{status_code: 200, body: body}} ->
        case SafeJson.decode(body) do
          {:ok, metadata} when is_map(metadata) ->
            validate_metadata(metadata)

          {:ok, _} ->
            {:error, {:invalid_metadata, "Discovery response is not a JSON object"}}

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

      {:ok, %{status_code: code}} ->
        {:error, {:http_error, code}}

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

  defp validate_metadata(metadata) do
    required = ["authorization_endpoint", "token_endpoint"]

    case Enum.filter(required, &(!Map.has_key?(metadata, &1))) do
      [] -> {:ok, metadata}
      missing -> {:error, {:incomplete_metadata, "Missing: #{Enum.join(missing, ", ")}"}}
    end
  end

  # -- PKCE Generation --

  defp generate_pkce do
    # Generate 32-byte random verifier, base64url encode
    verifier =
      :crypto.strong_rand_bytes(32)
      |> Base.url_encode64(padding: false)

    # SHA256 hash the verifier, base64url encode for challenge
    challenge =
      :crypto.hash(:sha256, verifier)
      |> Base.url_encode64(padding: false)

    {:ok, verifier, challenge}
  end

  defp generate_state do
    state =
      :crypto.strong_rand_bytes(16)
      |> Base.url_encode64(padding: false)

    {:ok, state}
  end

  # -- Authorization URL Building --

  defp build_authorization_url(metadata, cfg, state, challenge) do
    auth_endpoint = metadata["authorization_endpoint"]

    params =
      %{
        "response_type" => "code",
        "client_id" => cfg.client_id,
        "redirect_uri" => cfg.redirect_uri,
        "scope" => Enum.join(cfg.scopes, " "),
        "state" => state,
        "code_challenge" => challenge,
        "code_challenge_method" => "S256"
      }
      |> Map.merge(resource_param(cfg))

    query = URI.encode_query(params)
    {:ok, "#{auth_endpoint}?#{query}"}
  end

  # RFC 8707 resource indicator. Present on the authorization request, the
  # code exchange, AND the refresh request - the MCP spec requires it on all
  # three so each issued token stays bound to the same resource. Omitted
  # entirely when the caller didn't configure one (pre-RFC 8707 servers).
  defp resource_param(cfg) do
    case Map.get(cfg, :resource) do
      resource when is_binary(resource) and resource != "" -> %{"resource" => resource}
      _ -> %{}
    end
  end

  # -- State Verification --

  defp verify_state(params, expected_state) do
    case Map.get(params, "state") do
      ^expected_state -> :ok
      nil -> {:error, :missing_state}
      _ -> {:error, :state_mismatch}
    end
  end

  defp extract_code(params) do
    case Map.get(params, "code") do
      nil -> {:error, :missing_code}
      code when is_binary(code) -> {:ok, code}
      _ -> {:error, :invalid_code}
    end
  end

  # -- Token Exchange --

  defp exchange_code(metadata, cfg, code, code_verifier) do
    token_endpoint = metadata["token_endpoint"]

    body_params =
      %{
        "grant_type" => "authorization_code",
        "code" => code,
        "redirect_uri" => cfg.redirect_uri,
        "client_id" => cfg.client_id,
        "code_verifier" => code_verifier
      }
      |> Map.merge(resource_param(cfg))

    # Add client_secret if provided (confidential client)
    body_params =
      if cfg[:client_secret] do
        Map.put(body_params, "client_secret", cfg.client_secret)
      else
        body_params
      end

    make_token_request(token_endpoint, body_params)
  end

  # -- Token Refresh --

  defp refresh_with_server(metadata, cfg, refresh_token) do
    token_endpoint = metadata["token_endpoint"]

    body_params =
      %{
        "grant_type" => "refresh_token",
        "refresh_token" => refresh_token,
        "client_id" => cfg.client_id
      }
      |> Map.merge(resource_param(cfg))

    # Add client_secret if provided
    body_params =
      if cfg[:client_secret] do
        Map.put(body_params, "client_secret", cfg.client_secret)
      else
        body_params
      end

    make_token_request(token_endpoint, body_params)
  end

  # -- HTTP Request Helper --

  defp make_token_request(token_endpoint, body_params) do
    headers = [
      {"Content-Type", "application/x-www-form-urlencoded"},
      {"Accept", "application/json"}
    ]

    body = URI.encode_query(body_params)

    case Http.Client.impl().post(token_endpoint, body, headers,
           recv_timeout: 15_000,
           timeout: 15_000
         ) do
      {:ok, %{status_code: 200, body: response_body}} ->
        case SafeJson.decode(response_body) do
          {:ok, tokens} when is_map(tokens) ->
            {:ok, tokens}

          {:ok, _} ->
            {:error, {:invalid_response, "Token response is not a JSON object"}}

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

      {:ok, %{status_code: code, body: _error_body}} ->
        {:error, {:http_error, code}}

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

  # -- Token Normalization --

  defp normalize_tokens(tokens) do
    # Calculate expires_at from provider token response.
    # Prefer a provider-supplied "expires_at" where present; otherwise use "expires_in".
    expires_at =
      cond do
        is_integer(Map.get(tokens, "expires_at")) ->
          Map.get(tokens, "expires_at")

        true ->
          case Map.get(tokens, "expires_in") do
            secs when is_integer(secs) and secs > 0 ->
              System.os_time(:second) + secs

            _ ->
              # Treat 0 or nil as "unknown" (no expiry information). Don't set expires_at.
              nil
          end
      end

    # Debug: show raw fields used to compute expiry when debug enabled (redact access_token)
    if Util.Env.mcp_debug_enabled?() do
      raw = %{
        "raw_expires_in" => Map.get(tokens, "expires_in"),
        "raw_expires_at" => Map.get(tokens, "expires_at"),
        "computed_expires_at" => expires_at
      }

      # Redact access token before printing
      dbg =
        Map.put(
          raw,
          "access_token",
          if(Map.has_key?(tokens, "access_token"), do: "<redacted>", else: nil)
        )

      UI.debug("[MCP Debug] token fields used for expiry calculation")
      UI.printf_debug(dbg)
    end

    %{
      access_token: Map.fetch!(tokens, "access_token"),
      token_type: Map.get(tokens, "token_type", "Bearer"),
      expires_at: expires_at,
      refresh_token: Map.get(tokens, "refresh_token"),
      scope: Map.get(tokens, "scope")
    }
  end
end