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, pgHexToBytes } 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: pgHexToBytes(String(row.ciphertext ?? '\\x')),
    nonce: pgHexToBytes(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[]) => 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 — decrypt + append (with retry for envelope race).
  const handleInsert = useCallback(
    async (row: Record<string, unknown>) => {
      if (!deviceId) return;
      const msg = rowToMessage(row);
      let decrypted: DecryptedMessage = { ...msg, plaintext: null };
      for (let attempt = 0; attempt < 6; attempt++) {
        const [d] = await decryptBatch([msg]);
        if (d) {
          decrypted = d;
          if (d.plaintext !== null) break;
        }
        await new Promise((r) => window.setTimeout(r, 120 * (attempt + 1)));
      }
      setState((prev) => {
        if (prev.messages.some((m) => m.id === decrypted.id)) return prev;
        return { ...prev, messages: [...prev.messages, decrypted] };
      });
    },
    [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);
          }
        },
      )
      .subscribe();

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

  const send = useCallback(
    async (text: string, images: File[] = []) => {
      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 + envelopes in one helper).
      const msg = await sendEncryptedMessage({
        client: supabase,
        conversationId,
        plaintext: trimmed,
        senderUserId: userId,
        senderDeviceId: deviceId,
        senderPrivateKey: priv,
        ...(handles.length > 0 ? { attachmentHandles: handles } : {}),
      });

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