defmodule Store.APIUsage do
  @moduledoc """
  Module for recording and checking API usage data. Cross OS process
  coordination is handled with regard to file reads/writes. It is up to the
  caller to ensure that requests are internally ordered for consistency.
  """

  @path "usage.json"
  @tag "[usage]"
  @re_reset ~r/^(\d+(?:\.\d+)?)(ms|s)$/
  @debug_env_var "FNORD_DEBUG_API_USAGE"

  @type model_usage :: %{
          updated_at: non_neg_integer,
          requests_max: non_neg_integer,
          requests_left: non_neg_integer,
          requests_reset: non_neg_integer,
          tokens_max: non_neg_integer,
          tokens_left: non_neg_integer,
          tokens_reset: non_neg_integer
        }

  @type usage :: %{
          optional(binary) => model_usage
        }

  @type status :: non_neg_integer

  @type ok_result ::
          {:ok,
           %{
             body: map,
             headers: list,
             status: status
           }}

  @type http_result ::
          ok_result
          | {:http_error, {status, String.t()}}
          | {:transport_error, any}

  @doc """
  Returns the path to the usage store file.
  """
  @spec store_path() :: binary
  def store_path do
    Path.join(Store.store_home(), @path)
  end

  @doc """
  Records API usage for a provided model. If model is nil, returns http_result
  unchanged. For 2xx {:ok, %{headers: headers, status: code}} results and a
  binary model, records usage data and updates the store. Other responses or
  errors are returned unchanged.
  """
  @spec record_for_model(binary | nil, http_result) :: http_result
  def record_for_model(nil, http_result), do: http_result
  def record_for_model(_, {:http_error, _} = http_result), do: http_result
  def record_for_model(_, {:transport_error, _} = http_result), do: http_result

  def record_for_model(_, {:ok, %{status: code} = payload})
      when code < 200 or
             code >= 300 do
    {:ok, payload}
  end

  # HTTP success responses
  def record_for_model(model, {:ok, %{headers: headers} = payload}) do
    ensure_store_file()
    headers_map = Enum.into(headers, %{})

    with {:ok, usage} <- collect_usage_data(headers_map),
         :ok <- update_file(model, usage) do
      case Util.Env.get_env(@debug_env_var, nil) do
        "" -> :ok
        "0" -> :ok
        nil -> :ok
        _ -> UI.debug(@tag, "#{model}: #{inspect(usage, pretty: true)}")
      end
    else
      _reason -> :ok
    end

    {:ok, payload}
  end

  @doc """
  Checks if a request can be made for the given model based on the last
  recorded usage data for the model. If a request can be made, returns `:ok`.
  Otherwise, returns `{:wait, milliseconds}` indicating how long to wait before
  attempting another request.
  """
  @spec check(binary) ::
          :ok
          | {:wait, non_neg_integer}
          | {:error, term}
  def check(model) do
    ensure_store_file()

    with :ok <- can_request?(model) do
      :ok
    end
  end

  # ----------------------------------------------------------------------------
  # Internals
  # ----------------------------------------------------------------------------
  @spec can_request?(binary) ::
          :ok
          | {:wait, non_neg_integer}
          | {:error, term}
  defp can_request?(model) do
    FileLock.with_lock(store_path(), fn ->
      with {:ok, data} <- read_file() do
        usage = Map.get(data, model, %{})
        updated_at = Map.get(usage, :updated_at, 0)
        elapsed_time = now() - updated_at
        requests_reset = Map.get(usage, :requests_reset, 0)

        if requests_reset <= elapsed_time do
          :ok
        else
          {:wait, requests_reset - elapsed_time}
        end
      end
    end)
    |> case do
      {:ok, result} -> result
      {:callback_error, exc} -> {:error, exc}
      {:error, reason} -> {:error, reason}
    end
  end

  @spec update_file(binary, model_usage) :: :ok | {:error, term}
  defp update_file(model, usage) do
    FileLock.with_lock(store_path(), fn ->
      with {:ok, usage_data} <- read_file() do
        usage = Map.put(usage, :updated_at, now())
        payload = Map.put(usage_data, model, usage)
        write_file(payload)
      end
    end)
    |> case do
      {:ok, result} -> result
      {:callback_error, exc} -> {:error, exc}
      {:error, reason} -> {:error, reason}
    end
  end

  @spec read_file() :: {:ok, usage} | {:error, term}
  defp read_file do
    with {:ok, json} <- File.read(store_path()),
         {:ok, data} <- SafeJson.decode(json) do
      data
      |> Enum.map(fn {model, usage} ->
        {model, Util.string_keys_to_atoms(usage)}
      end)
      |> Enum.into(%{})
      |> then(&{:ok, &1})
    end
  end

  @spec write_file(map) :: :ok | {:error, term}
  defp write_file(data) do
    with {:ok, json} <- SafeJson.encode(data, pretty: true) do
      File.write(store_path(), json)
    end
  end

  @spec collect_usage_data(map) :: {:ok, model_usage} | {:error, atom, binary}
  defp collect_usage_data(headers) do
    with {:ok, requests_max} <- collect_usage_metric(headers, "x-ratelimit-limit-requests"),
         {:ok, requests_left} <- collect_usage_metric(headers, "x-ratelimit-remaining-requests"),
         {:ok, requests_reset} <- collect_reset_metric(headers, "x-ratelimit-reset-requests"),
         {:ok, tokens_max} <- collect_usage_metric(headers, "x-ratelimit-limit-tokens"),
         {:ok, tokens_left} <- collect_usage_metric(headers, "x-ratelimit-remaining-tokens"),
         {:ok, tokens_reset} <- collect_reset_metric(headers, "x-ratelimit-reset-tokens") do
      {:ok,
       %{
         updated_at: now(),
         requests_max: requests_max,
         requests_left: requests_left,
         requests_reset: requests_reset,
         tokens_max: tokens_max,
         tokens_left: tokens_left,
         tokens_reset: tokens_reset
       }}
    end
  end

  @spec collect_usage_metric(map, binary) :: {:ok, non_neg_integer} | {:error, atom, binary}
  defp collect_usage_metric(headers, key) do
    headers
    |> Map.fetch(key)
    |> case do
      {:ok, value} ->
        value
        |> Integer.parse()
        |> case do
          {int_value, _} -> {:ok, int_value}
          :error -> {:error, :invalid_integer}
        end

      :error ->
        {:error, :not_found, key}
    end
  end

  @spec collect_reset_metric(map, binary) :: {:ok, non_neg_integer} | {:error, atom, binary}
  defp collect_reset_metric(headers, key) do
    headers
    |> Map.fetch(key)
    |> case do
      # In "2ms" or "1s" format
      {:ok, value} -> parse_reset_metric(value)
      :error -> {:error, :not_found, key}
    end
  end

  @spec parse_reset_metric(binary) :: {:ok, non_neg_integer} | {:error, atom, binary}
  defp parse_reset_metric(value) do
    with [amount, unit] <- Regex.run(@re_reset, value, capture: :all_but_first),
         {float_value, _rest} <- Float.parse(amount) do
      ms_value =
        case unit do
          "ms" -> float_value
          "s" -> float_value * 1000.0
        end
        # Be conservative: never under-wait.
        |> Float.ceil()
        |> trunc()
        |> max(0)

      {:ok, ms_value}
    else
      _ -> {:error, :invalid_reset_format, value}
    end
  end

  @spec now() :: non_neg_integer
  defp now do
    DateTime.utc_now()
    |> DateTime.to_unix(:millisecond)
  end

  @spec ensure_store_file() :: :ok | {:error, term}
  defp ensure_store_file() do
    path = store_path()

    unless File.exists?(path) do
      File.write(path, "{}")
    end

    :ok
  end
end