siftlode/frontend/src/components/Messages.tsx

import { useEffect, useRef, useState } from "react";
import { useTranslation } from "react-i18next";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { ArrowLeft, Lock, Send, ShieldCheck, SquarePen } from "lucide-react";
import { api, type Conversation, type Message, type MessageUser, type MyKeyBundle } from "../lib/api";
import { useLiveQuery } from "../lib/useLiveQuery";
import { relativeTime } from "../lib/format";
import { onMessage } from "../lib/messagesSocket";
import * as e2ee from "../lib/e2ee";
import Avatar from "./Avatar";

const POLL_MS = 20000; // safety-net poll; live updates come over the WebSocket
const SYSTEM_ID = 0;

// Set-once cache of partner public keys (the server is the directory).
const pubCache = new Map<number, string>();
async function partnerPub(partnerId: number): Promise<string> {
  const cached = pubCache.get(partnerId);
  if (cached) return cached;
  const r = await api.messagePublicKey(partnerId);
  pubCache.set(partnerId, r.public_key);
  return r.public_key;
}

type View = "list" | "new" | { partnerId: number; name?: string; avatar?: string | null; system?: boolean };

export default function Messages({ meId }: { meId: number }) {
  const qc = useQueryClient();
  const [view, setView] = useState<View>("list");
  const [unlocked, setUnlocked] = useState(e2ee.isUnlocked());

  const keyQ = useQuery({ queryKey: ["message-key"], queryFn: api.messageMyKey, staleTime: 60_000 });
  const configured = keyQ.data?.configured ?? false;

  // Restore this device's unlocked key (set once per device, then survives reloads).
  useEffect(() => {
    e2ee.loadFromDevice(meId).then((ok) => ok && setUnlocked(true));
  }, [meId]);

  // Live delivery: a pushed message refreshes the conversation list, unread badge, and the
  // affected thread (which re-decrypts). Near-instant; the poll is only a fallback.
  useEffect(
    () =>
      onMessage((m) => {
        qc.invalidateQueries({ queryKey: ["conversations"] });
        qc.invalidateQueries({ queryKey: ["message-unread"] });
        const partner = m.kind === "system" ? SYSTEM_ID : m.sender_id === meId ? m.recipient_id : m.sender_id;
        qc.invalidateQueries({ queryKey: ["thread", partner] });
      }),
    [qc, meId]
  );

  const gate = !unlocked ? (
    <KeyGate
      meId={meId}
      configured={configured}
      bundle={keyQ.data}
      onUnlocked={() => {
        setUnlocked(true);
        qc.invalidateQueries({ queryKey: ["message-key"] });
      }}
    />
  ) : null;

  if (typeof view === "object") {
    return (
      <Thread
        meId={meId}
        partnerId={view.partnerId}
        isSystem={!!view.system}
        seedName={view.name}
        seedAvatar={view.avatar}
        unlocked={unlocked}
        gate={gate}
        onBack={() => setView("list")}
      />
    );
  }
  if (view === "new") {
    return <Directory onPick={(u) => setView({ partnerId: u.id, name: u.name, avatar: u.avatar_url })} onBack={() => setView("list")} />;
  }
  return (
    <ConversationList
      unlocked={unlocked}
      gate={gate}
      onOpen={(c) =>
        setView({ partnerId: c.partner.id, name: c.partner.name, avatar: c.partner.avatar_url, system: c.partner.id === SYSTEM_ID })
      }
      onNew={() => setView("new")}
    />
  );
}

