-- RPCs for the per-user key flow. SECURITY DEFINER so the lockout counter
-- is server-authoritative even if a malicious client suppresses
-- record_pin_attempt — try_unlock_user_key refuses to deliver ciphertext
-- while a lockout window is active.

-- 1) try_unlock_user_key ---------------------------------------------------

create or replace function public.try_unlock_user_key(p_user_id uuid)
returns jsonb
language plpgsql
security definer
set search_path = public
as $$
declare
  caller uuid := auth.uid();
  row    public.user_keys%rowtype;
begin
  if caller is null or caller <> p_user_id then
    raise exception 'not authenticated as %', p_user_id;
  end if;

  select * into row from public.user_keys where user_id = p_user_id;
  if not found then
    return jsonb_build_object('exists', false);
  end if;

  if row.locked_until is not null and row.locked_until > now() then
    return jsonb_build_object(
      'exists',       true,
      'locked',       true,
      'locked_until', row.locked_until
    );
  end if;

  return jsonb_build_object(
    'exists',                     true,
    'locked',                     false,
    'sealed_private_key',         encode(row.sealed_private_key, 'base64'),
    'salt',                       encode(row.salt, 'base64'),
    'kdf_params',                 row.kdf_params,
    'recovery_sealed_private_key',
      case when row.recovery_sealed_private_key is null
           then null
           else encode(row.recovery_sealed_private_key, 'base64') end,
    'recovery_salt',
      case when row.recovery_salt is null
           then null
           else encode(row.recovery_salt, 'base64') end,
    'failed_attempts',            row.failed_attempts,
    'failed_recovery_attempts',   row.failed_recovery_attempts,
    'recovery_locked_until',      row.recovery_locked_until,
    'key_version',                row.key_version
  );
end;
$$;

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

-- 2) record_pin_attempt ----------------------------------------------------

create or replace function public.record_pin_attempt(
  p_user_id  uuid,
  p_success  boolean,
  p_recovery boolean default false
)
returns jsonb
language plpgsql
security definer
set search_path = public
as $$
declare
  caller    uuid := auth.uid();
  attempts  int;
  cooldown  interval;
  col_attempts text;
  col_locked   text;
  threshold int;
begin
  if caller is null or caller <> p_user_id then
    raise exception 'not authenticated as %', p_user_id;
  end if;

  if p_recovery then
    col_attempts := 'failed_recovery_attempts';
    col_locked   := 'recovery_locked_until';
    threshold    := 20;
  else
    col_attempts := 'failed_attempts';
    col_locked   := 'locked_until';
    threshold    := 10;
  end if;

  if p_success then
    execute format(
      'update public.user_keys set %I = 0, %I = null, updated_at = now() where user_id = $1',
      col_attempts, col_locked
    ) using p_user_id;
    return jsonb_build_object('failed_attempts', 0, 'locked_until', null);
  end if;

  execute format(
    'update public.user_keys set %I = %I + 1, updated_at = now() where user_id = $1 returning %I',
    col_attempts, col_attempts, col_attempts
  ) using p_user_id into attempts;

  cooldown := case
    when attempts < 5      then interval '0 second'
    when attempts = 5      then interval '5 second'
    when attempts = 6      then interval '30 second'
    when attempts = 7      then interval '2 minute'
    when attempts = 8      then interval '10 minute'
    when attempts = 9      then interval '1 hour'
    when attempts >= threshold and not p_recovery
                           then interval '24 hour'
    when attempts >= threshold and p_recovery
                           then interval '100 year'
    else interval '0 second'
  end;

  if cooldown > interval '0 second' then
    execute format(
      'update public.user_keys set %I = now() + $2 where user_id = $1',
      col_locked
    ) using p_user_id, cooldown;
  end if;

  return jsonb_build_object(
    'failed_attempts', attempts,
    'locked_until',
      case when cooldown > interval '0 second' then now() + cooldown else null end
  );
end;
$$;

