src/app/components/analysis/AnalysisShell.tsx

import { useEffect, useMemo, useState } from "react";
import { useNavigate } from "react-router";
import { PanelRightOpen } from "lucide-react";
import { toast } from "sonner";
import KnowledgeManagement from "../KnowledgeManagement";
import {
  createAnalysisMessage,
  deleteAnalysis,
  getAnalysis,
  getAnalysisMessages,
  getDatabaseClients,
  getDocuments,
  listAnalyses,
  updateAnalysisDataBind,
  type Analysis,
  type AnalysisMessage,
  type DataBindItem,
} from "@/services/orchestrationApi";
import { AgenticApiError, getObservability, streamChat, streamHelp, type AgentStreamEvent } from "@/services/agenticApi";
import { AnalysisHeader } from "./AnalysisHeader";
import { AnalysisSidebar } from "./AnalysisSidebar";
import { AppNavigation, type AppMenuKey } from "./AppNavigation";
import { ChatInput } from "./ChatInput";
import { HomeDashboard } from "./HomeDashboard";
import { MessageList } from "./MessageList";
import { NewAnalysisDialog } from "./NewAnalysisDialog";
import { ReportSidebar } from "./ReportSidebar";
import { clearSession, getCurrentSession } from "./session";
import type { StreamState, UiMessage } from "./types";
import { cx, makeLocalId } from "./utils";

function toUiMessage(message: AnalysisMessage): UiMessage {
  return {
    id: message.id,
    role: message.role,
    content: message.content,
    created_at: message.created_at,
    status: "complete",
  };
}

const wait = (ms: number) => new Promise((resolve) => window.setTimeout(resolve, ms));

