-- Recoverable generation finalization and charge-before-dispatch fencing.
-- Additive so the previous release remains rollback-compatible.

alter table public.generation_jobs
  add column if not exists dispatch_ready_at timestamptz;

alter table public.generation_jobs
  add column if not exists finalized_at timestamptz;

alter table public.generation_jobs
  add column if not exists provider_dispatch_started_at timestamptz;

-- Existing rows were created before dispatch fencing and are therefore ready.
update public.generation_jobs
set dispatch_ready_at = coalesce(dispatch_ready_at, created_at, now())
where dispatch_ready_at is null;

-- Rows already terminal before this migration have completed their legacy
-- finalization path and must not be replayed into billing/webhook side effects.
update public.generation_jobs
set finalized_at = coalesce(finalized_at, completed_at, updated_at, now())
where finalized_at is null
  and status in ('succeeded', 'failed', 'expired', 'cancelled');

create index if not exists generation_jobs_finalize_claim_idx
  on public.generation_jobs(finalized_at, scheduled_at, locked_at, priority desc);

create or replace function public.claim_generation_jobs(
  p_worker_id text,
  p_limit integer default 1,
  p_lock_timeout_seconds integer default 300
)
returns setof public.generation_jobs
language plpgsql
set search_path = public, pg_temp
as $$
declare
  v_now timestamptz := now();
begin
  return query
  with candidates as (
    select id
    from public.generation_jobs
    where dispatch_ready_at is not null
      and finalized_at is null
      and status in ('queued', 'running', 'succeeded', 'failed', 'expired', 'cancelled')
      and coalesce(scheduled_at, created_at) <= v_now
      and (
        locked_at is null
        or locked_at < v_now - make_interval(secs => p_lock_timeout_seconds)
      )
    order by coalesce(priority, 0) desc, coalesce(scheduled_at, created_at) asc, created_at asc
    limit greatest(1, least(coalesce(p_limit, 1), 20))
    for update skip locked
  ),
  updated as (
    update public.generation_jobs
    set locked_at = v_now,
        locked_by = p_worker_id,
        started_at = coalesce(public.generation_jobs.started_at, v_now),
        updated_at = v_now
    where id in (select id from candidates)
    returning public.generation_jobs.*
  )
  select * from updated;
end;
$$;

revoke all on function public.claim_generation_jobs(text, integer, integer) from public;