ChatApp/apps/desktop/src/lib/useConversationMessages.ts

import { loadDevicePrivateKey } from '@chat-app/shared/auth';
import {
  type AttachmentHandle,
  type DecryptedMessage,
  decryptMessages,
  encryptAndUploadAttachment,
  fetchConversationMessages,
  insertAttachmentRow,
  MAX_ATTACHMENT_BYTES,
  type MessageWithCipher,
  sendEncryptedMessage,
} from '@chat-app/shared/chat';
import { bytesToPgHex, pgBytesToBytes } from '@chat-app/shared/supabase';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';

import {
  enqueueOutbox,
  getOutbox,
  isDue,
  markAttempt,
  type OutboxItem,
  removeOutbox,
  shouldGiveUp,
  subscribeOutbox,
} from './messageOutbox';
import { devLocalSecretStore } from './secretStore';
import { supabase } from './supabase';

interface State {
  messages: DecryptedMessage[];
  loading: boolean;
  error: string | null;
}

interface Args {
  conversationId: string | undefined;
  userId: string | undefined;
  deviceId: string | undefined;
}

type MessageChangePayload = {
  eventType: 'INSERT' | 'UPDATE' | 'DELETE';
  new: Record<string, unknown>;
  old: Record<string, unknown>;
};

function rowToMessage(row: Record<string, unknown>): MessageWithCipher {
  return {
    id: String(row.id),
    conversationId: String(row.conversation_id),
    senderId: String(row.sender_id),
    senderDeviceId: row.sender_device_id ? String(row.sender_device_id) : null,
    replyToId: row.reply_to_id ? String(row.reply_to_id) : null,
    editedAt: row.edited_at ? String(row.edited_at) : null,
    deletedAt: row.deleted_at ? String(row.deleted_at) : null,
    createdAt: String(row.created_at),
    ciphertext: pgBytesToBytes(String(row.ciphertext ?? '\\x')),
    nonce: pgBytesToBytes(String(row.nonce ?? '\\x')),
    keyVersion: typeof row.key_version === 'number' ? row.key_version : 1,
  };
}