function KeyGate({
  meId,
  configured,
  bundle,
  onUnlocked,
}: {
  meId: number;
  configured: boolean;
  bundle: MyKeyBundle | undefined;
  onUnlocked: () => void;
}) {
  const { t } = useTranslation();
  const [pass, setPass] = useState("");
  const [confirm, setConfirm] = useState("");
  const [err, setErr] = useState<string | null>(null);
  const [busy, setBusy] = useState(false);

  async function submit() {
    setErr(null);
    if (pass.length < 6) return setErr(t("messages.passTooShort"));
    if (!configured && pass !== confirm) return setErr(t("messages.passMismatch"));
    setBusy(true);
    try {
      if (configured) {
        await e2ee.unlock(meId, pass, bundle!);
      } else {
        const payload = await e2ee.setup(meId, pass);
        await api.setupMessageKey(payload);
      }
      onUnlocked();
    } catch {
      setErr(configured ? t("messages.wrongPass") : t("messages.setupFailed"));
    } finally {
      setBusy(false);
    }
  }

  return (
    <div className="glass rounded-2xl p-4 space-y-3">
      <div className="flex items-center gap-2">
        <ShieldCheck className="w-5 h-5 text-accent shrink-0" />
        <div className="font-semibold">{configured ? t("messages.unlockTitle") : t("messages.setupTitle")}</div>
      </div>
      <p className="text-sm text-muted">{configured ? t("messages.unlockBody") : t("messages.setupBody")}</p>
      <input
        type="password"
        value={pass}
        onChange={(e) => setPass(e.target.value)}
        onKeyDown={(e) => e.key === "Enter" && configured && submit()}
        placeholder={t("messages.passphrase")}
        className="w-full bg-card border border-border rounded-xl px-3 py-2 text-sm outline-none focus:border-accent"
      />
      {!configured && (
        <input
          type="password"
          value={confirm}
          onChange={(e) => setConfirm(e.target.value)}
          placeholder={t("messages.passphraseConfirm")}
          className="w-full bg-card border border-border rounded-xl px-3 py-2 text-sm outline-none focus:border-accent"
        />
      )}
      {err && <div className="text-sm text-red-400">{err}</div>}
      <button
        onClick={submit}
        disabled={busy || !pass}
        className="inline-flex items-center gap-2 px-4 py-2 rounded-xl font-semibold bg-accent text-accent-fg hover:opacity-90 transition disabled:opacity-40"
      >
        <Lock className="w-4 h-4" />
        {configured ? t("messages.unlock") : t("messages.setUp")}
      </button>
      {!configured && <p className="text-xs text-muted">{t("messages.setupWarn")}</p>}
    </div>
  );
}

function ConversationList({
  unlocked,
  gate,
  onOpen,
  onNew,
}: {
  unlocked: boolean;
  gate: React.ReactNode;
  onOpen: (c: Conversation) => void;
  onNew: () => void;
}) {
  const { t } = useTranslation();
  const q = useLiveQuery<{ items: Conversation[] }>(["conversations"], api.conversations, { intervalMs: POLL_MS });
  const items = q.data?.items ?? [];
  const [previews, setPreviews] = useState<Record<number, string>>({});

  // Decrypt the latest message of each user conversation for the list preview (system previews
  // are already plaintext). Runs whenever the conversations or unlock state change.
  useEffect(() => {
    if (!unlocked) return;
    let cancelled = false;
    (async () => {
      const out: Record<number, string> = {};
      for (const c of items) {
        const m = c.last_message;
        if (m.kind === "user" && m.ciphertext && m.iv) {
          try {
            const pub = await partnerPub(c.partner.id);
            out[c.partner.id] = await e2ee.decryptFrom(c.partner.id, pub, m.ciphertext, m.iv);
          } catch {
            out[c.partner.id] = "⚠";
          }
        }
      }
      if (!cancelled) setPreviews(out);
    })();
    return () => {
      cancelled = true;
    };
  }, [items, unlocked]);

  function previewText(c: Conversation): string {
    const m = c.last_message;
    if (m.kind === "system") return m.body ?? "";
    if (!unlocked) return "🔒 " + t("messages.encrypted");
    return previews[c.partner.id] ?? "…";
  }

  return (
    <div className="space-y-3">
      <div className="flex items-center justify-between">
        <div className="text-xs uppercase tracking-wide text-muted">{t("messages.conversations")}</div>
        {unlocked && (
          <button
            onClick={onNew}
            className="inline-flex items-center gap-1.5 text-xs px-2.5 py-1.5 rounded-lg bg-accent text-accent-fg hover:opacity-90 transition"
          >
            <SquarePen className="w-4 h-4" />
            {t("messages.new")}
          </button>
        )}
      </div>
      {gate}
      {q.isLoading && !q.data ? (
        <div className="p-8 text-muted text-center">{t("common.loading")}</div>
      ) : items.length === 0 ? (
        <div className="glass rounded-2xl p-10 text-center text-muted">{t("messages.empty")}</div>
      ) : (
        <div className="flex flex-col gap-1.5">
          {items.map((c) => (
            <button
              key={c.partner.id}
              onClick={() => onOpen(c)}
              className="flex items-center gap-3 p-2.5 rounded-xl text-left hover:bg-card transition"
            >
              <Avatar src={c.partner.avatar_url} fallback={c.partner.name} className="w-10 h-10 rounded-full shrink-0 text-sm" />
              <div className="min-w-0 flex-1">
                <div className="flex items-center gap-2">
                  <span className={`truncate ${c.unread > 0 ? "font-semibold" : "font-medium"}`}>{c.partner.name}</span>
                  <span className="ml-auto text-[11px] text-muted shrink-0">{relativeTime(c.last_message.created_at)}</span>
                </div>
                <div className="flex items-center gap-2">
                  <span className={`truncate text-sm ${c.unread > 0 ? "text-fg" : "text-muted"}`}>{previewText(c)}</span>
                  {c.unread > 0 && (
                    <span className="ml-auto shrink-0 min-w-5 h-5 px-1.5 grid place-items-center rounded-full bg-accent text-accent-fg text-[11px] font-semibold">
                      {c.unread}
                    </span>
                  )}
                </div>
              </div>
            </button>
          ))}
        </div>
      )}
    </div>
  );
}