revoke execute on function public.record_pin_attempt(uuid, boolean, boolean) from public, anon;
grant  execute on function public.record_pin_attempt(uuid, boolean, boolean) to authenticated;

-- 3) share_conv_keys: now accepts recipient_user_id in each bundle --------

drop function if exists public.share_conv_keys(uuid, uuid, int, jsonb);

create or replace function public.share_conv_keys(
  p_conv_id          uuid,
  p_sender_device_id uuid,    -- legacy; nullable if p_sender_user_id supplied
  p_sender_user_id   uuid,    -- new
  p_key_version      int,
  p_bundles          jsonb
) returns int
language plpgsql
security definer
set search_path = public
as $$
declare
  caller        uuid := auth.uid();
  bundle        jsonb;
  inserted      int := 0;
  recipient_uid uuid;
  recipient_did uuid;
  enc_key_hex   text;
  nonce_hex     text;
begin
  if caller is null then raise exception 'not authenticated'; 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;

  if p_sender_user_id is not null and p_sender_user_id <> caller then
    raise exception 'sender_user_id mismatch';
  end if;
  if p_sender_device_id is not null and not exists (
    select 1 from public.devices where id = p_sender_device_id and user_id = caller
  ) then
    raise exception 'sender_device % not owned by caller', p_sender_device_id;
  end if;

  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';

    if recipient_uid is null and recipient_did is not null then
      select user_id into recipient_uid from public.devices where id = recipient_did;
    end if;
    if recipient_uid is null then continue; end if;

    if not exists (
      select 1
        from public.conversation_members
       where conversation_id = p_conv_id
         and user_id         = recipient_uid
         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_key_version, p_sender_user_id, p_sender_device_id,
       decode(enc_key_hex, 'hex'),
       decode(nonce_hex, 'hex'))
    on conflict (conversation_id, recipient_user_id, key_version)
      where recipient_user_id is not null
      do nothing;

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

  return inserted;
end;
$$;

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

-- 4) reset_user_key: hard wipe + replace -----------------------------------

create or replace function public.reset_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 int
language plpgsql
security definer
set search_path = public
as $$
declare
  caller uuid := auth.uid();
  rows_deleted int;
begin
  if caller is null or caller <> p_user_id then
    raise exception 'not authenticated as %', p_user_id;
  end if;

  delete from public.conversation_keys
   where recipient_user_id = p_user_id
      or recipient_device_id in (select id from public.devices where user_id = p_user_id);
  get diagnostics rows_deleted = row_count;

  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,
    key_version                 = public.user_keys.key_version + 1,
    updated_at                  = now();

  return rows_deleted;
end;
$$;

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

-- 5) migrate_user_key_recipients ------------------------------------------

create or replace function public.migrate_user_key_recipients(
  p_conv_id    uuid,
  p_user_id    uuid,
  p_key_version int,
  p_bundles    jsonb -- [{ encrypted_key: hex, nonce: hex, sender_user_id: uuid }]
) returns int
language plpgsql
security definer
set search_path = public
as $$
declare
  caller uuid := auth.uid();
  bundle jsonb;
  inserted int := 0;
begin
  if caller is null or caller <> p_user_id then
    raise exception 'not authenticated as %', p_user_id;
  end if;

  if not exists (
    select 1 from public.conversation_members
     where conversation_id = p_conv_id and user_id = p_user_id and accepted = true
  ) then
    raise exception 'not a member';
  end if;

  for bundle in select * from jsonb_array_elements(p_bundles) loop
    insert into public.conversation_keys (
      conversation_id, recipient_user_id, key_version,
      sender_user_id, encrypted_key, nonce
    ) values (
      p_conv_id, p_user_id, p_key_version,
      nullif(bundle->>'sender_user_id', '')::uuid,
      decode(bundle->>'encrypted_key', 'hex'),
      decode(bundle->>'nonce', 'hex')
    )
    on conflict (conversation_id, recipient_user_id, key_version)
      where recipient_user_id is not null
      do nothing;
    if found then inserted := inserted + 1; end if;
  end loop;
  return inserted;
end;
$$;

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