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

  use Ecto.Migration

  import Oban.Pro.Migration, only: [partitioned?: 1]

  def up_ddl(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
           WHEN state = 'suspended' THEN '7'::jsonb
           END;
    $$ LANGUAGE SQL IMMUTABLE STRICT
    """

    create_if_not_exists table(:oban_workflows, primary_key: false, prefix: opts.prefix) do
      add :id, :text, primary_key: true
      add :name, :text
      add :parent_id, :text

      add :suspended, :integer, null: false, default: 0
      add :available, :integer, null: false, default: 0
      add :scheduled, :integer, null: false, default: 0
      add :executing, :integer, null: false, default: 0
      add :retryable, :integer, null: false, default: 0
      add :completed, :integer, null: false, default: 0
      add :cancelled, :integer, null: false, default: 0
      add :discarded, :integer, null: false, default: 0

      add :inserted_at, :utc_datetime_usec, null: false
      add :started_at, :utc_datetime_usec
      add :completed_at, :utc_datetime_usec

      add :meta, :map, null: false, default: %{}
    end

    execute """
    ALTER TABLE #{opts.quoted}.oban_workflows
    ADD COLUMN IF NOT EXISTS state TEXT GENERATED ALWAYS AS (
      CASE
        WHEN (suspended + available + scheduled + executing + retryable) > 0 THEN 'executing'
        WHEN cancelled > 0 AND NOT COALESCE((meta->>'ignore_cancelled')::boolean, false) THEN 'cancelled'
        WHEN discarded > 0 AND NOT COALESCE((meta->>'ignore_discarded')::boolean, false) THEN 'discarded'
        ELSE 'completed'
      END
    ) STORED
    """

    execute """
    CREATE OR REPLACE FUNCTION #{opts.quoted}.oban_workflows_update() RETURNS TRIGGER AS $$
    BEGIN
      UPDATE #{opts.quoted}.oban_workflows wf SET
        suspended = GREATEST(0, wf.suspended + COALESCE(d.suspended, 0)),
        available = GREATEST(0, wf.available + COALESCE(d.available, 0)),
        scheduled = GREATEST(0, wf.scheduled + COALESCE(d.scheduled, 0)),
        executing = GREATEST(0, wf.executing + COALESCE(d.executing, 0)),
        retryable = GREATEST(0, wf.retryable + COALESCE(d.retryable, 0)),
        completed = GREATEST(0, wf.completed + COALESCE(d.completed, 0)),
        cancelled = GREATEST(0, wf.cancelled + COALESCE(d.cancelled, 0)),
        discarded = GREATEST(0, wf.discarded + COALESCE(d.discarded, 0)),
        started_at = CASE
          WHEN wf.started_at IS NULL AND d.executing > 0 THEN NOW()
          ELSE wf.started_at
        END,
        completed_at = CASE
          WHEN wf.completed_at IS NULL
            AND (wf.suspended + wf.available + wf.scheduled + wf.executing + wf.retryable
                 + COALESCE(d.suspended, 0) + COALESCE(d.available, 0) + COALESCE(d.scheduled, 0)
                 + COALESCE(d.executing, 0) + COALESCE(d.retryable, 0)) <= 0
          THEN NOW()
          ELSE wf.completed_at
        END
      FROM (
        SELECT
          COALESCE(n.meta, o.meta)->>'workflow_id' AS workflow_id,
          SUM(CASE WHEN n.state = 'suspended' THEN 1 WHEN o.state = 'suspended' THEN -1 ELSE 0 END) AS suspended,
          SUM(CASE WHEN n.state = 'available' THEN 1 WHEN o.state = 'available' THEN -1 ELSE 0 END) AS available,
          SUM(CASE WHEN n.state = 'scheduled' THEN 1 WHEN o.state = 'scheduled' THEN -1 ELSE 0 END) AS scheduled,
          SUM(CASE WHEN n.state = 'executing' THEN 1 WHEN o.state = 'executing' THEN -1 ELSE 0 END) AS executing,
          SUM(CASE WHEN n.state = 'retryable' THEN 1 WHEN o.state = 'retryable' THEN -1 ELSE 0 END) AS retryable,
          SUM(CASE WHEN n.state = 'completed' THEN 1 WHEN o.state = 'completed' THEN -1 ELSE 0 END) AS completed,
          SUM(CASE WHEN n.state = 'cancelled' THEN 1 WHEN o.state = 'cancelled' THEN -1 ELSE 0 END) AS cancelled,
          SUM(CASE WHEN n.state = 'discarded' THEN 1 WHEN o.state = 'discarded' THEN -1 ELSE 0 END) AS discarded
        FROM new_table n
        FULL OUTER JOIN old_table o ON n.id = o.id
        WHERE (n.meta ? 'workflow_id' OR o.meta ? 'workflow_id')
          AND n.state IS DISTINCT FROM o.state
        GROUP BY COALESCE(n.meta, o.meta)->>'workflow_id'
      ) d
      WHERE wf.id = d.workflow_id;

      RETURN NULL;
    END;
    $$ LANGUAGE plpgsql;
    """

    execute """
    CREATE OR REPLACE FUNCTION #{opts.quoted}.oban_workflows_delete() RETURNS TRIGGER AS $$
    BEGIN
      UPDATE #{opts.quoted}.oban_workflows wf SET
        suspended = GREATEST(0, wf.suspended - COALESCE(d.suspended, 0)),
        available = GREATEST(0, wf.available - COALESCE(d.available, 0)),
        scheduled = GREATEST(0, wf.scheduled - COALESCE(d.scheduled, 0)),
        executing = GREATEST(0, wf.executing - COALESCE(d.executing, 0)),
        retryable = GREATEST(0, wf.retryable - COALESCE(d.retryable, 0)),
        completed = GREATEST(0, wf.completed - COALESCE(d.completed, 0)),
        cancelled = GREATEST(0, wf.cancelled - COALESCE(d.cancelled, 0)),
        discarded = GREATEST(0, wf.discarded - COALESCE(d.discarded, 0))
      FROM (
        SELECT
          o.meta->>'workflow_id' AS workflow_id,
          COUNT(*) FILTER (WHERE o.state = 'suspended') AS suspended,
          COUNT(*) FILTER (WHERE o.state = 'available') AS available,
          COUNT(*) FILTER (WHERE o.state = 'scheduled') AS scheduled,
          COUNT(*) FILTER (WHERE o.state = 'executing') AS executing,
          COUNT(*) FILTER (WHERE o.state = 'retryable') AS retryable,
          COUNT(*) FILTER (WHERE o.state = 'completed') AS completed,
          COUNT(*) FILTER (WHERE o.state = 'cancelled') AS cancelled,
          COUNT(*) FILTER (WHERE o.state = 'discarded') AS discarded
        FROM old_table o
        WHERE o.meta ? 'workflow_id'
        GROUP BY o.meta->>'workflow_id'
      ) d
      WHERE wf.id = d.workflow_id;

      RETURN NULL;
    END;
    $$ LANGUAGE plpgsql;
    """

    execute """
    CREATE OR REPLACE TRIGGER oban_workflows_update_trigger
    AFTER UPDATE ON #{opts.quoted}.oban_jobs
    REFERENCING OLD TABLE AS old_table NEW TABLE AS new_table
    FOR EACH STATEMENT
    EXECUTE FUNCTION #{opts.quoted}.oban_workflows_update();
    """

    execute """
    CREATE OR REPLACE TRIGGER oban_workflows_delete_trigger
    AFTER DELETE ON #{opts.quoted}.oban_jobs
    REFERENCING OLD TABLE AS old_table
    FOR EACH STATEMENT
    EXECUTE FUNCTION #{opts.quoted}.oban_workflows_delete();
    """
  end

  def up_idx(opts) do
    # Replace column-based indexes with expression indexes for unique and partition keys. This
    # enables querying meta fields directly without relying on generated columns. Partitioned
    # tables already use expression indexes from v1.5.0.
    if not opts.generated_columns and not partitioned?(opts) do
      drop_if_exists index(:oban_jobs, [],
                       name: :oban_jobs_unique_index,
                       concurrently: opts.concurrently,
                       prefix: opts.prefix
                     )

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

      create_if_not_exists index(
                             :oban_jobs,
                             ["(meta->>'uniq_key')"],
                             concurrently: opts.concurrently,
                             name: :oban_jobs_unique_index,
                             prefix: opts.prefix,
                             unique: true,
                             where:
                               "meta ? 'uniq_key' AND jsonb_contains(meta->'uniq_bmp', #{opts.quoted}.oban_state_to_bit(state))"
                           )

      create_if_not_exists index(
                             :oban_jobs,
                             ["(meta->>'partition_key')", :queue, :priority, :scheduled_at, :id],
                             concurrently: opts.concurrently,
                             name: :oban_jobs_partition_index,
                             prefix: opts.prefix,
                             where: "state = 'available' AND meta ? 'partition_key'"
                           )
    end

    # Rename old workflow/chain indexes that use on_hold
    execute "ALTER INDEX IF EXISTS #{opts.quoted}.oban_jobs_workflow_index RENAME TO oban_jobs_workflow_index_old"

    execute "ALTER INDEX IF EXISTS #{opts.quoted}.oban_jobs_sup_workflow_index RENAME TO oban_jobs_sup_workflow_index_old"

    execute "ALTER INDEX IF EXISTS #{opts.quoted}.oban_jobs_chain_index RENAME TO oban_jobs_chain_index_old"

    # Create new indexes without on_hold
    create_if_not_exists index(
                           :oban_jobs,
                           ["(meta->>'workflow_id')", :state, "(meta->>'name')"],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_workflow_index,
                           prefix: opts.prefix,
                           where: "meta ? 'workflow_id'"
                         )

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

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

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

    # Replace the non-partial state timestamp indexes from core Oban with partial indexes.
    # These partial indexes are smaller and more efficient for pruning operations.
    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_state_cancelled_at_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

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

    create_if_not_exists index(
                           :oban_jobs,
                           [:cancelled_at],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_cancelled_at_index,
                           prefix: opts.prefix,
                           where: "state = 'cancelled'"
                         )

    create_if_not_exists index(
                           :oban_jobs,
                           [:discarded_at],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_discarded_at_index,
                           prefix: opts.prefix,
                           where: "state = 'discarded'"
                         )

    create_if_not_exists index(
                           :oban_jobs,
                           [:completed_at],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_completed_at_index,
                           prefix: opts.prefix,
                           where: "state = 'completed'"
                         )

    # Add a staging index for efficiently finding scheduled/retryable jobs ready to run.
    # This speeds up the stage_jobs operation which transitions jobs to 'available'.
    create_if_not_exists index(
                           :oban_jobs,
                           [:scheduled_at, :id],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_staging_index,
                           prefix: opts.prefix,
                           where: "state IN ('scheduled', 'retryable')"
                         )

    # Workflow table indexes
    create_if_not_exists index(
                           :oban_workflows,
                           [:state, :name],
                           concurrently: opts.concurrently,
                           prefix: opts.prefix
                         )

    create_if_not_exists index(
                           :oban_workflows,
                           [:parent_id],
                           concurrently: opts.concurrently,
                           prefix: opts.prefix,
                           where: "parent_id IS NOT NULL"
                         )

    create_if_not_exists index(
                           :oban_workflows,
                           [:name],
                           concurrently: opts.concurrently,
                           name: :oban_workflows_unique_name,
                           prefix: opts.prefix,
                           unique: true,
                           where: "state != 'completed' AND meta ? 'unique'"
                         )
  end

  def down_ddl(opts) do
    execute "DROP TRIGGER IF EXISTS oban_workflows_update_trigger ON #{opts.quoted}.oban_jobs"
    execute "DROP TRIGGER IF EXISTS oban_workflows_delete_trigger ON #{opts.quoted}.oban_jobs"
    execute "DROP FUNCTION IF EXISTS #{opts.quoted}.oban_workflows_update()"
    execute "DROP FUNCTION IF EXISTS #{opts.quoted}.oban_workflows_delete()"
    drop_if_exists table(:oban_workflows, prefix: opts.prefix)
  end

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

    drop_if_exists index(:oban_workflows, [:parent_id],
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

    drop_if_exists index(:oban_workflows, [],
                     name: :oban_workflows_unique_name,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

    # Drop new indexes
    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_workflow_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_sup_workflow_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_chunk_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

    # Drop the partial state timestamp indexes and restore the original non-partial indexes
    drop_if_exists index(:oban_jobs, [],
                     name: :oban_jobs_staging_index,
                     concurrently: opts.concurrently,
                     prefix: opts.prefix
                   )

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

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

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

    create_if_not_exists index(
                           :oban_jobs,
                           [:state, :cancelled_at],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_state_cancelled_at_index,
                           prefix: opts.prefix
                         )

    create_if_not_exists index(
                           :oban_jobs,
                           [:state, :discarded_at],
                           concurrently: opts.concurrently,
                           name: :oban_jobs_state_discarded_at_index,
                           prefix: opts.prefix
                         )

    # Rename old indexes back
    execute "ALTER INDEX IF EXISTS #{opts.quoted}.oban_jobs_workflow_index_old RENAME TO oban_jobs_workflow_index"

    execute "ALTER INDEX IF EXISTS #{opts.quoted}.oban_jobs_sup_workflow_index_old RENAME TO oban_jobs_sup_workflow_index"

    execute "ALTER INDEX IF EXISTS #{opts.quoted}.oban_jobs_chain_index_old RENAME TO oban_jobs_chain_index"
  end
end