function Directory({ onPick, onBack }: { onPick: (u: MessageUser) => void; onBack: () => void }) {
  const { t } = useTranslation();
  const [filter, setFilter] = useState("");
  const q = useQuery({ queryKey: ["message-users"], queryFn: api.messageUsers });
  const users = (q.data ?? []).filter((u) => u.name.toLowerCase().includes(filter.trim().toLowerCase()));
  return (
    <div className="space-y-3">
      <div className="flex items-center gap-2">
        <button onClick={onBack} className="p-1.5 rounded-lg text-muted hover:text-fg hover:bg-card transition">
          <ArrowLeft className="w-4 h-4" />
        </button>
        <div className="text-xs uppercase tracking-wide text-muted">{t("messages.newTitle")}</div>
      </div>
      <input
        value={filter}
        onChange={(e) => setFilter(e.target.value)}
        placeholder={t("messages.searchPeople")}
        className="w-full bg-card border border-border rounded-xl px-3 py-2 text-sm outline-none focus:border-accent"
      />
      {q.isLoading ? (
        <div className="p-8 text-muted text-center">{t("common.loading")}</div>
      ) : users.length === 0 ? (
        <div className="glass rounded-2xl p-10 text-center text-muted">{t("messages.noPeople")}</div>
      ) : (
        <div className="flex flex-col gap-1">
          {users.map((u) => (
            <button
              key={u.id}
              onClick={() => onPick(u)}
              className="flex items-center gap-3 p-2.5 rounded-xl text-left hover:bg-card transition"
            >
              <Avatar src={u.avatar_url} fallback={u.name} className="w-9 h-9 rounded-full shrink-0 text-sm" />
              <span className="truncate font-medium">{u.name}</span>
            </button>
          ))}
        </div>
      )}
    </div>
  );
}

