defmodule Oban.Pro.Migrations.V150 do
  @moduledoc false

  use Ecto.Migration

  def up_ddl(opts) do
    if not partitioned?(opts) do
      execute """
      CREATE OR REPLACE FUNCTION #{opts.quoted}.oban_state_to_bit(state #{opts.quoted}.oban_job_state)
      RETURNS jsonb AS $$
      SELECT CASE
             WHEN state = 'scheduled' THEN '0'::jsonb
             WHEN state = 'available' THEN '1'::jsonb
             WHEN state = 'executing' THEN '2'::jsonb
             WHEN state = 'retryable' THEN '3'::jsonb
             WHEN state = 'completed' THEN '4'::jsonb
             WHEN state = 'cancelled' THEN '5'::jsonb
             WHEN state = 'discarded' THEN '6'::jsonb
             END;
      $$ LANGUAGE SQL IMMUTABLE STRICT
      """

      alter table(:oban_jobs, prefix: opts.prefix) do
        add_if_not_exists :uniq_key, :text,
          generated: """
          ALWAYS AS (CASE WHEN meta->'uniq_bmp' @> #{opts.quoted}.oban_state_to_bit(state) THEN meta->>'uniq_key' END) STORED
          """
      end
    end
  end

  def up_idx(opts) do
    if partitioned?(opts) do
      create_if_not_exists index(
                             :oban_jobs,
                             ["(meta->>'uniq_key')"],
                             concurrently: opts.concurrently,
                             name: :oban_jobs_unique_index,
                             prefix: opts.prefix,
                             where: "meta ? 'uniq_key'"
                           )
    else
      create_if_not_exists index(
                             :oban_jobs,
                             [:uniq_key],
                             concurrently: opts.concurrently,
                             name: :oban_jobs_unique_index,
                             prefix: opts.prefix,
                             unique: true,
                             where: "uniq_key IS NOT NULL"
                           )
    end

    create_if_not_exists index(
                           :oban_jobs,
                           [:state, "(meta->>'batch_id')", "(meta->>'callback')"],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_batch_index,
                           prefix: opts.prefix,
                           where: "meta ? 'batch_id'"
                         )

    create_if_not_exists index(
                           :oban_jobs,
                           [:state, "(meta->>'chain_id')", "(meta->>'on_hold')"],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_chain_index,
                           prefix: opts.prefix,
                           where: "meta ? 'chain_id'"
                         )
  end

  def down_ddl(opts) do
    alter table(:oban_jobs, prefix: opts.prefix) do
      remove_if_exists :uniq_key, :text
    end

    execute "DROP FUNCTION IF EXISTS #{opts.quoted}.oban_state_to_bit"
  end

  def down_idx(opts) do
    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_batch_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_chain_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_unique_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )
  end

  defp partitioned?(opts) do
    query = """
    SELECT EXISTS (
      SELECT 1
      FROM pg_tables
      WHERE schemaname = '#{opts.prefix}' AND tablename = 'oban_jobs_completed'
    )
    """

    {:ok, %{rows: [[partitioned?]]}} = repo().query(query, [], log: false)

    partitioned?
  end
end
