defmodule PolicyServiceWeb.Plugs.AuthenticationPlug do
  @moduledoc """
  Authentication plug for validating JWT tokens using Zitadel.

  This plug validates JWT tokens using the oidcc library and extracts
  user claims for use throughout the application.
  """

  import Plug.Conn
  require Logger

  @doc """
  Initializes the authentication plug.

  ## Options
  - :provider - The OIDCC provider configuration worker name (required)
  """
  def init(opts) do
    provider = Keyword.fetch!(opts, :provider)

    zitadel_config = Application.get_env(:policy_service, :zitadel, [])

    %{
      provider: provider,
      client_id: Keyword.get(zitadel_config, :client_id),
      client_secret: Keyword.get(zitadel_config, :client_secret),
      required_scopes: Keyword.get(zitadel_config, :required_scopes, [])
    }
  end

  @doc """
  Authenticates the request by validating the JWT token.

  Extracts the Authorization header, validates the JWT token using
  Zitadel's JWKS, and assigns the extracted claims to the connection.
  """
  def call(conn, config) do
    with {:ok, token} <- extract_token(conn),
         {:ok, validated_token} <- validate_token(token, config),
         {:ok, claims} <- extract_claims(validated_token) do
      conn
      |> assign(:authenticated, true)
      |> assign(:current_user, claims)
      |> assign(:user_id, claims.user_id)
      |> assign(:org_id, claims.org_id)
      |> assign(:roles, claims.roles)
      |> assign(:scopes, claims.scopes)
    else
      {:error, :missing_token} ->
        handle_missing_token(conn)

      {:error, :invalid_token} = error ->
        handle_invalid_token(conn, error)

      {:error, :expired_token} = error ->
        handle_expired_token(conn, error)

      {:error, :insufficient_scopes} = error ->
        handle_insufficient_scopes(conn, error)

      {:error, reason} = error ->
        handle_authentication_error(conn, error, reason)
    end
  end

  defp extract_token(conn) do
    case get_req_header(conn, "authorization") do
      ["Bearer " <> token] when token != "" ->
        {:ok, token}

      ["Bearer " <> token] when token == "" ->
        {:error, :missing_token}

      ["bearer " <> token] when token != "" ->
        {:ok, token}

      [] ->
        {:error, :missing_token}

      _ ->
        {:error, :invalid_token}
    end
  end

  defp validate_token(token, config) do
    try do
      case Oidcc.Token.validate_jwt(
             token,
             config.provider,
             %{}
           ) do
        {:ok, validated_token} ->
          validate_required_scopes(validated_token, config)

        {:error, :invalid_token} ->
          Logger.warning("Invalid JWT token provided")
          {:error, :invalid_token}

        {:error, :expired_token} ->
          Logger.warning("Expired JWT token provided")
          {:error, :expired_token}

        {:error, reason} ->
          Logger.error("Token validation failed: #{inspect(reason)}")
          {:error, :invalid_token}
      end
    rescue
      error ->
        Logger.error("Token validation error: #{inspect(error)}")
        {:error, :invalid_token}
    end
  end

  defp validate_required_scopes(validated_token, config) do
    required_scopes = config.required_scopes || []

    if required_scopes == [] do
      {:ok, validated_token}
    else
      token_scopes = PolicyServiceWeb.Plugs.ClaimsExtractor.get_scopes(validated_token.claims)

      if has_all_required_scopes?(token_scopes, required_scopes) do
        {:ok, validated_token}
      else
        missing_scopes = required_scopes -- token_scopes
        Logger.warning("Token missing required scopes: #{inspect(missing_scopes)}")
        {:error, :insufficient_scopes}
      end
    end
  end

  defp has_all_required_scopes?(token_scopes, required_scopes) do
    Enum.all?(required_scopes, fn scope -> scope in token_scopes end)
  end

  defp extract_claims(validated_token) do
    try do
      claims = PolicyServiceWeb.Plugs.ClaimsExtractor.extract_claims(validated_token)
      PolicyServiceWeb.Plugs.ClaimsExtractor.validate_claims!(validated_token.claims)
      {:ok, claims}
    rescue
      error ->
        Logger.error("Failed to extract claims: #{inspect(error)}")
        {:error, :invalid_claims}
    end
  end

  defp handle_missing_token(conn) do
    conn
    |> put_resp_content_type("application/json")
    |> send_resp(401, Jason.encode!(%{error: "Missing authentication token"}))
    |> halt()
  end

  defp handle_invalid_token(conn, _error) do
    conn
    |> put_resp_content_type("application/json")
    |> send_resp(401, Jason.encode!(%{error: "Invalid authentication token"}))
    |> halt()
  end

  defp handle_expired_token(conn, _error) do
    conn
    |> put_resp_content_type("application/json")
    |> send_resp(401, Jason.encode!(%{error: "Authentication token has expired"}))
    |> halt()
  end

  defp handle_insufficient_scopes(conn, _error) do
    conn
    |> put_resp_content_type("application/json")
    |> send_resp(
      403,
      Jason.encode!(%{error: "Insufficient permissions - required scopes not granted"})
    )
    |> halt()
  end

  defp handle_authentication_error(conn, _error, reason) do
    Logger.error("Authentication error: #{inspect(reason)}")

    conn
    |> put_resp_content_type("application/json")
    |> send_resp(401, Jason.encode!(%{error: "Authentication failed"}))
    |> halt()
  end
end