anydrop/web/src/hooks/useSignaling.ts

import { useEffect, useRef, useCallback } from "react";
import type {
  ServerMessage,
  TransferRequestMessage,
  TransferResponseMessage,
  TransferResumeMessage,
} from "@anydrop/shared";
import { SignalingClient } from "../lib/signaling";
import { PeerManager } from "../lib/peerManager";
import {
  sendFile,
  createTransferRequest,
  createTransferResponse,
  createFileReceiver,
} from "../lib/fileTransfer";
import { setupPushNotifications, showLocalNotification } from "../lib/notifications";
import { detectLocalIP } from "../lib/localIP";
import {
  saveTransferMeta,
  updateTransferBytes,
  cacheSendFile,
  getCachedSendFile,
  getAllPendingTransfers,
  getReceivedData,
  appendReceivedChunk,
  cleanupTransfer,
  cleanupOldTransfers,
} from "../lib/transferStore";
import { useStore } from "../stores/useStore";
import { useProfileStore } from "../stores/useProfileStore";

function downloadBlob(blob: Blob, fileName: string) {
  const url = URL.createObjectURL(blob);
  const a = document.createElement("a");
  a.href = url;
  a.download = fileName;
  document.body.appendChild(a);
  a.click();
  document.body.removeChild(a);
  URL.revokeObjectURL(url);
}

