require Protocol

defmodule Oban.Pro.Producer do
  @moduledoc false

  use Ecto.Schema

  import Ecto.Changeset

  alias Ecto.Changeset
  alias Oban.Pro.Utils

  defmodule ListOrMap do
    @moduledoc false

    # This is a custom type to allow a graceful transition from the legacy rate limit windows
    # structure to the new, flattened map.

    @behaviour Ecto.Type

    @impl true
    def type, do: :any

    @impl true
    def cast(data) when is_map(data), do: {:ok, data}
    def cast([data | _]), do: {:ok, data}
    def cast(_data), do: :error

    @impl true
    def load(data) when is_map(data), do: {:ok, data}
    def load([data | _]), do: {:ok, data}
    def load([]), do: {:ok, %{}}
    def load(_data), do: :error

    @impl true
    def dump(data) when is_map(data), do: {:ok, data}
    def dump([data | _]), do: {:ok, data}
    def dump([]), do: {:ok, %{}}
    def dump(_data), do: :error

    @impl true
    def equal?(a, b), do: a == b

    @impl true
    def embed_as(_data), do: :dump
  end

  @type t :: %__MODULE__{
          uuid: Ecto.UUID.t(),
          name: binary(),
          node: binary(),
          queue: binary(),
          meta: map(),
          started_at: DateTime.t(),
          updated_at: DateTime.t()
        }

  @type meta :: %__MODULE__.Meta{local_limit: pos_integer(), paused: boolean()}

  @primary_key {:uuid, :binary_id, autogenerate: false}
  schema "oban_producers" do
    field :name, :string
    field :node, :string
    field :queue, :string
    field :started_at, :utc_datetime_usec
    field :updated_at, :utc_datetime_usec

    field :ack_async, :boolean, virtual: true, default: true
    field :ack_tab, :any, virtual: true
    field :refresh_interval, :integer, virtual: true, default: :timer.seconds(30)
    field :xact_delay, :integer, default: :timer.seconds(1), virtual: true
    field :xact_retry, :integer, default: 5, virtual: true
    field :xact_timeout, :integer, default: :timer.seconds(30), virtual: true

    embeds_one :meta, Meta, on_replace: :update, primary_key: false do
      @moduledoc false

      field :local_limit, :integer, default: 10
      field :paused, :boolean, default: false
      field :shutdown_started_at, :utc_datetime_usec

      embeds_one :global_limit, GlobalLimit, on_replace: :update, primary_key: false do
        @moduledoc false

        field :allowed, :integer
        field :burst, :boolean, default: false
        field :tracked, ListOrMap, default: %{}

        embeds_one :partition, Partition, on_replace: :update, primary_key: false do
          @moduledoc false

          field :fields, {:array, Ecto.Enum}, values: [:args, :meta, :worker]
          field :keys, {:array, :string}
        end
      end

      embeds_one :rate_limit, RateLimit, on_replace: :update, primary_key: false do
        @moduledoc false

        field :algorithm, Ecto.Enum,
          values: ~w(sliding_window fixed_window token_bucket)a,
          default: :sliding_window

        field :allowed, :integer
        field :period, :integer
        field :window_time, :integer, skip_default_validation: true
        field :windows, ListOrMap, default: %{}

        embeds_one :partition, Partition, on_replace: :update, primary_key: false do
          @moduledoc false

          field :fields, {:array, Ecto.Enum}, values: [:args, :meta, :worker]
          field :keys, {:array, :string}
        end
      end
    end
  end

  @spec new(map() | Keyword.t()) :: Changeset.t(t())
  def new(params) when is_list(params) or is_map(params) do
    params =
      params
      |> Map.new()
      |> Map.put_new_lazy(:uuid, &Oban.Pro.UUIDv7.generate/0)
      |> Map.update!(:name, &to_clean_string/1)
      |> Map.update!(:node, &to_clean_string/1)
      |> Map.update!(:queue, &to_clean_string/1)

    %__MODULE__{}
    |> cast(params, ~w(name node queue refresh_interval started_at updated_at uuid)a)
    |> cast_embed(:meta, required: true, with: &meta_changeset/2)
    |> validate_required(~w(name node queue)a)
    |> validate_length(:name, min: 1)
    |> validate_length(:node, min: 1)
    |> validate_length(:queue, min: 1)
    |> Utils.enforce_keys(params, __MODULE__)
  end

  @spec update_meta(t(), atom(), term()) :: Changeset.t()
  def update_meta(schema, key, value) do
    update_meta(schema, %{key => value})
  end

  @spec update_meta(t(), map()) :: Changeset.t()
  def update_meta(schema, changes) do
    meta =
      schema.meta
      |> meta_changeset(changes)
      |> apply_changes()
      |> Ecto.embedded_dump(:json)

    schema
    |> cast(%{meta: meta}, [])
    |> cast_embed(:meta, with: &meta_changeset/2)
  end

  @allowed ~w(local_limit paused shutdown_started_at)a

  @doc false
  def meta_changeset(schema, params, allowed \\ @allowed) do
    params =
      params
      |> Map.new()
      |> Map.delete(:limit)
      |> Map.put_new_lazy(:local_limit, fn -> default_local_limit(params, schema) end)

    params =
      if Map.get(params, :global_limit) do
        Map.update!(params, :global_limit, &cast_global_limit/1)
      else
        params
      end

    params =
      if Map.get(params, :rate_limit) do
        Map.update!(params, :rate_limit, &cast_rate_limit/1)
      else
        params
      end

    schema
    |> cast(params, allowed)
    |> cast_embed(:global_limit, with: &global_changeset/2)
    |> cast_embed(:rate_limit, with: &rate_changeset/2)
    |> validate_required(~w(local_limit)a)
    |> validate_number(:local_limit, greater_than: 0)
    |> validate_single_partitioner()
    |> Utils.enforce_keys(params, schema.__struct__)
  end

  @doc false
  @spec default_local_limit(map(), Ecto.Schema.t()) :: non_neg_integer()
  def default_local_limit(params, schema) do
    cond do
      is_integer(params[:limit]) ->
        params[:limit]

      is_integer(params[:global_limit]) ->
        params[:global_limit]

      is_integer(get_in(params, [:global_limit, :allowed])) ->
        get_in(params, [:global_limit, :allowed])

      true ->
        schema.local_limit
    end
  end

  @doc false
  @spec validate_single_partitioner(Changeset.t()) :: Changeset.t()
  def validate_single_partitioner(changeset) do
    case {get_field(changeset, :global_limit), get_field(changeset, :rate_limit)} do
      {%{partition: %{}}, %{partition: %{}}} ->
        add_error(changeset, :global_limit, "only one limiter may have partitioning")

      _ ->
        changeset
    end
  end

  @doc false
  @spec global_changeset(Ecto.Schema.t(), integer() | map() | Keyword.t()) :: Changeset.t()
  def global_changeset(schema, params) do
    params = Map.update(params, :partition, nil, &normalize_partition/1)

    schema
    |> cast(params, ~w(allowed burst)a)
    |> cast_embed(:partition, with: &partition_changeset/2)
    |> validate_number(:allowed, greater_than: 0)
    |> Utils.enforce_keys(params, schema.__struct__)
  end

  @doc false
  @spec rate_changeset(Ecto.Schema.t(), map() | Keyword.t()) :: Changeset.t()
  def rate_changeset(schema, params) do
    params =
      params
      |> Map.update(:period, nil, &Utils.cast_period/1)
      |> Map.update(:partition, nil, &normalize_partition/1)
      |> Map.put_new_lazy(:window_time, fn -> DateTime.to_unix(DateTime.utc_now(), :second) end)

    fields = schema.__struct__.__schema__(:fields) -- [:partition]

    schema
    |> cast(params, fields)
    |> cast_embed(:partition, with: &partition_changeset/2)
    |> validate_required(~w(allowed period)a)
    |> validate_number(:allowed, greater_than: 0)
    |> validate_number(:period, greater_than: 0)
    |> clear_windows_on_algorithm_change()
    |> Utils.enforce_keys(params, schema.__struct__)
  end

  defp clear_windows_on_algorithm_change(changeset) do
    case get_change(changeset, :algorithm) do
      nil -> changeset
      _new_algorithm -> force_change(changeset, :windows, %{})
    end
  end

  @doc false
  @spec partition_changeset(Ecto.Schema.t(), map() | Keyword.t()) :: Changeset.t()
  def partition_changeset(schema, params) do
    params =
      params
      |> Map.new()
      |> Map.update(:keys, [], &for(key <- &1, do: to_string(key)))

    schema
    |> cast(params, ~w(fields keys)a)
    |> validate_required(~w(fields)a)
    |> Utils.enforce_keys(params, schema.__struct__)
  end

  # Global Limit Helpers

  @doc false
  @spec cast_global_limit(pos_integer() | list() | map()) :: map()
  def cast_global_limit(limit) when is_integer(limit) do
    %{allowed: limit}
  end

  def cast_global_limit(opts), do: cast_rate_limit(opts)

  # Rate Limit Helpers

  @doc false
  @spec cast_rate_limit(list() | map()) :: map()
  def cast_rate_limit([[key | _] | _] = opts) when is_binary(key) do
    opts
    |> Enum.map(&List.to_tuple/1)
    |> Map.new()
    |> cast_rate_limit()
  end

  def cast_rate_limit(opts)
      when is_map_key(opts, "algorithm") or
             is_map_key(opts, "allowed") or
             is_map_key(opts, "burst") or
             is_map_key(opts, "period") or
             is_map_key(opts, "partition") do
    Map.new(opts, fn
      {"algorithm", algorithm} ->
        {:algorithm, String.to_existing_atom(algorithm)}

      {"allowed", allowed} ->
        {:allowed, allowed}

      {"burst", allowed} ->
        {:burst, allowed}

      {"period", [time, unit]} ->
        {:period, {time, unit}}

      {"period", period} ->
        {:period, period}

      {"partition", [["fields", fields]]} ->
        {:partition, fields: fields}

      {"partition", [["fields", fields], ["keys", keys]]} ->
        {:partition, fields: fields, keys: keys}

      {"partition", field} when is_binary(field) ->
        {:partition, fields: [field]}
    end)
  end

  def cast_rate_limit(opts), do: opts

  def normalize_partition(partition) do
    cond do
      is_nil(partition) ->
        nil

      is_map(partition) ->
        partition

      Keyword.keyword?(partition) and Keyword.has_key?(partition, :fields) ->
        partition

      true ->
        partition
        |> List.wrap()
        |> Enum.reduce([fields: [], keys: []], fn
          {field, keys}, opts when field in ~w(args meta)a ->
            opts
            |> Keyword.put(:keys, List.wrap(keys))
            |> Keyword.update!(:fields, &[field | &1])

          field, opts ->
            Keyword.update!(opts, :fields, &[field | &1])
        end)
    end
  end

  # Helpers

  defp to_clean_string(value) when is_reference(value) do
    inspect(value)
  end

  defp to_clean_string(value) do
    value
    |> to_string()
    |> String.trim_leading("Elixir.")
  end
end

encoder = if Code.ensure_loaded?(JSON.Encoder), do: JSON.Encoder, else: Jason.Encoder

Protocol.derive(encoder, Oban.Pro.Producer, except: [:__meta__])
Protocol.derive(encoder, Oban.Pro.Producer.Meta.GlobalLimit)
Protocol.derive(encoder, Oban.Pro.Producer.Meta.GlobalLimit.Partition)
Protocol.derive(encoder, Oban.Pro.Producer.Meta.RateLimit)
Protocol.derive(encoder, Oban.Pro.Producer.Meta.RateLimit.Partition)