src/app/components/analysis/AnalysisShell.tsx

import { useEffect, useMemo, useState } from "react";
import { useNavigate } from "react-router";
import { FileText, Menu, PanelRightOpen, X } 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 { 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",
    messageId: message.message_id ?? undefined,
  };
}


function normalizeAssistantContent(content: string) {
  return content
    .replace(/([A-Za-z0-9])(?=\*\*)/g, "$1 ")
    .replace(/(\*\*[^*\n]+\*\*)(?=[A-Za-z0-9])/g, "$1 ")
    .replace(/\*\*([^*\n]*?)\*\*/g, (_, inner) => `**${inner.trim()}**`)
    .trim();
}

export function AnalysisShell() {
  const navigate = useNavigate();
  const session = useMemo(() => getCurrentSession(), []);
  const [activeMenu, setActiveMenu] = useState<AppMenuKey>("home");
  const [mobileNavOpen, setMobileNavOpen] = useState(false);
  const [mobileReportOpen, setMobileReportOpen] = useState(false);
  const [desktopNavCollapsed, setDesktopNavCollapsed] = useState(false);
  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 [scrollTo, setScrollTo] = useState<"bottom" | "last-ai" | null>(null);
  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)));

    try {
      const traceability = await getObservability(analysisId, messageId);
      setMessages((prev) => prev.map((msg) => (msg.id === uiMessageId ? { ...msg, traceability, traceabilityLoading: false } : msg)));
    } catch (err) {
      const traceabilityError =
        err instanceof AgenticApiError && err.status === 404
          ? "Traceability not available for this answer."
          : err instanceof Error
          ? err.message
          : "Failed to load traceability";
      setMessages((prev) => prev.map((msg) => (msg.id === uiMessageId ? { ...msg, traceabilityLoading: false, traceabilityError } : msg)));
    }
  };
  useEffect(() => {
    if (!activeAnalysis?.id || loadingMessages) return;

    messages
      .filter(
        (message) =>
          message.role === "ai" &&
          message.status === "complete" &&
          Boolean(message.messageId) &&
          !message.traceability &&
          !message.traceabilityLoading &&
          !message.traceabilityError
      )
      .forEach((message) => {
        fetchTraceability(activeAnalysis.id, message.messageId, message.id);
      });
  }, [activeAnalysis?.id, loadingMessages, messages]);

  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",
      },
    ]);

    setScrollTo("bottom");
    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;

      const completedContent = normalizeAssistantContent(finalContent);
      let traceTargetId = aiId;
      let resolvedMessageId = doneMessageId;
      setScrollTo("last-ai");
      setMessages((prev) =>
        prev.map((msg) =>
          msg.id === aiId
            ? { ...msg, content: completedContent || msg.content, status: "complete", messageId: doneMessageId, traceabilityLoading: Boolean(doneMessageId) }
            : msg
        )
      );

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

      await fetchTraceability(activeAnalysis.id, resolvedMessageId, 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 changeMenu = (menu: AppMenuKey) => {
    setActiveMenu(menu);
    setMobileNavOpen(false);
    setMobileReportOpen(false);
  };

  const activeMenuLabel =
    activeMenu === "analysis-agent" ? "Analysis Agent" : activeMenu === "knowledge" ? "Knowledge" : "Home";

  const renderAnalysisAgent = () => (
    <div className={cx("grid h-full min-h-0 grid-cols-1", reportCollapsed ? "lg:grid-cols-[minmax(0,1fr)_3.5rem]" : "lg:grid-cols-[minmax(0,1fr)_23rem]")}>
      <main className="flex h-full min-h-0 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 overflow-x-hidden overscroll-contain">
          {loadingMessages ? (
            <div className="flex h-full items-center justify-center text-sm text-slate-500">Loading messages</div>
          ) : (
            <MessageList messages={messages} scrollTo={scrollTo} onScrollDone={() => setScrollTo(null)} />
          )}
        </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="hidden min-h-0 overflow-y-auto border-l border-slate-200 bg-white p-4 lg:block">
          <ReportSidebar analysis={activeAnalysis} userId={session?.user_id} onCollapse={() => setReportCollapsed(true)} />
        </aside>
      )}
    </div>
  );

  return (
    <div className="workspace-shell min-h-screen bg-slate-100 text-slate-900">
      <header className="sticky top-0 z-40 flex h-14 items-center justify-between border-b border-slate-200 bg-white px-3 lg:hidden">
        <button
          type="button"
          title="Open workspace menu"
          aria-label="Open workspace menu"
          onClick={() => setMobileNavOpen(true)}
          className="flex h-10 w-10 items-center justify-center rounded-md text-slate-600 hover:bg-slate-100 hover:text-slate-950"
        >
          <Menu className="h-5 w-5" />
        </button>
        <div className="min-w-0 text-center">
          <p className="truncate text-sm font-semibold text-slate-950">{activeMenuLabel}</p>
          <p className="text-xs text-slate-500">Data Eyond</p>
        </div>
        {activeMenu === "analysis-agent" ? (
          <button
            type="button"
            title="Open report panel"
            aria-label="Open report panel"
            onClick={() => setMobileReportOpen(true)}
            className="flex h-10 w-10 items-center justify-center rounded-md text-slate-600 hover:bg-slate-100 hover:text-slate-950"
          >
            <FileText className="h-5 w-5" />
          </button>
        ) : (
          <div className="h-10 w-10" />
        )}
      </header>

      {mobileNavOpen && (
        <div className="fixed inset-0 z-50 lg:hidden" role="dialog" aria-modal="true" aria-label="Workspace navigation">
          <button
            type="button"
            aria-label="Close workspace menu overlay"
            onClick={() => setMobileNavOpen(false)}
            className="absolute inset-0 bg-slate-950/40"
          />
          <div className="relative h-full w-72 max-w-[86vw] bg-white shadow-2xl">
            <AppNavigation
              active={activeMenu}
              userName={session?.name}
              analyses={analyses}
              activeAnalysisId={activeAnalysis?.id}
              loadingAnalyses={loadingAnalyses}
              onChange={changeMenu}
              onLogout={logout}
              onNewAnalysis={() => {
                setNewAnalysisOpen(true);
                setMobileNavOpen(false);
              }}
              onSelectAnalysis={(analysis) => {
                selectAnalysis(analysis);
                setMobileNavOpen(false);
              }}
              onDeleteAnalysis={handleDelete}
            />
            <button
              type="button"
              title="Close workspace menu"
              aria-label="Close workspace menu"
              onClick={() => setMobileNavOpen(false)}
              className="absolute right-3 top-3 flex h-9 w-9 items-center justify-center rounded-md text-slate-500 hover:bg-slate-100 hover:text-slate-950"
            >
              <X className="h-4 w-4" />
            </button>
          </div>
        </div>
      )}
      {mobileReportOpen && activeMenu === "analysis-agent" && (
        <div className="fixed inset-0 z-50 lg:hidden" role="dialog" aria-modal="true" aria-label="Report panel">
          <button
            type="button"
            aria-label="Close report panel overlay"
            onClick={() => setMobileReportOpen(false)}
            className="absolute inset-0 bg-slate-950/40"
          />
          <div className="absolute inset-x-0 bottom-0 max-h-[86vh] overflow-y-auto rounded-t-xl bg-white p-4 shadow-2xl">
            <ReportSidebar analysis={activeAnalysis} userId={session?.user_id} onCollapse={() => setMobileReportOpen(false)} />
          </div>
        </div>
      )}

      <div className={cx("lg:grid", activeMenu === "analysis-agent" ? "h-[calc(100vh-3.5rem)] min-h-0 overflow-hidden lg:h-screen" : "min-h-screen", desktopNavCollapsed ? "lg:grid-cols-[4.5rem_minmax(0,1fr)]" : "lg:grid-cols-[16rem_minmax(0,1fr)]")}>
        <div className="hidden lg:block">
          <AppNavigation
            active={activeMenu}
            userName={session?.name}
            collapsed={desktopNavCollapsed}
            analyses={analyses}
            activeAnalysisId={activeAnalysis?.id}
            loadingAnalyses={loadingAnalyses}
            onChange={setActiveMenu}
            onLogout={logout}
            onToggleCollapsed={() => setDesktopNavCollapsed((value) => !value)}
            onNewAnalysis={() => setNewAnalysisOpen(true)}
            onSelectAnalysis={selectAnalysis}
            onDeleteAnalysis={handleDelete}
          />
        </div>

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

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