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 { 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>;
} {
  const [state, setState] = useState<State>({ messages: [], loading: true, error: null });
  const privateKeyRef = useRef<Uint8Array | null>(null);

  // 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 {
      setState((prev) => ({ ...prev, loading: true }));
      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) {
        const [decrypted] = await decryptBatch([partial]);
        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 };
        });
      }
    },
    [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();

    return () => {
      void supabase.removeChannel(channel);
    };
  }, [conversationId, userId, deviceId, refresh, handleInsert, handleUpdate, handleDelete]);

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

      // 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],
  );

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

// 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);
  }
}