ThothII/frontend/src/shell/AppShell.tsx

import { useSessionStream } from "../stream/useSessionStream";
import { useSessionStore } from "../store/sessionStore";
import { WidgetHost } from "./WidgetHost";
import { CentralStatus } from "./CentralStatus";
import { ModelActivityPanel } from "./ModelActivityPanel";
import { NavSessions } from "./NavSessions";
import { SessionMenu } from "./SessionMenu";
import { SessionDocumentsPanel } from "./SessionDocumentsPanel";
import { RenameDialog } from "./RenameDialog";
import { DeleteConfirmDialog } from "./DeleteConfirmDialog";
import { StopConfirmDialog } from "./StopConfirmDialog";
import { SteerInput, ComposerFooter } from "./SteerInput";
import { WorkflowBar } from "./WorkflowBar";
import { Pencil, ArrowLeft, ArrowRight } from "lucide-react";
import { Button } from "../components/ui/button";
import { Toaster } from "../components/ui/sonner";
import { toast } from "sonner";
import {
  closeSession, listSessions, resumeSession, getSession,
  renameSession, setSessionGroup, archiveSession, unarchiveSession, deleteSession,
} from "../api/sessions";
import type { SessionSummary } from "../api/types";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { useEffect, useMemo, useRef, useState } from "react";

/**
 * The page is designed to live INSIDE the Omics Portal chrome (its left sidebar
 * + topbar) once embedded, so this shell intentionally has no left rail and no
 * top header of its own. The session rail sits on the RIGHT, mirroring the
 * portal's left sidebar, and the workflow phases ride a slim strip above the chat.
 */
