-- 0.18.3 hotfix: split user_keys upload into a non-destructive `upsert_user_key`
-- and the existing destructive `reset_user_key`. Add `rotate_conv_key` so the
-- client can escape "Awaiting key" deadlocks by minting a fresh per-conversation
-- key and wrapping it for every member (per-user where possible, per-device as
-- a legacy fallback for peers still on 0.17.x).
--
-- Why: `reset_user_key` was being called from EVERY upload path
-- (setupNewUserIdentity, changePin, regenerateRecoveryCode), wiping every
-- legacy `conversation_keys` row for the user before the migration could
-- re-wrap them. Users ended up with `user_keys` set, zero un-migrated
-- bundles, and no way to send or read.

-- 1) upsert_user_key — same UPSERT as reset_user_key but WITHOUT the DELETE.
--    Safe to call on every PIN-set / PIN-change / recovery-regen.

create or replace function public.upsert_user_key(
  p_user_id             uuid,
  p_public_key_b64      text,
  p_sealed_private_b64  text,
  p_salt_b64            text,
  p_kdf_params          jsonb,
  p_recovery_sealed_b64 text default null,
  p_recovery_salt_b64   text default null
) returns void
language plpgsql
security definer
set search_path = public
as $$
declare
  caller uuid := auth.uid();
begin
  if caller is null or caller <> p_user_id then
    raise exception 'not authenticated as %', p_user_id;
  end if;

  insert into public.user_keys (
    user_id, public_key, sealed_private_key, salt, kdf_params,
    recovery_sealed_private_key, recovery_salt,
    failed_attempts, locked_until,
    failed_recovery_attempts, recovery_locked_until,
    key_version, created_at, updated_at
  ) values (
    p_user_id,
    decode(p_public_key_b64, 'base64'),
    decode(p_sealed_private_b64, 'base64'),
    decode(p_salt_b64, 'base64'),
    p_kdf_params,
    case when p_recovery_sealed_b64 is null then null else decode(p_recovery_sealed_b64, 'base64') end,
    case when p_recovery_salt_b64   is null then null else decode(p_recovery_salt_b64,   'base64') end,
    0, null, 0, null,
    1, now(), now()
  )
  on conflict (user_id) do update set
    public_key                  = excluded.public_key,
    sealed_private_key          = excluded.sealed_private_key,
    salt                        = excluded.salt,
    kdf_params                  = excluded.kdf_params,
    recovery_sealed_private_key = excluded.recovery_sealed_private_key,
    recovery_salt               = excluded.recovery_salt,
    failed_attempts             = 0,
    locked_until                = null,
    failed_recovery_attempts    = 0,
    recovery_locked_until       = null,
    -- Don't bump key_version here — the public key is unchanged.
    updated_at                  = now();
end;
$$;

revoke execute on function public.upsert_user_key(uuid, text, text, text, jsonb, text, text) from public, anon;
grant  execute on function public.upsert_user_key(uuid, text, text, text, jsonb, text, text) to authenticated;

-- 2) rotate_conv_key — atomically bumps active_key_version and inserts a new
--    set of bundles. Each bundle may carry recipient_user_id (per-user wrap)
--    OR recipient_device_id (per-device fallback for peers on the legacy
--    client). Caller must ensure the new version is strictly greater than
--    the current one (we lock the row to prevent races).

create or replace function public.rotate_conv_key(
  p_conv_id        uuid,
  p_sender_user_id uuid,
  p_new_version    int,
  p_bundles        jsonb
) returns int
language plpgsql
security definer
set search_path = public
as $$
declare
  caller         uuid := auth.uid();
  cur_version    int;
  bundle         jsonb;
  inserted       int := 0;
  recipient_uid  uuid;
  recipient_did  uuid;
  member_user_id uuid;
  enc_key_hex    text;
  nonce_hex      text;
begin
  if caller is null or caller <> p_sender_user_id then
    raise exception 'not authenticated as %', p_sender_user_id;
  end if;

  if not exists (
    select 1 from public.conversation_members
     where conversation_id = p_conv_id
       and user_id         = caller
       and accepted        = true
  ) then
    raise exception 'caller is not an accepted member of %', p_conv_id;
  end if;

  -- Lock the conversation row so concurrent rotations don't race the version bump.
  select active_key_version into cur_version
    from public.conversations
   where id = p_conv_id
   for update;
  if cur_version is null then
    raise exception 'conversation % not found', p_conv_id;
  end if;
  if p_new_version <= cur_version then
    raise exception 'new key version % must be greater than current %',
      p_new_version, cur_version;
  end if;

  update public.conversations
     set active_key_version = p_new_version
   where id = p_conv_id;

  -- Insert each bundle. We don't auto-derive recipient_user_id from the
  -- device anymore — for per-device fallback rows the column stays NULL so
  -- multiple devices of the same user can each get their own bundle.
  for bundle in select * from jsonb_array_elements(p_bundles) loop
    recipient_uid := nullif(bundle->>'recipient_user_id', '')::uuid;
    recipient_did := nullif(bundle->>'recipient_device_id', '')::uuid;
    enc_key_hex   := bundle->>'encrypted_key';
    nonce_hex     := bundle->>'nonce';

    -- Validate membership regardless of mode.
    if recipient_uid is not null then
      member_user_id := recipient_uid;
    elsif recipient_did is not null then
      select user_id into member_user_id from public.devices where id = recipient_did;
    else
      continue;
    end if;
    if member_user_id is null then continue; end if;
    if not exists (
      select 1 from public.conversation_members
       where conversation_id = p_conv_id
         and user_id         = member_user_id
         and accepted        = true
    ) then continue; end if;

    insert into public.conversation_keys
      (conversation_id, recipient_user_id, recipient_device_id,
       key_version, sender_user_id, sender_device_id,
       encrypted_key, nonce)
    values
      (p_conv_id, recipient_uid, recipient_did,
       p_new_version, p_sender_user_id, null,
       decode(enc_key_hex, 'hex'),
       decode(nonce_hex, 'hex'))
    on conflict do nothing;

    if found then inserted := inserted + 1; end if;
  end loop;

  return inserted;
end;
$$;

revoke execute on function public.rotate_conv_key(uuid, uuid, int, jsonb) from public, anon;
grant  execute on function public.rotate_conv_key(uuid, uuid, int, jsonb) to authenticated;