defmodule Gemini.Auth.JWT do
  @moduledoc """
  JWT token generation and management for Google Cloud service accounts.

  This module handles JWT creation and signing for Vertex AI authentication
  based on the documentation in v1.md. It supports both service account key files
  and the Google Cloud IAM signJwt API.
  """

  alias Joken.Signer

  @iam_credentials_url "https://iamcredentials.googleapis.com/v1"
  # 1 hour in seconds
  @default_token_lifetime 3600

  @type service_account_key :: %{
          type: String.t(),
          project_id: String.t(),
          private_key_id: String.t(),
          private_key: String.t(),
          client_email: String.t(),
          client_id: String.t(),
          auth_uri: String.t(),
          token_uri: String.t(),
          auth_provider_x509_cert_url: String.t(),
          client_x509_cert_url: String.t()
        }

  @type jwt_payload :: %{
          required(:iss) => String.t(),
          required(:aud) => String.t(),
          required(:sub) => String.t(),
          required(:iat) => integer(),
          required(:exp) => integer(),
          optional(:scope) => String.t(),
          optional(atom()) => any()
        }

  @doc """
  Create a JWT payload for Vertex AI authentication.

  ## Parameters
  - `service_account_email`: The email of the service account (issuer)
  - `audience`: The audience for the JWT (must match deployment config)
  - `opts`: Optional parameters
    - `:lifetime` - Token lifetime in seconds (default: 3600)
    - `:issued_at` - Custom issued at time (default: current time)

  ## Examples

      iex> payload = Gemini.Auth.JWT.create_payload(
      ...>   "my-service@project.iam.gserviceaccount.com",
      ...>   "my-app-audience"
      ...> )
      iex> payload.iss
      "my-service@project.iam.gserviceaccount.com"
  """
  @spec create_payload(String.t(), String.t(), keyword()) :: jwt_payload()
  def create_payload(service_account_email, audience, opts \\ []) do
    now = Keyword.get(opts, :issued_at, System.system_time(:second))
    lifetime = Keyword.get(opts, :lifetime, @default_token_lifetime)

    %{
      iss: service_account_email,
      aud: audience,
      # Both aud and sub should be the same according to v1.md
      sub: audience,
      iat: now,
      exp: now + lifetime
    }
  end

  @doc """
  Sign a JWT payload using a service account private key.

  This method signs the JWT locally using the private key from the service account JSON file.

  ## Examples

      iex> key = %{private_key: "-----BEGIN PRIVATE KEY-----...", client_email: "..."}
      iex> payload = %{iss: "...", aud: "...", sub: "...", iat: 123, exp: 456}
      iex> {:ok, token} = Gemini.Auth.JWT.sign_with_key(payload, key)
  """
  @spec sign_with_key(jwt_payload(), service_account_key()) ::
          {:ok, String.t()} | {:error, term()}
  def sign_with_key(payload, %{private_key: private_key}) do
    # Create RS256 signer from the private key
    signer = Signer.create("RS256", %{"pem" => private_key})

    # Generate the token
    case Joken.generate_and_sign(%{}, payload, signer) do
      {:ok, token, _claims} -> {:ok, token}
      {:error, reason} -> {:error, reason}
    end
  rescue
    error -> {:error, error}
  end

  @doc """
  Sign a JWT payload using Google Cloud IAM signJwt API.

  This method uses the Google Cloud IAM service to sign the JWT, which requires
  the caller to have the `roles/iam.serviceAccountTokenCreator` role.

  ## Parameters
  - `payload`: The JWT payload to sign
  - `service_account_email`: The service account email
  - `access_token`: A valid access token for authentication

  ## Examples

      iex> payload = %{iss: "...", aud: "...", sub: "...", iat: 123, exp: 456}
      iex> {:ok, token} = Gemini.Auth.JWT.sign_with_iam_api(
      ...>   payload,
      ...>   "my-service@project.iam.gserviceaccount.com",
      ...>   "ya29.access-token"
      ...> )
  """
  @spec sign_with_iam_api(jwt_payload(), String.t(), String.t()) ::
          {:ok, String.t()} | {:error, term()}
  def sign_with_iam_api(payload, service_account_email, access_token) do
    url = "#{@iam_credentials_url}/projects/-/serviceAccounts/#{service_account_email}:signJwt"

    headers = [
      {"Authorization", "Bearer #{access_token}"},
      {"Content-Type", "application/json"}
    ]

    # Convert payload to JSON string as required by the API
    body =
      Jason.encode!(%{
        "payload" => Jason.encode!(payload)
      })

    case Req.post(url, headers: headers, body: body) do
      {:ok, response} ->
        handle_iam_response(response)

      {:error, reason} ->
        {:error, "Request failed: #{reason}"}
    end
  end

  @doc """
  Load and parse a service account key file.

  ## Examples

      iex> {:ok, key} = Gemini.Auth.JWT.load_service_account_key("/path/to/key.json")
      iex> key.client_email
      "my-service@project.iam.gserviceaccount.com"
  """
  @spec load_service_account_key(String.t()) :: {:ok, service_account_key()} | {:error, term()}
  def load_service_account_key(file_path) do
    case File.read(file_path) do
      {:ok, content} ->
        case Jason.decode(content) do
          {:ok, key_data} ->
            # Convert string keys to atom keys for easier access
            parsed_key =
              for {key, val} <- key_data, into: %{} do
                {String.to_atom(key), val}
              end

            {:ok, parsed_key}

          {:error, %Jason.DecodeError{} = reason} ->
            {:error, "Failed to parse JSON: #{Exception.message(reason)}"}
        end

      {:error, reason} ->
        {:error, "Failed to read file: #{reason}"}
    end
  end

  @doc """
  Validate a JWT payload has all required fields.

  ## Examples

      iex> payload = %{iss: "test", aud: "test", sub: "test", iat: 123, exp: 456}
      iex> Gemini.Auth.JWT.validate_payload(payload)
      :ok
  """
  @spec validate_payload(jwt_payload()) :: :ok | {:error, String.t()}
  def validate_payload(%{iss: iss, aud: aud, sub: sub, iat: iat, exp: exp})
      when is_binary(iss) and is_binary(aud) and is_binary(sub) and
             is_integer(iat) and is_integer(exp) do
    now = System.system_time(:second)

    case find_payload_error(%{aud: aud, sub: sub, iat: iat, exp: exp}, now) do
      nil -> :ok
      {:error, _} = error -> error
    end
  end

  def validate_payload(_payload) do
    {:error, "Missing required JWT claims: iss, aud, sub, iat, exp"}
  end

  @doc """
  Create a signed JWT token using the most appropriate method.

  This function automatically chooses between local signing (if private key is available)
  or IAM API signing (if access token is provided).

  ## Examples

      # Using service account key file
      iex> {:ok, token} = Gemini.Auth.JWT.create_signed_token(
      ...>   "my-service@project.iam.gserviceaccount.com",
      ...>   "my-app-audience",
      ...>   service_account_key: "/path/to/key.json"
      ...> )

      # Using IAM API
      iex> {:ok, token} = Gemini.Auth.JWT.create_signed_token(
      ...>   "my-service@project.iam.gserviceaccount.com",
      ...>   "my-app-audience",
      ...>   access_token: "ya29.access-token"
      ...> )
  """
  @spec create_signed_token(String.t(), String.t(), keyword()) ::
          {:ok, String.t()} | {:error, term()}
  def create_signed_token(service_account_email, audience, opts \\ []) do
    payload = create_payload(service_account_email, audience, opts)

    with :ok <- validate_payload(payload) do
      cond do
        service_account_key = Keyword.get(opts, :service_account_key) ->
          with {:ok, key} <- load_service_account_key(service_account_key) do
            sign_with_key(payload, key)
          end

        service_account_data = Keyword.get(opts, :service_account_data) ->
          sign_with_key(payload, service_account_data)

        access_token = Keyword.get(opts, :access_token) ->
          sign_with_iam_api(payload, service_account_email, access_token)

        true ->
          {:error,
           "Either service_account_key, service_account_data, or access_token must be provided"}
      end
    end
  end

  @doc """
  Extract service account email from a service account key.

  ## Examples

      iex> key = %{client_email: "my-service@project.iam.gserviceaccount.com"}
      iex> Gemini.Auth.JWT.get_service_account_email(key)
      "my-service@project.iam.gserviceaccount.com"
  """
  @spec get_service_account_email(service_account_key()) :: String.t()
  def get_service_account_email(%{client_email: email}), do: email

  defp handle_iam_response(%Req.Response{status: 200, body: body}), do: parse_signed_jwt(body)

  defp handle_iam_response(%Req.Response{status: status, body: body}) do
    error_body = if is_binary(body), do: body, else: inspect(body)
    {:error, "HTTP #{status}: #{error_body}"}
  end

  defp parse_signed_jwt(%{"signedJwt" => signed_jwt}), do: {:ok, signed_jwt}

  defp parse_signed_jwt(body) when is_binary(body) do
    case Jason.decode(body) do
      {:ok, decoded} -> parse_signed_jwt(decoded)
      {:error, reason} -> {:error, "Failed to parse response: #{reason}"}
    end
  end

  defp parse_signed_jwt(response),
    do: {:error, "Unexpected response format: #{inspect(response)}"}

  defp find_payload_error(payload, now) do
    validators = [
      &validate_aud_sub/2,
      &validate_exp_gt_iat/2,
      &validate_iat_not_future/2,
      &validate_exp_not_past/2
    ]

    Enum.find_value(validators, fn validator -> validator.(payload, now) end)
  end

  defp validate_aud_sub(%{aud: aud, sub: sub}, _now) do
    if aud == sub, do: nil, else: {:error, "aud and sub claims must be identical for Vertex AI"}
  end

  defp validate_exp_gt_iat(%{exp: exp, iat: iat}, _now) do
    if exp > iat, do: nil, else: {:error, "exp must be greater than iat"}
  end

  defp validate_iat_not_future(%{iat: iat}, now) do
    if iat <= now + 60, do: nil, else: {:error, "iat cannot be in the future"}
  end

  defp validate_exp_not_past(%{exp: exp}, now) do
    if exp >= now, do: nil, else: {:error, "Token has expired"}
  end
end