export function AppShell() {
  const [activeSessionId, setActiveSessionId] = useState<string | null>(null);
  const { data: sessions = [] } = useQuery<SessionSummary[]>({
    queryKey: ["sessions"], queryFn: listSessions, refetchInterval: 10_000,
  });
  const composerRef = useRef<HTMLTextAreaElement>(null);

  const queryClient = useQueryClient();
  const [panelSession, setPanelSession] = useState<SessionSummary | null>(null);
  const [showActivity, setShowActivity] = useState(false);
  const [activeOpen, setActiveOpen] = useState(true);
  const [archiveOpen, setArchiveOpen] = useState(false);
  const [renameTarget, setRenameTarget] = useState<SessionSummary | null>(null);
  const [deleteTarget, setDeleteTarget] = useState<SessionSummary | null>(null);
  const [stopConfirm, setStopConfirm] = useState(false);
  const [collapsedGroups, setCollapsedGroups] = useState<Record<string, boolean>>({});
  const [renameGroupTarget, setRenameGroupTarget] = useState<string | null>(null);

  const groups = useMemo(
    () => [...new Set(sessions.map((s) => s.group).filter((g): g is string => !!g))].sort(),
    [sessions],
  );
  const activeList = sessions.filter((s) => !s.archived);
  const ungroupedActive = activeList.filter((s) => !s.group);
  const archivedList = sessions.filter((s) => s.archived);
  const refresh = () => queryClient.invalidateQueries({ queryKey: ["sessions"] });
  const activeSession = sessions.find((s) => s.id === activeSessionId) ?? null;
  const finalized = activeSession?.status === "finalized";

  function openPanel(id: string) {
    const s = sessions.find((x) => x.id === id);
    if (s) { setPanelSession(s); setShowActivity(false); }
  }
  function toggleActivity() {
    setShowActivity((v) => {
      const next = !v;
      if (next) setPanelSession(null);
      return next;
    });
  }
  async function doResume(id: string) {
    const s = sessions.find((x) => x.id === id) ?? null;
    // Optimistic switch: change to the session view IMMEDIATELY so the click feels
    // instant (the resume POST spawns a Pi process and can take seconds). The
    // working spinner shows straight away; the backend calls run after.
    setPanelSession(null);
    resetSession();
    // Optimistic: the resume POST is about to hand the ball to the harness.
    setAgentActive(true);
    setActiveSessionId(id);
    try {
      await resumeSession(id);
      // Optimistic phase paint: colour the re-entry phase before the first gate.
      // The manifest's `phase` is the 1-based current phase (1..8).
      try {
        const m = (await getSession(id)) as { phase?: number };
        if (typeof m.phase === "number" && m.phase >= 1 && m.phase <= 8) {
          setPhase(`F${m.phase}`);
        }
      } catch {
        /* non-fatal: the first gate will set the phase */
      }
    } catch {
      // Revert the optimistic switch and restore the panel.
      resetSession();
      setActiveSessionId(null);
      if (s) setPanelSession(s);
      toast.error("Failed to resume session.");
    }
  }
  async function move(s: SessionSummary, group: string) {
    try {
      await setSessionGroup(s.id, group); refresh();
    } catch {
      toast.error("Failed to move session.");
    }
  }
  async function newGroup(s: SessionSummary) {
    const name = window.prompt("New group:");
    if (name && name.trim()) {
      try {
        await setSessionGroup(s.id, name.trim()); refresh();
      } catch {
        toast.error("Failed to update group.");
      }
    }
  }
  async function renameGroup(oldName: string, newName: string) {
    const trimmed = newName.trim();
    if (!trimmed || trimmed === oldName) return;
    try {
      for (const s of sessions.filter((x) => x.group === oldName)) {
        await setSessionGroup(s.id, trimmed);
      }
      refresh();
    } catch {
      toast.error("Failed to rename group.");
    }
  }
  async function toggleArchive(s: SessionSummary) {
    try {
      await (s.archived ? unarchiveSession(s.id) : archiveSession(s.id));
      if (panelSession?.id === s.id) setPanelSession(null);
      refresh();
    } catch {
      toast.error(s.archived ? "Failed to restore session." : "Failed to archive session.");
    }
  }

  function menuFor(s: SessionSummary) {
    return (
      <SessionMenu
        session={s}
        groups={groups}
        onResume={() => doResume(s.id)}
        onView={() => setPanelSession(s)}
        onRename={() => setRenameTarget(s)}
        onMove={(g) => move(s, g)}
        onNewGroup={() => newGroup(s)}
        onArchiveToggle={() => toggleArchive(s)}
        onDelete={() => setDeleteTarget(s)}
      />
    );
  }

  // The harness "holds the ball" whenever a session is live, no widget is waiting on
  // the human, AND the Pi turn is still in flight (agentActive). Without the last
  // condition the final workflow step — which ends with no follow-up gate — would
  // leave the spinner on forever.
  const pendingWidget = useSessionStore((s) => s.pendingWidget);
  const resetSession = useSessionStore((s) => s.resetSession);
  const setPhase = useSessionStore((s) => s.setPhase);
  const setAgentActive = useSessionStore((s) => s.setAgentActive);
  const lastSystemEvent = useSessionStore((s) => s.lastSystemEvent);
  const agentActive = useSessionStore((s) => s.agentActive);
  const working = Boolean(activeSessionId) && !pendingWidget && agentActive;
  // Processing time counts only while the harness works, not while a finalized
  // session sits idle or a gate awaits the reviewer (pendingWidget).
  const running = working && !finalized;

  useSessionStream(activeSessionId);

  // A backend "session_exit" system event (e.g. the replay server emitting it
  // when the reviewer picks "Esci") asks us to leave the live session view and
  // return to the landing state. We deliberately do NOT also POST /close here —
  // stopSession already does, and this effect must stay side-effect-free on the
  // backend if the session is already inactive.
  useEffect(() => {
    if (lastSystemEvent?.type === "system_event" && (lastSystemEvent as any).event === "session_exit") {
      stopSession();
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [lastSystemEvent]);

  function startNewSession() {
    resetSession();
    setActiveSessionId(null);
    requestAnimationFrame(() => composerRef.current?.focus());
  }

  async function stopSession() {
    if (!activeSessionId) return;
    try {
      await closeSession(activeSessionId);
    } finally {
      resetSession();
      setActiveSessionId(null);
    }
  }

  return (
    <div className="flex h-screen bg-background text-foreground">
      {panelSession && (
        <SessionDocumentsPanel
          session={panelSession}
          onClose={() => setPanelSession(null)}
          onResume={doResume}
        />
      )}
      {showActivity && <ModelActivityPanel onClose={() => setShowActivity(false)} />}

      {/* Conversation column */}
      <div className="flex min-w-0 flex-1 flex-col">
        {activeSessionId && (
          <div className="relative shrink-0 border-b border-border/70 bg-background/80 px-6 py-2.5 backdrop-blur">
            {/* Arrow toggle for the left Model-activity panel: → opens it, ← closes it. */}
            <button
              type="button"
              onClick={toggleActivity}
              aria-expanded={showActivity}
              aria-label={showActivity ? "Hide model activity" : "Show model activity"}
              title={showActivity ? "Hide model activity" : "Show model activity"}
              className="absolute left-4 top-1/2 grid size-7 -translate-y-1/2 place-items-center rounded-md border border-border text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
            >
              {showActivity ? <ArrowLeft className="size-4" /> : <ArrowRight className="size-4" />}
            </button>
            <div className="mx-auto max-w-4xl">
              <WorkflowBar
                finalized={finalized}
                running={running}
                resetKey={activeSessionId}
              />
            </div>
          </div>
        )}

        <main className="flex-1 overflow-y-auto px-6 py-8">
          <div className="mx-auto flex max-w-4xl flex-col gap-6">
            {activeSessionId ? (
              <>
                <CentralStatus working={working} />
                <WidgetHost sessionId={activeSessionId} />
              </>
            ) : (
              <EmptyState />
            )}
          </div>
        </main>

        {/* Sticky composer: the box stays pinned while the conversation scrolls above. */}
        <div className="shrink-0 px-6 pb-6 pt-2">
          <div className="mx-auto max-w-4xl">
            <div className="rounded-2xl border border-border/80 bg-card shadow-md transition-colors focus-within:border-primary/50 focus-within:ring-3 focus-within:ring-ring/15">
              <div className="px-2.5 py-2">
                <SteerInput
                  sessionId={activeSessionId}
                  onSessionCreated={setActiveSessionId}
                  onStop={() => setStopConfirm(true)}
                  inputRef={composerRef}
                />
              </div>
              <ComposerFooter />
            </div>
          </div>
        </div>
      </div>

      {/* Right session rail — symmetric to the portal's left sidebar */}
      <aside className="flex w-[15vw] shrink-0 flex-col border-l border-border bg-sidebar">
        <div className="relative px-4 pb-3 pt-5 text-center">
          <h1 className="font-heading text-xl font-semibold leading-none tracking-tight text-foreground">
            Thoth<span className="text-primary">II</span>
          </h1>
          <p className="thot-label mt-1.5">
            Datamart Builder with
            <br />
            Human In The Loop
          </p>
        </div>

        <div className="flex flex-col gap-2 px-4 pb-3">
          <Button
            variant="default"
            size="sm"
            className="w-full"
            onClick={startNewSession}
          >
            New session
          </Button>
        </div>

        {/* L1 — rail title */}
        <div className="px-4 pb-1.5 pt-1">
          <span className="thot-label text-[0.8rem] font-bold tracking-[0.18em] text-primary">
            Sessions
          </span>
        </div>
        <div className="flex-1 overflow-y-auto px-2 pb-4">
          {/* L2 — section toggle */}
          <button
            type="button"
            onClick={() => setActiveOpen((v) => !v)}
            aria-expanded={activeOpen}
            className="thot-label flex w-full items-center gap-1 px-1 pb-1 pt-1 text-left text-foreground/65 hover:text-foreground"
          >
            <span className="select-none">{activeOpen ? "▾" : "▸"}</span>
            <span>Active sessions</span>
          </button>
          {activeOpen && (
            <div className="flex flex-col gap-3 pb-2">
              {groups.map((g) => (
                <div key={g}>
                  <div className="group/gh flex items-center gap-1">
                    {/* L2 — group header (sub-level of the section) */}
                    <button
                      type="button"
                      onClick={() => setCollapsedGroups((c) => ({ ...c, [g]: !c[g] }))}
                      aria-expanded={!collapsedGroups[g]}
                      className="thot-label flex min-w-0 flex-1 items-center gap-1 px-1 pb-1 text-left hover:text-foreground"
                    >
                      <span className="select-none">{collapsedGroups[g] ? "▸" : "▾"}</span>
                      <span className="truncate">{g}</span>
                    </button>
                    <button
                      type="button"
                      aria-label={`Rename group ${g}`}
                      onClick={() => setRenameGroupTarget(g)}
                      className="rounded-md p-0.5 text-muted-foreground opacity-0 transition-opacity hover:bg-accent group-hover/gh:opacity-100"
                    >
                      <Pencil className="size-3" />
                    </button>
                  </div>
                  {!collapsedGroups[g] && (
                    <NavSessions
                      sessions={activeList.filter((s) => s.group === g)}
                      activeSessionId={activeSessionId}
                      onOpenPanel={openPanel}
                      menuFor={menuFor}
                    />
                  )}
                </div>
              ))}
              {/* Ungrouped sessions list directly after the groups — no "No group" label.
                  When there are no groups at all, still render it so its empty state can
                  teach first-time users. */}
              {(ungroupedActive.length > 0 || groups.length === 0) && (
                <NavSessions
                  sessions={ungroupedActive}
                  activeSessionId={activeSessionId}
                  onOpenPanel={openPanel}
                  menuFor={menuFor}
                />
              )}
            </div>
          )}

          {/* L2 — section toggle */}
          <button
            type="button"
            onClick={() => setArchiveOpen((v) => !v)}
            aria-expanded={archiveOpen}
            className="thot-label mt-2 flex w-full items-center gap-1 px-1 pb-1 pt-1 text-left text-foreground/65 hover:text-foreground"
          >
            <span className="select-none">{archiveOpen ? "▾" : "▸"}</span>
            <span className="tabular-nums">Archive ({archivedList.length})</span>
          </button>
          {archiveOpen && (
            <NavSessions sessions={archivedList} activeSessionId={activeSessionId} onOpenPanel={openPanel} menuFor={menuFor} />
          )}
        </div>
      </aside>

      <Toaster />

      <StopConfirmDialog
        open={stopConfirm}
        onOpenChange={setStopConfirm}
        onConfirm={stopSession}
      />

      {renameTarget && (
        <RenameDialog
          open
          initial={renameTarget.name ?? ""}
          onOpenChange={(o) => { if (!o) setRenameTarget(null); }}
          onSubmit={async (name) => {
            try {
              await renameSession(renameTarget.id, name);
              setRenameTarget(null);
              refresh();
            } catch {
              toast.error("Failed to rename session.");
            }
          }}
        />
      )}
      {renameGroupTarget && (
        <RenameDialog
          open
          initial={renameGroupTarget}
          onOpenChange={(o) => { if (!o) setRenameGroupTarget(null); }}
          onSubmit={async (name) => { await renameGroup(renameGroupTarget, name); setRenameGroupTarget(null); }}
          title="Rename group"
        />
      )}
      {deleteTarget && (
        <DeleteConfirmDialog
          open
          label={deleteTarget.name || deleteTarget.question}
          onOpenChange={(o) => { if (!o) setDeleteTarget(null); }}
          onConfirm={async () => {
            try {
              await deleteSession(deleteTarget.id);
              if (panelSession?.id === deleteTarget.id) setPanelSession(null);
              if (activeSessionId === deleteTarget.id) { resetSession(); setActiveSessionId(null); }
              setDeleteTarget(null);
              refresh();
            } catch {
              toast.error("Failed to delete session.");
            }
          }}
        />
      )}
    </div>
  );
}

function EmptyState() {
  return (
    <div className="flex min-h-[72vh] flex-col items-center justify-center text-center">
      <span className="font-heading text-5xl font-semibold tracking-tight text-foreground">
        Thoth<span className="text-primary">II</span>
      </span>
      <p className="mt-4 max-w-md text-balance text-[0.95rem] leading-relaxed text-muted-foreground">
        Build datamarts from your data warehouse through a guided,
        human-in-the-loop workflow. Type your question in the box below to begin
        a session.
      </p>
    </div>
  );
}