export function useSignaling(joinCode?: string) {
  const signalingRef = useRef<SignalingClient | null>(null);
  const peerManagerRef = useRef<PeerManager | null>(null);
  const pendingFilesRef = useRef<Map<string, { files: File[] }>>(new Map());
  const fileReceiverRef = useRef<ReturnType<typeof createFileReceiver> | null>(null);

  const {
    setConnection,
    setConnected,
    addPeer,
    removePeer,
    setPeers,
    setPublicRoom,
    addTransfer,
    updateTransfer,
    setIncomingRequest,
    setError,
  } = useStore();

  // ── beforeunload: warn during active transfers ──
  useEffect(() => {
    const handler = (e: BeforeUnloadEvent) => {
      const { transfers } = useStore.getState();
      const active = transfers.some(
        (t) => t.status === "transferring" || t.status === "pending",
      );
      if (active) {
        e.preventDefault();
      }
    };
    window.addEventListener("beforeunload", handler);
    return () => window.removeEventListener("beforeunload", handler);
  }, []);

  // ── Cleanup old transfers on mount ──
  useEffect(() => {
    cleanupOldTransfers();
  }, []);

  useEffect(() => {
    fileReceiverRef.current = createFileReceiver({
      onData: () => {},
      onComplete: (fileId, blob, fileName) => {
        updateTransfer(fileId, { progress: 1, status: "done" });
        cleanupTransfer(fileId);
        downloadBlob(blob, fileName);
      },
      onProgress: (fileId, progress) => {
        updateTransfer(fileId, { progress, status: "transferring" });
      },
      onChunkReceived: (fileId, chunk, totalReceived) => {
        // Persist to IndexedDB for resume after refresh (fire-and-forget)
        appendReceivedChunk(fileId, chunk).catch(() => {});
        updateTransferBytes(fileId, totalReceived).catch(() => {});
      },
      onTransferRequest: (msg: TransferRequestMessage) => {
        setIncomingRequest({
          peerId: "",
          displayName: "",
          files: msg.files,
          text: msg.text,
        });
      },
      onTransferResponse: (msg: TransferResponseMessage) => {
        if (msg.accepted) {
          startSending();
        } else {
          const store = useStore.getState();
          for (const transfer of store.transfers) {
            if (transfer.direction === "send" && transfer.status === "pending") {
              updateTransfer(transfer.id, { status: "error" });
              cleanupTransfer(transfer.id);
            }
          }
        }
      },
      onTransferResume: (msg: TransferResumeMessage) => {
        // Peer is asking us to resume sending from an offset
        handleResumeRequest(msg);
      },
      onText: (text) => {
        setIncomingRequest({
          peerId: "",
          displayName: "",
          files: [],
          text,
        });
      },
    });
  }, []);

  const handleResumeRequest = useCallback(
    async (msg: TransferResumeMessage) => {
      const pm = peerManagerRef.current;
      if (!pm) return;

      for (const { id, bytesReceived } of msg.files) {
        const cachedFile = await getCachedSendFile(id);
        if (!cachedFile) continue;

        // Find which peer to send to from the stored metadata
        const meta = await import("../lib/transferStore").then((m) =>
          m.getTransferMeta(id),
        );
        if (!meta) continue;

        const peer = pm.getPeer(meta.peerId);
        if (!peer) continue;

        updateTransfer(id, {
          progress: bytesReceived / meta.fileSize,
          status: "transferring",
        });

        sendFile(peer, cachedFile, id, (progress) => {
          updateTransfer(id, {
            progress,
            status: progress >= 1 ? "done" : "transferring",
          });
          if (progress >= 1) {
            cleanupTransfer(id);
          }
        }, bytesReceived);
      }
    },
    [updateTransfer],
  );

  const startSending = useCallback(() => {
    const store = useStore.getState();
    for (const [peerId, { files }] of pendingFilesRef.current) {
      const pm = peerManagerRef.current;
      if (!pm) continue;
      const peer = pm.getPeer(peerId);
      if (!peer) continue;

      for (const file of files) {
        const transfer = store.transfers.find(
          (t) =>
            t.direction === "send" &&
            t.fileName === file.name &&
            t.status === "pending",
        );
        if (transfer) {
          sendFile(peer, file, transfer.id, (progress) => {
            updateTransfer(transfer.id, {
              progress,
              status: progress >= 1 ? "done" : "transferring",
            });
            if (progress >= 1) {
              cleanupTransfer(transfer.id);
            }
          });
        }
      }
    }
    pendingFilesRef.current.clear();
  }, [updateTransfer]);

  useEffect(() => {
    let cancelled = false;
    const profile = useProfileStore.getState();

    const handleMessage = (msg: ServerMessage) => {
      switch (msg.type) {
        case "welcome":
          setConnection(msg.peerId, msg.roomId);
          setPeers(msg.peers);

          if (msg.vapidPublicKey && signalingRef.current) {
            setupPushNotifications(
              msg.vapidPublicKey,
              profile.deviceId,
              signalingRef.current,
            );
          }

          // Check for pending transfers to resume
          checkPendingResumes();
          break;
        case "peer-joined":
          addPeer({
            peerId: msg.peerId,
            displayName: msg.displayName,
            deviceType: msg.deviceType,
            avatar: msg.avatar,
          });
          showLocalNotification(
            "AnyDrop",
            `${msg.displayName} est à proximité`,
          );
          break;
        case "peer-left":
          removePeer(msg.peerId);
          peerManagerRef.current?.closePeer(msg.peerId);
          break;
        case "signal":
          peerManagerRef.current?.handleSignal(msg.from, msg.data);
          break;
        case "public-room-created":
          setPublicRoom(msg.code, msg.url, msg.expiresAt);
          break;
        case "error":
          setError(msg.message);
          break;
      }
    };

    const checkPendingResumes = async () => {
      const pending = await getAllPendingTransfers();
      if (pending.length === 0) return;

      for (const meta of pending) {
        // Restore transfer in UI
        addTransfer({
          id: meta.id,
          peerId: meta.peerId,
          fileName: meta.fileName,
          fileSize: meta.fileSize,
          mime: meta.mime,
          direction: meta.direction,
          progress: meta.bytesTransferred / meta.fileSize,
          status: "pending",
        });

        if (meta.direction === "receive") {
          // We were receiving — need to tell the sender to resume
          const existingBlob = await getReceivedData(meta.id);
          const bytesReceived = existingBlob?.size ?? 0;

          if (bytesReceived > 0) {
            // Restore partial data in file receiver
            fileReceiverRef.current?.restoreReceiving(
              meta.id,
              meta.fileName,
              meta.fileSize,
              meta.mime,
              bytesReceived,
              existingBlob!,
            );
          }

          // When we reconnect with this peer, send a resume request
          const pm = peerManagerRef.current;
          if (pm) {
            const sendResumeWhenReady = () => {
              const peer = pm.getPeer(meta.peerId);
              if (peer && (peer as any).connected) {
                const resumeMsg: TransferResumeMessage = {
                  type: "transfer-resume",
                  files: [{ id: meta.id, bytesReceived }],
                };
                peer.send(JSON.stringify(resumeMsg));
              } else {
                setTimeout(sendResumeWhenReady, 500);
              }
            };
            // Delay to let peer connection establish
            setTimeout(sendResumeWhenReady, 2000);
          }
        }
      }
    };

    const init = async () => {
      const localIP = await detectLocalIP().catch(() => undefined);
      if (cancelled) return;

      const signaling = new SignalingClient(
        handleMessage,
        {
          deviceId: profile.deviceId,
          deviceName: profile.deviceName,
          deviceType: profile.deviceType,
          avatar: profile.avatar || undefined,
          localIP,
          groupId: profile.groupId || undefined,
        },
        joinCode,
      );
      signalingRef.current = signaling;

      const peerManager = new PeerManager(signaling, {
        onConnect: (peerId) => {
          console.log(`P2P connected with ${peerId}`);
        },
        onClose: (peerId) => {
          console.log(`P2P disconnected from ${peerId}`);
        },
        onData: (peerId, data) => {
          const store = useStore.getState();
          const peerInfo = store.peers.find((p) => p.peerId === peerId);

          if (typeof data === "string") {
            try {
              const msg = JSON.parse(data);
              if (msg.type === "transfer-request") {
                setIncomingRequest({
                  peerId,
                  displayName: peerInfo?.displayName || peerId,
                  files: msg.files,
                  text: msg.text,
                });
                const fileCount = msg.files?.length || 0;
                const label =
                  fileCount > 1
                    ? `${fileCount} fichiers`
                    : msg.files?.[0]?.name || "un fichier";
                showLocalNotification(
                  "Transfert entrant",
                  `${peerInfo?.displayName || "Quelqu'un"} veut vous envoyer ${label}`,
                );
                return;
              }
              if (msg.type === "transfer-response") {
                if (msg.accepted) {
                  startSending();
                }
                return;
              }
              if (msg.type === "transfer-resume") {
                handleResumeRequest(msg);
                return;
              }
              if (msg.type === "text") {
                setIncomingRequest({
                  peerId,
                  displayName: peerInfo?.displayName || peerId,
                  files: [],
                  text: msg.content,
                });
                showLocalNotification(
                  "Texte reçu",
                  `${peerInfo?.displayName || "Quelqu'un"} vous a envoyé du texte`,
                );
                return;
              }
            } catch {
              // Not JSON
            }
          }
          fileReceiverRef.current?.handleData(data);
        },
        onError: (peerId, err) => {
          console.error(`P2P error with ${peerId}:`, err);
        },
      });
      peerManagerRef.current = peerManager;

      signaling.connect();
      setConnected(true);
    };

    init();

    return () => {
      cancelled = true;
      signalingRef.current?.disconnect();
      peerManagerRef.current?.closeAll();
      setConnected(false);
    };
  }, [joinCode]);

  const sendFiles = useCallback(
    async (peerId: string, files: File[]) => {
      const pm = peerManagerRef.current;
      const signaling = signalingRef.current;
      if (!pm || !signaling) return;

      let peer = pm.getPeer(peerId);
      if (!peer) {
        peer = pm.createPeer(peerId, true);
      }

      const request = createTransferRequest(files);
      const store = useStore.getState();
      const peerInfo = store.peers.find((p) => p.peerId === peerId);

      for (let i = 0; i < request.files.length; i++) {
        const fileMeta = request.files[i];
        addTransfer({
          id: fileMeta.id,
          peerId,
          fileName: fileMeta.name,
          fileSize: fileMeta.size,
          mime: fileMeta.mime,
          direction: "send",
          progress: 0,
          status: "pending",
        });

        // Persist file and metadata for resume
        await cacheSendFile(fileMeta.id, files[i]);
        await saveTransferMeta({
          id: fileMeta.id,
          peerId,
          peerDisplayName: peerInfo?.displayName || peerId,
          fileName: fileMeta.name,
          fileSize: fileMeta.size,
          mime: fileMeta.mime,
          direction: "send",
          bytesTransferred: 0,
          createdAt: Date.now(),
        });
      }

      pendingFilesRef.current.set(peerId, { files });

      const sendRequest = () => {
        const p = pm.getPeer(peerId);
        if (p && (p as any)._channel?.readyState === "open") {
          p.send(JSON.stringify(request));
        } else {
          setTimeout(sendRequest, 100);
        }
      };

      if (peer && (peer as any).connected) {
        sendRequest();
      } else {
        peer.on("connect", sendRequest);
      }
    },
    [addTransfer],
  );

  const sendText = useCallback((peerId: string, text: string) => {
    const pm = peerManagerRef.current;
    if (!pm) return;

    let peer = pm.getPeer(peerId);
    if (!peer) {
      peer = pm.createPeer(peerId, true);
    }

    const sendMsg = () => {
      const p = pm.getPeer(peerId);
      if (p) {
        p.send(
          JSON.stringify({ type: "text", id: `t-${Date.now()}`, content: text }),
        );
      }
    };

    if (peer && (peer as any).connected) {
      sendMsg();
    } else {
      peer.on("connect", sendMsg);
    }
  }, []);

  const acceptTransfer = useCallback(
    async (peerId: string) => {
      const pm = peerManagerRef.current;
      if (!pm) return;

      const peer = pm.getPeer(peerId);
      if (!peer) return;

      const store = useStore.getState();
      const request = store.incomingRequest;
      if (request) {
        const peerInfo = store.peers.find((p) => p.peerId === peerId);
        for (const file of request.files) {
          addTransfer({
            id: file.id,
            peerId,
            fileName: file.name,
            fileSize: file.size,
            mime: file.mime,
            direction: "receive",
            progress: 0,
            status: "pending",
          });

          // Persist receive metadata for resume
          await saveTransferMeta({
            id: file.id,
            peerId,
            peerDisplayName: peerInfo?.displayName || peerId,
            fileName: file.name,
            fileSize: file.size,
            mime: file.mime,
            direction: "receive",
            bytesTransferred: 0,
            createdAt: Date.now(),
          });
        }
      }

      peer.send(JSON.stringify(createTransferResponse(true)));
      setIncomingRequest(null);
    },
    [addTransfer, setIncomingRequest],
  );

  const rejectTransfer = useCallback(
    (peerId: string) => {
      const pm = peerManagerRef.current;
      if (!pm) return;

      const peer = pm.getPeer(peerId);
      if (peer) {
        peer.send(JSON.stringify(createTransferResponse(false)));
      }
      setIncomingRequest(null);
    },
    [setIncomingRequest],
  );

  const createPublicRoom = useCallback(() => {
    signalingRef.current?.send({ type: "create-public-room" });
  }, []);

  const wakePeer = useCallback((deviceId: string) => {
    signalingRef.current?.send({ type: "wake-peer", deviceId });
  }, []);

  return {
    sendFiles,
    sendText,
    acceptTransfer,
    rejectTransfer,
    createPublicRoom,
    wakePeer,
  };
}