export function useConversationMessages({ conversationId, userId, deviceId }: Args): State & {
  send: (text: string, images?: File[], replyToId?: string | null) => Promise<void>;
  refresh: () => Promise<void>;
  pending: OutboxItem[];
  retryPending: (id: string) => void;
  cancelPending: (id: string) => void;
} {
  const [state, setState] = useState<State>({ messages: [], loading: true, error: null });
  const [pending, setPending] = useState<OutboxItem[]>(() =>
    conversationId ? getOutbox(conversationId) : [],
  );
  const privateKeyRef = useRef<Uint8Array | null>(null);

  useEffect(() => {
    if (!conversationId) {
      setPending([]);
      return;
    }
    return subscribeOutbox((byConv) => {
      setPending(byConv[conversationId] ?? []);
    });
  }, [conversationId]);

  // Load own private key once per (user, device).
  useEffect(() => {
    privateKeyRef.current = null;
    if (!userId || !deviceId) return;
    void loadDevicePrivateKey(devLocalSecretStore, userId, deviceId).then((pk) => {
      privateKeyRef.current = pk;
    });
  }, [userId, deviceId]);

  const decryptBatch = useCallback(
    async (messages: MessageWithCipher[]): Promise<DecryptedMessage[]> => {
      const priv = privateKeyRef.current;
      if (!priv || !deviceId || messages.length === 0) {
        return messages.map((m) => ({ ...m, plaintext: null }));
      }
      return decryptMessages({
        client: supabase,
        messages,
        ownDeviceId: deviceId,
        ownPrivateKey: priv,
      });
    },
    [deviceId],
  );

  const refresh = useCallback(async () => {
    if (!conversationId) return;
    try {
      // Only show the loading spinner on the *initial* fetch — subsequent
      // refreshes (focus/visibility) replace messages in-place to avoid
      // flickering an empty state on every wake.
      setState((prev) =>
        prev.messages.length === 0 ? { ...prev, loading: true } : prev,
      );
      const rows = await fetchConversationMessages(supabase, conversationId);
      const decrypted = await decryptBatch(rows);
      setState({ messages: decrypted, loading: false, error: null });
    } catch (err: unknown) {
      setState((prev) => ({
        ...prev,
        loading: false,
        error: err instanceof Error ? err.message : 'failed to load messages',
      }));
    }
  }, [conversationId, decryptBatch]);

  // Realtime INSERT handler — refetches the row via REST so we get the
  // canonical bytea encoding (postgres_changes payloads serialize bytea
  // differently and decoding them inline is brittle). Then decrypt + append.
  // Skips if the message is already in state (e.g. optimistic insert from our
  // own send), so the sender's cached copy isn't overwritten with a flicker.
  const handleInsert = useCallback(
    async (row: Record<string, unknown>) => {
      if (!conversationId || !deviceId) return;
      const id = String(row.id);

      let alreadyHave = false;
      setState((prev) => {
        if (prev.messages.some((m) => m.id === id)) alreadyHave = true;
        return prev;
      });
      if (alreadyHave) return;

      let decrypted: DecryptedMessage | null = null;
      for (let attempt = 0; attempt < 6; attempt++) {
        const { data, error } = await supabase
          .from('messages')
          .select(
            'id, conversation_id, sender_id, sender_device_id, reply_to_id, edited_at, deleted_at, created_at, ciphertext, nonce, key_version',
          )
          .eq('id', id)
          .maybeSingle();
        if (error) {
          console.warn('handleInsert refetch failed', error);
          return;
        }
        if (!data) {
          await new Promise((r) => window.setTimeout(r, 120 * (attempt + 1)));
          continue;
        }
        // db-types snapshot predates the sender-key columns; cast to bypass.
        const r = data as unknown as {
          id: string;
          conversation_id: string;
          sender_id: string;
          sender_device_id: string | null;
          reply_to_id: string | null;
          edited_at: string | null;
          deleted_at: string | null;
          created_at: string;
          ciphertext: string;
          nonce: string;
          key_version: number;
        };
        const msg: MessageWithCipher = {
          id: r.id,
          conversationId: r.conversation_id,
          senderId: r.sender_id,
          senderDeviceId: r.sender_device_id,
          replyToId: r.reply_to_id,
          editedAt: r.edited_at,
          deletedAt: r.deleted_at,
          createdAt: r.created_at,
          ciphertext: pgBytesToBytes(String(r.ciphertext)),
          nonce: pgBytesToBytes(String(r.nonce)),
          keyVersion: r.key_version,
        };
        const [d] = await decryptBatch([msg]);
        if (d) {
          decrypted = d;
          if (d.plaintext !== null) break;
        }
        await new Promise((r) => window.setTimeout(r, 200 * (attempt + 1)));
      }
      if (!decrypted) return;
      setState((prev) => {
        if (prev.messages.some((m) => m.id === decrypted!.id)) return prev;
        return { ...prev, messages: [...prev.messages, decrypted!] };
      });
    },
    [conversationId, deviceId, decryptBatch],
  );

  const handleUpdate = useCallback(
    async (row: Record<string, unknown>) => {
      const partial = rowToMessage(row);
      setState((prev) => {
        const idx = prev.messages.findIndex((m) => m.id === partial.id);
        if (idx === -1) return prev;
        const existing = prev.messages[idx];
        if (!existing) return prev;
        const next = [...prev.messages];
        next[idx] = {
          ...existing,
          editedAt: partial.editedAt,
          deletedAt: partial.deletedAt,
        };
        return { ...prev, messages: next };
      });
      if (partial.editedAt && !partial.deletedAt) {
        // Realtime bytea encoding varies (base64 vs hex, even null for
        // unchanged columns on some configs). Refetch via REST to get the
        // canonical `\x…` hex then decrypt — same pattern as handleInsert.
        if (!conversationId || !deviceId) return;
        let decrypted: DecryptedMessage | null = null;
        for (let attempt = 0; attempt < 6; attempt++) {
          const { data, error } = await supabase
            .from('messages')
            .select(
              'id, conversation_id, sender_id, sender_device_id, reply_to_id, edited_at, deleted_at, created_at, ciphertext, nonce, key_version',
            )
            .eq('id', partial.id)
            .maybeSingle();
          if (error) {
            console.warn('handleUpdate refetch failed', error);
            return;
          }
          if (!data) {
            await new Promise((r) => window.setTimeout(r, 120 * (attempt + 1)));
            continue;
          }
          const r = data as unknown as {
            id: string;
            conversation_id: string;
            sender_id: string;
            sender_device_id: string | null;
            reply_to_id: string | null;
            edited_at: string | null;
            deleted_at: string | null;
            created_at: string;
            ciphertext: string;
            nonce: string;
            key_version: number;
          };
          const msg: MessageWithCipher = {
            id: r.id,
            conversationId: r.conversation_id,
            senderId: r.sender_id,
            senderDeviceId: r.sender_device_id,
            replyToId: r.reply_to_id,
            editedAt: r.edited_at,
            deletedAt: r.deleted_at,
            createdAt: r.created_at,
            ciphertext: pgBytesToBytes(String(r.ciphertext)),
            nonce: pgBytesToBytes(String(r.nonce)),
            keyVersion: r.key_version,
          };
          const [d] = await decryptBatch([msg]);
          if (d) {
            decrypted = d;
            if (d.plaintext !== null) break;
          }
          await new Promise((r) => window.setTimeout(r, 200 * (attempt + 1)));
        }
        if (!decrypted) return;
        setState((prev) => {
          const idx = prev.messages.findIndex((m) => m.id === decrypted!.id);
          if (idx === -1) return prev;
          const next = [...prev.messages];
          next[idx] = decrypted!;
          return { ...prev, messages: next };
        });
      }
    },
    [conversationId, deviceId, decryptBatch],
  );

  const handleDelete = useCallback((row: Record<string, unknown>) => {
    const id = String(row.id);
    setState((prev) => ({
      ...prev,
      messages: prev.messages.filter((m) => m.id !== id),
    }));
  }, []);

  useEffect(() => {
    if (!conversationId || !userId || !deviceId) return;
    void refresh();

    const channel = supabase
      .channel('conv:' + conversationId)
      .on(
        'postgres_changes',
        {
          event: '*',
          schema: 'public',
          table: 'messages',
          filter: 'conversation_id=eq.' + conversationId,
        },
        (payload: MessageChangePayload) => {
          if (payload.eventType === 'INSERT') {
            void handleInsert(payload.new);
          } else if (payload.eventType === 'UPDATE') {
            void handleUpdate(payload.new);
          } else if (payload.eventType === 'DELETE') {
            handleDelete(payload.old);
          }
        },
      )
      // When a peer device wraps the conversation-key for us (e.g. we just
      // registered a fresh device), re-decrypt the visible messages.
      .on(
        'postgres_changes',
        {
          event: 'INSERT',
          schema: 'public',
          table: 'conversation_keys',
          filter: 'conversation_id=eq.' + conversationId,
        },
        (payload: { new: { recipient_device_id?: string } }) => {
          if (payload.new?.recipient_device_id === deviceId) {
            void refresh();
          }
        },
      )
      .subscribe();

    // Refresh + reconnect on wake from background throttle (mostly Windows
    // WebView2). Without this, messages inserted while the window is
    // minimised never arrive until the user explicitly reloads. Throttled to
    // at most once per AWAKE_THROTTLE_MS so the inevitable cluster of
    // visibility/online events on focus does not cause a flicker storm.
    let lastAwakeRefresh = 0;
    const AWAKE_THROTTLE_MS = 30_000;
    const onAwake = () => {
      if (document.visibilityState !== 'visible') return;
      const now = Date.now();
      if (now - lastAwakeRefresh < AWAKE_THROTTLE_MS) return;
      lastAwakeRefresh = now;
      void refresh();
      try {
        channel.subscribe();
      } catch {
        /* already live */
      }
    };
    document.addEventListener('visibilitychange', onAwake);
    window.addEventListener('online', onAwake);

    return () => {
      document.removeEventListener('visibilitychange', onAwake);
      window.removeEventListener('online', onAwake);
      void supabase.removeChannel(channel);
    };
  }, [conversationId, userId, deviceId, refresh, handleInsert, handleUpdate, handleDelete]);

  const sendText = useCallback(
    async (convId: string, uid: string, did: string, priv: Uint8Array, text: string, replyToId: string | null): Promise<void> => {
      const msg = await sendEncryptedMessage({
        client: supabase,
        conversationId: convId,
        plaintext: text,
        senderUserId: uid,
        senderDeviceId: did,
        senderPrivateKey: priv,
        ...(replyToId ? { replyToId } : {}),
      });
      setState((prev) => {
        if (prev.messages.some((m) => m.id === msg.id)) return prev;
        return {
          ...prev,
          messages: [
            ...prev.messages,
            { ...msg, plaintext: text } as DecryptedMessage,
          ],
        };
      });
    },
    [],
  );

  const send = useCallback(
    async (text: string, images: File[] = [], replyToId: string | null = null) => {
      const trimmed = text.trim();
      if ((!trimmed && images.length === 0) || !conversationId || !userId || !deviceId) return;
      const priv = privateKeyRef.current;
      if (!priv) throw new Error('private key not loaded');

      // Text-only path is retryable — if the network is down or the server
      // rejects transiently, stash in the outbox and keep the UI optimistic.
      // Attachments can't be deferred (large payloads, uploaded separately),
      // so those still surface the error immediately.
      if (images.length === 0) {
        try {
          await sendText(conversationId, userId, deviceId, priv, trimmed, replyToId);
        } catch (err: unknown) {
          const msg = err instanceof Error ? err.message : 'send failed';
          enqueueOutbox({
            conversationId,
            text: trimmed,
            replyToId,
            error: msg,
          });
        }
        return;
      }

      // 1. Upload + encrypt each image. Collect handles + raw blob nonces
      //    (so the public attachment row can reference the blob-level nonce).
      const handles: AttachmentHandle[] = [];
      const blobNonceHexByHandleId = new Map<string, string>();
      for (const file of images) {
        if (file.size > MAX_ATTACHMENT_BYTES) {
          throw new Error('attachment exceeds max size (10 MB)');
        }
        const dims = await readImageDimensions(file);
        const res = await encryptAndUploadAttachment({
          client: supabase,
          conversationId,
          file,
          mimeType: file.type || 'application/octet-stream',
          sizeBytes: file.size,
          ...(dims.width !== undefined ? { width: dims.width } : {}),
          ...(dims.height !== undefined ? { height: dims.height } : {}),
        });
        handles.push(res.handle);
        blobNonceHexByHandleId.set(res.handle.id, bytesToPgHex(res.nonce));
      }

      // 2. Send message (inserts messages + per-conversation key bundles).
      const msg = await sendEncryptedMessage({
        client: supabase,
        conversationId,
        plaintext: trimmed,
        senderUserId: userId,
        senderDeviceId: deviceId,
        senderPrivateKey: priv,
        ...(handles.length > 0 ? { attachmentHandles: handles } : {}),
        ...(replyToId ? { replyToId } : {}),
      });

      // 3. Optimistic insert — we already have the plaintext in hand and the
      //    server returned the row id, so add the message to local state
      //    immediately. Realtime will then no-op (handleInsert dedupes by id).
      const attachmentsPayload =
        handles.length === 0
          ? trimmed
          : JSON.stringify({ v: 1, text: trimmed, attachments: handles });
      setState((prev) => {
        if (prev.messages.some((m) => m.id === msg.id)) return prev;
        return {
          ...prev,
          messages: [
            ...prev.messages,
            {
              ...msg,
              plaintext: attachmentsPayload,
            } as DecryptedMessage,
          ],
        };
      });

      // 4. Insert public attachment metadata rows pointing at the new message.
      for (const h of handles) {
        const blobNonce = blobNonceHexByHandleId.get(h.id) ?? '\\x';
        await insertAttachmentRow(supabase, msg.id, h, blobNonce);
      }
    },
    [conversationId, userId, deviceId, sendText],
  );

  // Drain outbox: retry due items, remove on success, record attempt on fail.
  // Runs on online event, window focus, and a 10s interval.
  useEffect(() => {
    if (!conversationId || !userId || !deviceId) return;

    let draining = false;
    const drain = async (): Promise<void> => {
      if (draining) return;
      const priv = privateKeyRef.current;
      if (!priv) return;
      if (typeof navigator !== 'undefined' && navigator.onLine === false) return;
      draining = true;
      try {
        const items = getOutbox(conversationId).filter(isDue);
        for (const item of items) {
          if (shouldGiveUp(item)) continue;
          try {
            await sendText(
              conversationId,
              userId,
              deviceId,
              priv,
              item.text,
              item.replyToId,
            );
            removeOutbox(conversationId, item.id);
          } catch (err: unknown) {
            markAttempt(
              conversationId,
              item.id,
              err instanceof Error ? err.message : 'send failed',
            );
          }
        }
      } finally {
        draining = false;
      }
    };

    const interval = window.setInterval(() => {
      void drain();
    }, 10_000);
    const onOnline = () => void drain();
    window.addEventListener('online', onOnline);
    window.addEventListener('focus', onOnline);
    // Kick once immediately on mount for stale queued items.
    void drain();

    return () => {
      window.clearInterval(interval);
      window.removeEventListener('online', onOnline);
      window.removeEventListener('focus', onOnline);
    };
  }, [conversationId, userId, deviceId, sendText]);

  const retryPending = useCallback(
    (id: string) => {
      if (!conversationId) return;
      // Force due now; the drain loop will pick it up on next tick.
      markAttempt(conversationId, id, null);
      // Manual kick: mutate nextAttemptAt by re-enqueuing? Simpler — just
      // trigger a drain-ish by queueing a microtask. The interval picks up
      // due items within 10s, but for UX we also eagerly try here.
      const priv = privateKeyRef.current;
      if (!priv || !userId || !deviceId) return;
      const item = getOutbox(conversationId).find((x) => x.id === id);
      if (!item) return;
      void (async () => {
        try {
          await sendText(conversationId, userId, deviceId, priv, item.text, item.replyToId);
          removeOutbox(conversationId, id);
        } catch (err: unknown) {
          markAttempt(
            conversationId,
            id,
            err instanceof Error ? err.message : 'send failed',
          );
        }
      })();
    },
    [conversationId, userId, deviceId, sendText],
  );

  const cancelPending = useCallback(
    (id: string) => {
      if (!conversationId) return;
      removeOutbox(conversationId, id);
    },
    [conversationId],
  );

  return useMemo(
    () => ({ ...state, send, refresh, pending, retryPending, cancelPending }),
    [state, send, refresh, pending, retryPending, cancelPending],
  );
}

// Best-effort image dimension probe. Falls back silently on non-images.
async function readImageDimensions(file: File): Promise<{ width?: number; height?: number }> {
  if (!file.type.startsWith('image/')) return {};
  const url = URL.createObjectURL(file);
  try {
    return await new Promise<{ width?: number; height?: number }>((resolve) => {
      const img = new Image();
      img.onload = () => resolve({ width: img.naturalWidth, height: img.naturalHeight });
      img.onerror = () => resolve({});
      img.src = url;
    });
  } finally {
    URL.revokeObjectURL(url);
  }
}