function Thread({
  meId,
  partnerId,
  isSystem,
  seedName,
  seedAvatar,
  unlocked,
  gate,
  onBack,
}: {
  meId: number;
  partnerId: number;
  isSystem: boolean;
  seedName?: string;
  seedAvatar?: string | null;
  unlocked: boolean;
  gate: React.ReactNode;
  onBack: () => void;
}) {
  const { t } = useTranslation();
  const qc = useQueryClient();
  const [draft, setDraft] = useState("");
  const [plain, setPlain] = useState<Record<number, string>>({});
  const bottomRef = useRef<HTMLDivElement>(null);
  const needsKey = !isSystem && !unlocked;

  const q = useLiveQuery<{ partner: MessageUser; items: Message[] }>(
    ["thread", partnerId],
    () => api.thread(partnerId),
    { intervalMs: POLL_MS, enabled: !needsKey }
  );
  const partner = q.data?.partner;
  const items = q.data?.items ?? [];

  // Decrypt user-message bodies for display (system messages are already plaintext).
  useEffect(() => {
    if (isSystem || !unlocked) return;
    let cancelled = false;
    (async () => {
      const out: Record<number, string> = {};
      for (const m of items) {
        if (m.ciphertext && m.iv) {
          try {
            const pub = await partnerPub(partnerId);
            out[m.id] = await e2ee.decryptFrom(partnerId, pub, m.ciphertext, m.iv);
          } catch {
            out[m.id] = "⚠ " + t("messages.cantDecrypt");
          }
        }
      }
      if (!cancelled) setPlain(out);
    })();
    return () => {
      cancelled = true;
    };
  }, [items, unlocked, isSystem, partnerId, t]);

  useEffect(() => {
    if (q.dataUpdatedAt) {
      qc.invalidateQueries({ queryKey: ["conversations"] });
      qc.invalidateQueries({ queryKey: ["message-unread"] });
    }
  }, [q.dataUpdatedAt, qc]);

  useEffect(() => {
    bottomRef.current?.scrollIntoView({ block: "end" });
  }, [items.length, plain]);

  const send = useMutation({
    mutationFn: async (text: string) => {
      const pub = await partnerPub(partnerId);
      const { ciphertext, iv } = await e2ee.encryptFor(partnerId, pub, text);
      return api.sendMessage(partnerId, ciphertext, iv);
    },
    onSuccess: () => {
      setDraft("");
      qc.invalidateQueries({ queryKey: ["thread", partnerId] });
      qc.invalidateQueries({ queryKey: ["conversations"] });
    },
  });

  const submit = () => {
    const text = draft.trim();
    if (text && !send.isPending) send.mutate(text);
  };

  const name = partner?.name ?? seedName ?? "";
  return (
    <div className="flex flex-col h-[60vh] min-h-80">
      <div className="flex items-center gap-3 pb-3 border-b border-border">
        <button onClick={onBack} className="p-1.5 rounded-lg text-muted hover:text-fg hover:bg-card transition">
          <ArrowLeft className="w-4 h-4" />
        </button>
        <Avatar src={partner?.avatar_url ?? seedAvatar} fallback={name} className="w-9 h-9 rounded-full shrink-0 text-sm" />
        <span className="font-semibold truncate">{name}</span>
        {!isSystem && <ShieldCheck className="w-4 h-4 text-muted shrink-0" aria-label={t("messages.encryptedHint")} />}
      </div>

      {needsKey ? (
        <div className="flex-1 grid place-items-center py-6">
          <div className="w-full max-w-sm">{gate}</div>
        </div>
      ) : (
        <>
          <div className="flex-1 overflow-y-auto py-3 flex flex-col gap-2">
            {items.length === 0 ? (
              <div className="m-auto text-sm text-muted">{t("messages.threadEmpty")}</div>
            ) : (
              items.map((m) => {
                const mine = m.sender_id === meId;
                const text = m.kind === "system" ? m.body ?? "" : plain[m.id] ?? "…";
                return (
                  <div key={m.id} className={`flex ${mine ? "justify-end" : "justify-start"}`}>
                    <div
                      className={`max-w-[78%] rounded-2xl px-3 py-2 text-sm whitespace-pre-wrap break-words ${
                        mine ? "bg-accent text-accent-fg rounded-br-sm" : "bg-card rounded-bl-sm"
                      }`}
                    >
                      {text}
                      <div className={`text-[10px] mt-1 ${mine ? "text-accent-fg/70" : "text-muted"}`}>
                        {relativeTime(m.created_at)}
                      </div>
                    </div>
                  </div>
                );
              })
            )}
            <div ref={bottomRef} />
          </div>

          {isSystem ? (
            <div className="pt-3 border-t border-border text-center text-xs text-muted">{t("messages.systemReadOnly")}</div>
          ) : (
            <div className="flex items-end gap-2 pt-3 border-t border-border">
              <textarea
                value={draft}
                onChange={(e) => setDraft(e.target.value)}
                onKeyDown={(e) => {
                  if (e.key === "Enter" && !e.shiftKey) {
                    e.preventDefault();
                    submit();
                  }
                }}
                rows={1}
                placeholder={t("messages.writePlaceholder")}
                className="flex-1 resize-none max-h-32 bg-card border border-border rounded-xl px-3 py-2 text-sm outline-none focus:border-accent"
              />
              <button
                onClick={submit}
                disabled={!draft.trim() || send.isPending}
                className="shrink-0 p-2.5 rounded-xl bg-accent text-accent-fg hover:opacity-90 transition disabled:opacity-40"
                aria-label={t("messages.send")}
                title={t("messages.send")}
              >
                <Send className="w-4 h-4" />
              </button>
            </div>
          )}
        </>
      )}
    </div>
  );
}