export function AnalysisShell() {
  const navigate = useNavigate();
  const session = useMemo(() => getCurrentSession(), []);
  const [activeMenu, setActiveMenu] = useState<AppMenuKey>("home");
  const [analyses, setAnalyses] = useState<Analysis[]>([]);
  const [activeAnalysis, setActiveAnalysis] = useState<Analysis | null>(null);
  const [messages, setMessages] = useState<UiMessage[]>([]);
  const [loadingAnalyses, setLoadingAnalyses] = useState(true);
  const [loadingMessages, setLoadingMessages] = useState(false);
  const [streamState, setStreamState] = useState<StreamState>("idle");
  const [newAnalysisOpen, setNewAnalysisOpen] = useState(false);
  const [reportCollapsed, setReportCollapsed] = useState(false);
  const [staleSources, setStaleSources] = useState<DataBindItem[]>([]);

  const selectAnalysis = async (analysis: Analysis) => {
    setActiveAnalysis(analysis);
    setLoadingMessages(true);
    try {
      const [detail, history] = await Promise.all([getAnalysis(analysis.id), getAnalysisMessages(analysis.id)]);
      setActiveAnalysis(detail);
      setAnalyses((prev) => prev.map((item) => (item.id === detail.id ? detail : item)));
      setMessages(history.map(toUiMessage));
    } catch (err) {
      toast.error(err instanceof Error ? err.message : "Failed to open analysis");
      setMessages([]);
    } finally {
      setLoadingMessages(false);
    }
  };

  const loadAnalyses = async () => {
    setLoadingAnalyses(true);
    try {
      const result = await listAnalyses({ status: "active", page: 1, limit: 50 });
      setAnalyses(result.analyses);
      if (!activeAnalysis && result.analyses.length > 0) {
        await selectAnalysis(result.analyses[0]);
      }
    } catch (err) {
      toast.error(err instanceof Error ? err.message : "Failed to load analyses");
    } finally {
      setLoadingAnalyses(false);
    }
  };

  useEffect(() => {
    loadAnalyses();
  }, []);

  useEffect(() => {
    const checkStaleSources = async () => {
      if (!session?.user_id || !activeAnalysis?.data_bind.length) {
        setStaleSources([]);
        return;
      }

      try {
        const [documents, databaseClients] = await Promise.all([getDocuments(session.user_id), getDatabaseClients(session.user_id)]);
        const documentIds = new Set(documents.map((document) => document.id));
        const databaseIds = new Set(databaseClients.map((client) => client.id));
        setStaleSources(
          activeAnalysis.data_bind.filter((source) => {
            if (source.group_type === "document") return !documentIds.has(source.id);
            if (source.group_type === "database") return !databaseIds.has(source.id);
            return false;
          })
        );
      } catch {
        setStaleSources([]);
      }
    };

    checkStaleSources();
  }, [activeAnalysis?.id, activeAnalysis?.data_bind_version, activeMenu, session?.user_id]);

  const fetchTraceability = async (analysisId: string, messageId: string | undefined, uiMessageId: string) => {
    if (!messageId) return;
    setMessages((prev) => prev.map((msg) => (msg.id === uiMessageId ? { ...msg, traceabilityLoading: true, traceabilityError: null } : msg)));

    for (let attempt = 0; attempt < 3; attempt += 1) {
      try {
        const traceability = await getObservability(analysisId, messageId);
        setMessages((prev) => prev.map((msg) => (msg.id === uiMessageId ? { ...msg, traceability, traceabilityLoading: false } : msg)));
        return;
      } catch (err) {
        if (err instanceof AgenticApiError && err.status === 404 && attempt < 2) {
          await wait(900);
          continue;
        }
        const traceabilityError =
          err instanceof AgenticApiError && err.status === 404
            ? "Traceability is still being prepared. Try again after the next answer."
            : err instanceof Error
            ? err.message
            : "Failed to load traceability";
        setMessages((prev) => prev.map((msg) => (msg.id === uiMessageId ? { ...msg, traceabilityLoading: false, traceabilityError } : msg)));
        return;
      }
    }
  };

  const runAgentStream = async (kind: "chat" | "help", prompt?: string) => {
    if (!session?.user_id || !activeAnalysis || streamState !== "idle") return;

    const aiId = makeLocalId(kind === "chat" ? "ai" : "help");
    let finalContent = "";
    let doneMessageId: string | undefined;
    let streamError: string | null = null;

    setMessages((prev) => [
      ...prev,
      {
        id: aiId,
        role: "ai",
        content: "",
        created_at: new Date().toISOString(),
        status: "streaming",
        statusText: kind === "chat" ? "Reading analysis context" : "Preparing help",
      },
    ]);

    setStreamState(kind === "chat" ? "streaming-chat" : "streaming-help");
    try {
      const onEvent = (event: AgentStreamEvent) => {
        if (event.type === "chunk") {
          finalContent += event.data;
          setMessages((prev) => prev.map((msg) => (msg.id === aiId ? { ...msg, content: finalContent, statusText: undefined } : msg)));
        }
        if (event.type === "status") {
          setMessages((prev) => prev.map((msg) => (msg.id === aiId ? { ...msg, statusText: event.data } : msg)));
        }
        if (event.type === "sources") {
          setMessages((prev) => prev.map((msg) => (msg.id === aiId ? { ...msg, sources: event.data } : msg)));
        }
        if (event.type === "done") {
          doneMessageId = typeof event.data.message_id === "string" ? event.data.message_id : undefined;
        }
        if (event.type === "error") {
          streamError = event.data;
          setMessages((prev) => prev.map((msg) => (msg.id === aiId ? { ...msg, status: "error", statusText: event.data } : msg)));
        }
      };

      if (kind === "chat" && prompt) {
        await streamChat({ user_id: session.user_id, analysis_id: activeAnalysis.id, message: prompt }, onEvent);
      } else {
        await streamHelp({ user_id: session.user_id, analysis_id: activeAnalysis.id }, onEvent);
      }

      if (streamError) return;

      let traceTargetId = aiId;
      setMessages((prev) =>
        prev.map((msg) => (msg.id === aiId ? { ...msg, status: "complete", messageId: doneMessageId, traceabilityLoading: Boolean(doneMessageId) } : msg))
      );

      if (finalContent.trim()) {
        const saved = await createAnalysisMessage(activeAnalysis.id, { role: "ai", content: finalContent.trim() });
        traceTargetId = saved.id;
        setMessages((prev) =>
          prev.map((msg) =>
            msg.id === aiId
              ? { ...toUiMessage(saved), messageId: doneMessageId, sources: msg.sources, traceabilityLoading: Boolean(doneMessageId) }
              : msg
          )
        );
      }

      await fetchTraceability(activeAnalysis.id, doneMessageId, traceTargetId);
    } catch (err) {
      const message = err instanceof Error ? err.message : "Agent stream failed";
      setMessages((prev) => prev.map((msg) => (msg.id === aiId ? { ...msg, status: "error", statusText: message, traceabilityLoading: false } : msg)));
    } finally {
      setStreamState("idle");
    }
  };

  const handleSend = async (content: string) => {
    if (!activeAnalysis || !session?.user_id || streamState !== "idle") return;
    const localId = makeLocalId("user");
    setMessages((prev) => [
      ...prev,
      { id: localId, role: "user", content, created_at: new Date().toISOString(), status: "sending" },
    ]);

    try {
      const saved = await createAnalysisMessage(activeAnalysis.id, { role: "user", content });
      setMessages((prev) => prev.map((msg) => (msg.id === localId ? toUiMessage(saved) : msg)));
    } catch (err) {
      setMessages((prev) => prev.map((msg) => (msg.id === localId ? { ...msg, status: "error", statusText: "Failed to save message" } : msg)));
      toast.error(err instanceof Error ? err.message : "Failed to save message");
      return;
    }

    await runAgentStream("chat", content);
  };

  const handleHelp = async () => {
    await runAgentStream("help");
  };

  const handleCreated = (analysis: Analysis) => {
    setAnalyses((prev) => [analysis, ...prev.filter((item) => item.id !== analysis.id)]);
    setActiveMenu("analysis-agent");
    selectAnalysis(analysis);
  };

  const handleDelete = async (analysis: Analysis) => {
    if (!window.confirm(`Delete analysis "${analysis.analysis_title}"?`)) return;
    try {
      await deleteAnalysis(analysis.id);
      const next = analyses.filter((item) => item.id !== analysis.id);
      setAnalyses(next);
      if (activeAnalysis?.id === analysis.id) {
        setActiveAnalysis(null);
        setMessages([]);
        setStaleSources([]);
        if (next[0]) selectAnalysis(next[0]);
      }
    } catch (err) {
      toast.error(err instanceof Error ? err.message : "Failed to delete analysis");
    }
  };

  const handleUpdateDataBind = async (items: DataBindItem[]) => {
    if (!activeAnalysis) return;
    const updated = await updateAnalysisDataBind(activeAnalysis.id, activeAnalysis.data_bind_version, items);
    setActiveAnalysis(updated);
    setAnalyses((prev) => prev.map((item) => (item.id === updated.id ? updated : item)));
    setStaleSources([]);
  };

  const logout = () => {
    clearSession();
    navigate("/login", { replace: true });
  };

  const renderAnalysisAgent = () => (
    <div className={cx("grid min-h-screen grid-cols-1", reportCollapsed ? "lg:grid-cols-[18rem_minmax(0,1fr)_3.5rem]" : "lg:grid-cols-[18rem_minmax(0,1fr)_23rem]")}>
      <AnalysisSidebar
        analyses={analyses}
        activeId={activeAnalysis?.id}
        loading={loadingAnalyses}
        onNew={() => setNewAnalysisOpen(true)}
        onSelect={selectAnalysis}
        onDelete={handleDelete}
      />

      <main className="flex min-h-[70vh] min-w-0 flex-col bg-slate-50">
        <AnalysisHeader analysis={activeAnalysis} staleSources={staleSources} onUpdateDataBind={handleUpdateDataBind} />
        <div className="min-h-0 flex-1 overflow-y-auto">
          {loadingMessages ? (
            <div className="flex h-full items-center justify-center text-sm text-slate-500">Loading messages</div>
          ) : (
            <MessageList messages={messages} />
          )}
        </div>
        <ChatInput disabled={!activeAnalysis} streaming={streamState !== "idle"} onSend={handleSend} onHelp={handleHelp} />
      </main>

      {reportCollapsed ? (
        <aside className="hidden min-h-0 border-l border-slate-200 bg-white p-2 lg:block">
          <button
            type="button"
            title="Restore report sidebar"
            aria-label="Restore report sidebar"
            onClick={() => setReportCollapsed(false)}
            className="flex h-10 w-10 items-center justify-center rounded-md text-slate-500 hover:bg-slate-100 hover:text-slate-900"
          >
            <PanelRightOpen className="h-4 w-4" />
          </button>
        </aside>
      ) : (
        <aside className="min-h-0 overflow-y-auto border-l border-slate-200 bg-white p-4">
          <ReportSidebar analysis={activeAnalysis} userId={session?.user_id} onCollapse={() => setReportCollapsed(true)} />
        </aside>
      )}
    </div>
  );

  return (
    <div className="grid min-h-screen grid-cols-1 bg-slate-100 text-slate-900 lg:grid-cols-[16rem_minmax(0,1fr)]">
      <AppNavigation active={activeMenu} userName={session?.name} onChange={setActiveMenu} onLogout={logout} />

      <div className="min-w-0">
        {activeMenu === "home" && <HomeDashboard hasAnalyses={analyses.length > 0} onNavigate={setActiveMenu} />}
        {activeMenu === "knowledge" && <KnowledgeManagement open onClose={() => setActiveMenu("home")} variant="page" />}
        {activeMenu === "analysis-agent" && renderAnalysisAgent()}
      </div>

      <NewAnalysisDialog open={newAnalysisOpen} onClose={() => setNewAnalysisOpen(false)} onCreated={handleCreated} />
    </div>
  );
}