PTSRaiAgent/web/src/dialog/DialogContext.tsx

import { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState, type ReactNode } from "react";

import { api, sendEventStream } from "../api/client";
import type { DialogEvent, Reply } from "../api/types";
import { useSession } from "../session/SessionContext";
import { sessionStore, type ChatEntry, type DialogSnapshot } from "../session/store";

/*
 * One conversation with the dialog engine per signed-in patient, shared by every screen:
 * focus screens render `pending` (the latest unanswered batch), «Сообщения» renders `entries`.
 * The engine owns the logic (consent, onboarding, scales, check-ins, meds, crisis) — exactly as in Telegram.
 */

const OUTBOX_POLL_MS = 6000;

let seq = 0;
const newId = () => `${Date.now().toString(36)}-${(seq++).toString(36)}`;
const asEntries = (replies: Reply[]): ChatEntry[] => replies.map((reply) => ({ id: newId(), from: "assistant", reply }));

/** The agent at work: its steps and the answer as it is being written (cleared when the final reply lands). */
export interface LiveReply {
  status: string[];
  text: string;
}

interface DialogApi extends DialogSnapshot {
  busy: boolean;
  live: LiveReply | null;
  error: boolean;
  started: boolean;
  /** Sends an event; `echo` is what the patient «said» (button label / text) for the history. */
  send: (event: DialogEvent, echo?: string) => Promise<boolean>;
  start: () => Promise<boolean>;
  markRead: () => void;
}

const Ctx = createContext<DialogApi | null>(null);
const EMPTY: DialogSnapshot = { entries: [], pending: [], unread: 0 };

export function DialogProvider({ children }: { children: ReactNode }) {
  const { session } = useSession();
  const token = session?.token ?? null;
  const [owner, setOwner] = useState(token);
  const [snap, setSnap] = useState<DialogSnapshot>(() => (token ? sessionStore.loadDialog(token) : EMPTY));
  const [busy, setBusy] = useState(false);
  const [live, setLive] = useState<LiveReply | null>(null);
  const [error, setError] = useState(false);
  const tokenRef = useRef(token);
  tokenRef.current = token;

  // Switching accounts: reset to that token's snapshot during render, so it is never saved under another token.
  if (owner !== token) {
    setOwner(token);
    setSnap(token ? sessionStore.loadDialog(token) : EMPTY);
    setError(false);
  }

  useEffect(() => {
    if (token && owner === token) sessionStore.saveDialog(token, snap);
  }, [token, owner, snap]);

  const send = useCallback(
    async (event: DialogEvent, echo?: string) => {
      const t = tokenRef.current;
      if (!t) return false;
      setBusy(true);
      setError(false);
      try {
        let replies: Reply[];
        if (event.type === "text") {
          // Free text may go to the agent: stream its steps and the answer so the viewer sees it working.
          const said = echo;
          if (said) {
            // Show the patient's message right away, while the agent is still working.
            setSnap((prev) => ({
              ...prev,
              entries: [...prev.entries, { id: newId(), from: "patient" as const, text: said }],
            }));
            echo = undefined;
          }
          setLive({ status: [], text: "" });
          try {
            replies = await sendEventStream(t, event, (e) =>
              setLive((prev) => {
                const cur = prev ?? { status: [], text: "" };
                if (e.kind === "status") return { ...cur, status: [...cur.status, e.text] };
                if (e.kind === "delta") return { ...cur, text: cur.text + e.text };
                if (e.kind === "reset") return { ...cur, text: "" };
                return cur;
              }),
            );
          } catch {
            replies = await api.sendEvent(t, event); // proxies without SSE support: plain request
          }
        } else {
          replies = await api.sendEvent(t, event);
        }
        setSnap((prev) => ({
          entries: [
            ...prev.entries,
            ...(echo ? [{ id: newId(), from: "patient" as const, text: echo }] : []),
            ...asEntries(replies),
          ],
          pending: replies,
          unread: 0,
        }));
        return true;
      } catch {
        setError(true);
        return false;
      } finally {
        setLive(null);
        setBusy(false);
      }
    },
    [],
  );

  const start = useCallback(() => send({ type: "start" }), [send]);
  const markRead = useCallback(() => setSnap((prev) => (prev.unread ? { ...prev, unread: 0 } : prev)), []);

  // System-initiated messages (reminders, check-ins, scales) arrive through the outbox.
  useEffect(() => {
    if (!token) return;
    const timer = window.setInterval(() => {
      api
        .outbox(token)
        .then((replies) => {
          if (!replies.length || tokenRef.current !== token) return;
          setSnap((prev) => ({
            entries: [...prev.entries, ...asEntries(replies)],
            pending: replies,
            unread: prev.unread + replies.length,
          }));
        })
        .catch(() => undefined);
    }, OUTBOX_POLL_MS);
    return () => window.clearInterval(timer);
  }, [token]);

  const value = useMemo<DialogApi>(
    () => ({ ...snap, busy, live, error, started: snap.entries.length > 0, send, start, markRead }),
    [snap, busy, live, error, send, start, markRead],
  );
  return <Ctx.Provider value={value}>{children}</Ctx.Provider>;
}

export function useDialog(): DialogApi {
  const value = useContext(Ctx);
  if (!value) throw new Error("useDialog outside DialogProvider");
  return value;
}