turbostarter/apps/web/src/modules/copilot/components/CopilotPanel.tsx

"use client";

import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport } from "ai";
import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
import { AnimatePresence } from "motion/react";
import { useQuery } from "@tanstack/react-query";
import { toast } from "sonner";

import { cn } from "@turbostarter/ui";
import { Button } from "@turbostarter/ui-web/button";
import { Icons } from "@turbostarter/ui-web/icons";
import { ScrollArea } from "@turbostarter/ui-web/scroll-area";

import { api } from "~/lib/api/client";
import { MemoizedMarkdown } from "~/modules/common/markdown/memoized-markdown";
import { Prose } from "~/modules/common/prose";
import { useGraphStore } from "~/modules/diagram/stores/useGraphStore";
import { flowToGraph } from "~/modules/diagram/lib/graph-converter";
import { useGraphMutation, persistGraphData } from "../hooks/useGraphMutation";
import { useProposalDiff } from "../hooks/useProposalDiff";
import { BadgeChip } from "./BadgeChip";

import type { DiagramType } from "~/modules/diagram/types/graph";

type ProposalStatus = "idle" | "pending" | "accepted" | "rejected";

/** Shared accept handler — used by CopilotPanel, AssistantBubble, ProposalBar, DiagramCanvas */
export function acceptCurrentProposal(diagramId: string) {
  const store = useGraphStore.getState();
  const patch = store.proposedPatch;
  store.acceptProposal();
  if (patch) {
    persistGraphData(diagramId, patch);
  }
}

/** Shared reject handler */
export function rejectCurrentProposal() {
  useGraphStore.getState().rejectProposal();
}

// Type helper for tool invocation parts from AI SDK
interface ToolPart {
  type: string;
  toolCallId: string;
  state: "input-streaming" | "input-available" | "output-available" | "output-error";
  output?: unknown;
  input?: unknown;
}

function isGenerateDiagramTool(part: { type: string }): part is ToolPart {
  return part.type === "tool-generateDiagram";
}

interface CopilotPanelProps {
  diagramId: string;
  diagramType: DiagramType;
  initialDescription?: string;
  isSharedView?: boolean;
}

export function CopilotPanel({ diagramId, diagramType, initialDescription, isSharedView }: CopilotPanelProps) {
  const chatId = useMemo(() => `copilot-${diagramId}`, [diagramId]);
  const [input, setInput] = useState("");
  const scrollRef = useRef<HTMLDivElement>(null);
  const inputRef = useRef<HTMLTextAreaElement>(null);
  const userScrolledRef = useRef(false);
  const appliedToolCallIds = useRef(new Set<string>());
  const hasSentInitial = useRef(false);

  const { proposeGraphPatch } = useGraphMutation(diagramId, diagramType);
  const proposalStatus = useGraphStore((s) => s.proposalStatus);
  const lastProposalOutcome = useGraphStore((s) => s.lastProposalOutcome);
  const { changeSummary } = useProposalDiff();

  // Subscribe to prefillChat from hover affordances / command palette
  const prefillChat = useGraphStore((s) => s.prefillChat);
  useEffect(() => {
    if (prefillChat) {
      setInput(prefillChat.text);
      // Position cursor at end of pre-filled text
      requestAnimationFrame(() => {
        const el = inputRef.current;
        if (el) {
          el.focus();
          el.setSelectionRange(el.value.length, el.value.length);
        }
      });
      useGraphStore.getState().clearPrefillChat();
    }
  }, [prefillChat]);

  // Subscribe to selected nodes for badge chips
  const selectedNodeIds = useGraphStore((s) => s.selectedNodeIds);
  const storeNodes = useGraphStore((s) => s.nodes);
  const storeEdges = useGraphStore((s) => s.edges);

  const selectedElements = useMemo(() => {
    if (selectedNodeIds.length === 0) return [];
    const nodeMap = new Map(storeNodes.map((n) => [n.id, n]));
    return selectedNodeIds
      .map((id) => {
        const node = nodeMap.get(id);
        if (!node) return null;
        const data = node.data as { type?: string; label?: string };
        return { id: node.id, type: data.type ?? "unknown", label: data.label ?? node.id };
      })
      .filter((e): e is { id: string; type: string; label: string } => e !== null);
  }, [selectedNodeIds, storeNodes]);

  const handleDismissBadge = useCallback((nodeId: string) => {
    const current = useGraphStore.getState().selectedNodeIds;
    useGraphStore.getState().setSelectedNodeIds(current.filter((id) => id !== nodeId));
  }, []);

  // Scope indicator: show what context the AI will see
  const scopeInfo = useMemo(() => {
    if (selectedElements.length === 0) return null;
    const selectedIds = new Set(selectedNodeIds);
    const connectedCount = storeEdges.filter(
      (e) => selectedIds.has(e.source) || selectedIds.has(e.target),
    ).length;
    const label =
      selectedElements.length === 1
        ? selectedElements[0]!.label
        : `${selectedElements.length} elements`;
    return `Context: ${label} + ${connectedCount} connected edge${connectedCount !== 1 ? "s" : ""}`;
  }, [selectedElements, selectedNodeIds, storeEdges]);

  // Dynamic placeholder text
  const placeholder = useMemo(() => {
    if (selectedElements.length === 0) return "Describe what you want to build...";
    if (selectedElements.length === 1) return `Describe changes to ${selectedElements[0]!.label}...`;
    return `Describe changes to ${selectedElements.length} elements...`;
  }, [selectedElements]);

  // Fetch existing chat history on mount (H1 fix)
  const { data: initialMessages } = useQuery({
    queryKey: ["copilot", "messages", chatId],
    queryFn: async () => {
      const res = await api.ai.copilot.messages.$get({
        query: { chatId },
      });
      if (!res.ok) return [];
      return res.json();
    },
    staleTime: Infinity,
  });

  // Memoize transport to avoid recreation on every render (M2 fix)
  const transport = useMemo(
    () =>
      new DefaultChatTransport({
        api: api.ai.copilot.$url().toString(),
        prepareSendMessagesRequest: ({ messages, id }) => {
          const lastMessage = messages.at(-1);

          // Serialize current graph state for AI context
          const currentNodes = useGraphStore.getState().nodes;
          const currentEdges = useGraphStore.getState().edges;
          const graphContext =
            currentNodes.length > 0
              ? JSON.stringify(
                  flowToGraph(currentNodes, currentEdges, {
                    version: "1",
                    title: "",
                    diagramType,
                  }),
                )
              : undefined;

          // Build selected element context for targeted modifications
          const currentSelectedIds = useGraphStore.getState().selectedNodeIds;
          const selectedEls =
            currentSelectedIds.length > 0
              ? currentSelectedIds
                  .map((nid) => {
                    const node = currentNodes.find((n) => n.id === nid);
                    if (!node) return null;
                    const data = node.data as { type?: string; label?: string };
                    return { id: node.id, type: data.type ?? "unknown", label: data.label ?? node.id };
                  })
                  .filter(Boolean)
              : undefined;

          return {
            body: {
              ...lastMessage,
              chatId: id,
              diagramId,
              diagramType,
              graphContext,
              selectedElements: selectedEls,
            },
          };
        },
      }),
    [diagramId, diagramType],
  );

  const { messages, sendMessage, status, error, stop, setMessages } = useChat({
    id: chatId,
    transport,
    onError: (err) => {
      console.error("[copilot]", err);
      toast.error("Failed to get AI response");
    },
  });

  // Seed chat with persisted history once loaded
  useEffect(() => {
    if (initialMessages && initialMessages.length > 0 && messages.length === 0) {
      setMessages(
        initialMessages.map((m) => ({
          ...m,
          createdAt: new Date(),
        })),
      );
    }
  }, [initialMessages, messages.length, setMessages]);

  // Auto-send initial description from wizard (chat-first onboarding)
  useEffect(() => {
    if (
      initialDescription &&
      !hasSentInitial.current &&
      messages.length === 0 &&
      !initialMessages?.length
    ) {
      hasSentInitial.current = true;
      void sendMessage({ text: initialDescription, metadata: {} });
      // Clean only the desc URL param without navigation
      const url = new URL(window.location.href);
      url.searchParams.delete("desc");
      window.history.replaceState({}, "", url.pathname + url.search);
    }
  }, [initialDescription, messages.length, initialMessages, sendMessage]);

  // Detect and apply graph patches from tool invocations
  useEffect(() => {
    for (const message of messages) {
      if (message.role !== "assistant") continue;
      for (const part of message.parts) {
        if (
          isGenerateDiagramTool(part) &&
          part.state === "output-available" &&
          !appliedToolCallIds.current.has(part.toolCallId)
        ) {
          appliedToolCallIds.current.add(part.toolCallId);
          const result = part.output as
            | { success: true; data: Parameters<typeof applyGraphPatch>[0] }
            | { success: false; errors: string[] };

          if (result.success) {
            proposeGraphPatch(result.data);
          } else {
            toast.error("Diagram generation failed: invalid graph structure");
            console.error("[copilot] Graph validation errors:", result.errors);
          }
        }
      }
    }
  }, [messages, proposeGraphPatch]);

  const isSubmitting = status === "submitted" || status === "streaming";

  // Check if currently generating a diagram (tool call in progress)
  const isGeneratingDiagram = useMemo(() => {
    const lastMessage = messages.at(-1);
    if (!lastMessage || lastMessage.role !== "assistant") return false;
    return lastMessage.parts.some(
      (p) =>
        isGenerateDiagramTool(p) &&
        (p.state === "input-streaming" || p.state === "input-available"),
    );
  }, [messages]);

  // Auto-scroll on new content, but pause if user scrolled up
  useEffect(() => {
    if (userScrolledRef.current) return;
    const viewport = scrollRef.current?.querySelector(
      "[data-radix-scroll-area-viewport]",
    );
    if (viewport) {
      viewport.scrollTop = viewport.scrollHeight;
    }
  }, [messages]);

  // Detect user scroll
  useEffect(() => {
    const viewport = scrollRef.current?.querySelector(
      "[data-radix-scroll-area-viewport]",
    );
    if (!viewport) return;

    let timeoutId: NodeJS.Timeout;
    const handleScroll = () => {
      const isAtBottom =
        Math.abs(
          viewport.scrollHeight - viewport.scrollTop - viewport.clientHeight,
        ) < 80;
      userScrolledRef.current = !isAtBottom;
      clearTimeout(timeoutId);
      timeoutId = setTimeout(() => {
        userScrolledRef.current = false;
      }, 2000);
    };

    viewport.addEventListener("scroll", handleScroll);
    return () => {
      viewport.removeEventListener("scroll", handleScroll);
      clearTimeout(timeoutId);
    };
  }, []);

  const handleSend = useCallback(() => {
    const text = input.trim();
    if (!text || isSubmitting) return;

    void sendMessage({
      text,
      metadata: {},
    });
    setInput("");
    userScrolledRef.current = false;
  }, [input, isSubmitting, sendMessage]);

  const handleAccept = useCallback(() => {
    acceptCurrentProposal(diagramId);
  }, [diagramId]);

  const handleReject = useCallback(() => {
    rejectCurrentProposal();
  }, []);

  const handleKeyDown = useCallback(
    (e: React.KeyboardEvent) => {
      const currentProposalStatus = useGraphStore.getState().proposalStatus;

      // Proposal pending + Enter + empty textarea → accept
      if (e.key === "Enter" && !e.shiftKey && currentProposalStatus === "pending" && !input.trim()) {
        e.preventDefault();
        handleAccept();
        return;
      }
      // Proposal pending + Escape → reject (don't clear badges)
      if (e.key === "Escape" && currentProposalStatus === "pending") {
        e.preventDefault();
        handleReject();
        return;
      }
      // Normal Enter → send message
      if (e.key === "Enter" && !e.shiftKey) {
        e.preventDefault();
        handleSend();
        return;
      }
      // Normal Escape → clear badges
      if (e.key === "Escape" && selectedNodeIds.length > 0) {
        e.preventDefault();
        useGraphStore.getState().setSelectedNodeIds([]);
      }
    },
    [handleSend, handleAccept, handleReject, selectedNodeIds.length, input],
  );

  return (
    <div className="flex h-full flex-col">
      {/* Messages area */}
      <ScrollArea ref={scrollRef} className="flex-1">
        <div className="flex flex-col gap-1 p-3">
          {messages.length === 0 && (
            <EmptyState isSharedView={isSharedView} />
          )}

          {messages.map((message) => (
            <div
              key={message.id}
              className={cn(
                "flex flex-col gap-1",
                message.role === "user" ? "items-end" : "items-start",
              )}
            >
              {message.role === "user" ? (
                <UserBubble message={message} />
              ) : (
                <AssistantBubble
                  message={message}
                  isStreaming={
                    status === "streaming" &&
                    message.id === messages.at(-1)?.id
                  }
                  diagramId={diagramId}
                  proposalStatus={proposalStatus}
                  lastProposalOutcome={lastProposalOutcome}
                  changeSummary={changeSummary}
                />
              )}
            </div>
          ))}

          {/* Typing indicator */}
          {status === "submitted" && (
            <div className="flex items-start gap-2 py-2">
              <div className="flex gap-1">
                <span className="size-1.5 rounded-full bg-muted-foreground/50 animate-bounce [animation-delay:0ms]" />
                <span className="size-1.5 rounded-full bg-muted-foreground/50 animate-bounce [animation-delay:150ms]" />
                <span className="size-1.5 rounded-full bg-muted-foreground/50 animate-bounce [animation-delay:300ms]" />
              </div>
            </div>
          )}

          {/* Diagram generation indicator */}
          {isGeneratingDiagram && (
            <div className="flex items-center gap-2 rounded-lg bg-muted/50 px-3 py-2 text-xs text-muted-foreground">
              <Icons.Loader2 className="size-3 animate-spin" />
              Generating diagram...
            </div>
          )}

          {/* Error display */}
          {error && (
            <div className="rounded-lg bg-destructive/10 px-3 py-2 text-xs text-destructive">
              Something went wrong. Please try again.
            </div>
          )}
        </div>
      </ScrollArea>

      {/* Input area */}
      <div className="shrink-0 border-t border-border p-3">
        {/* Badge chips for selected elements */}
        {selectedElements.length > 0 && (
          <div className="mb-2 flex flex-wrap gap-1" role="list" aria-label="Selected elements">
            <AnimatePresence mode="popLayout">
              {selectedElements.map((el) => (
                <BadgeChip
                  key={el.id}
                  nodeId={el.id}
                  label={el.label}
                  nodeType={el.type}
                  onDismiss={handleDismissBadge}
                />
              ))}
            </AnimatePresence>
            {scopeInfo && (
              <p className="w-full text-[10px] text-muted-foreground/60 mt-0.5">
                {scopeInfo}
              </p>
            )}
          </div>
        )}
        <div className="relative">
          <textarea
            ref={inputRef}
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={handleKeyDown}
            placeholder={placeholder}
            rows={1}
            className="w-full resize-none rounded-lg border border-border bg-background px-3 py-2 pr-20 text-sm placeholder:text-muted-foreground focus:outline-none focus:ring-1 focus:ring-ring"
          />
          <div className="absolute right-1.5 bottom-1.5 flex items-center gap-1">
            {/* Mic button placeholder (Epic 5) */}
            <Button
              size="icon"
              variant="ghost"
              className="size-7"
              disabled
              title="Voice input (coming soon)"
            >
              <Icons.Mic className="size-3.5 text-muted-foreground" />
            </Button>

            {/* Send / Stop button */}
            {isSubmitting ? (
              <Button
                size="icon"
                variant="ghost"
                className="size-7"
                onClick={() => stop()}
              >
                <Icons.Square className="size-3 fill-current" />
              </Button>
            ) : (
              <Button
                size="icon"
                variant="ghost"
                className="size-7"
                disabled={!input.trim()}
                onClick={handleSend}
              >
                <Icons.ArrowUp className="size-3.5" />
              </Button>
            )}
          </div>
        </div>
      </div>
    </div>
  );
}

function EmptyState({ isSharedView }: { isSharedView?: boolean }) {
  if (isSharedView) {
    return (
      <div className="flex flex-col items-center justify-center py-12 text-center">
        <Icons.MessageSquare className="mb-3 size-8 text-muted-foreground/30" />
        <p className="text-sm font-medium text-muted-foreground">
          Join the conversation
        </p>
        <p className="mt-1 text-xs text-muted-foreground/60">
          This diagram was shared with you. Type below to start collaborating.
        </p>
      </div>
    );
  }

  return (
    <div className="flex flex-col items-center justify-center py-12 text-center">
      <Icons.Sparkles className="mb-3 size-8 text-muted-foreground/30" />
      <p className="text-sm font-medium text-muted-foreground">
        What are you designing today?
      </p>
      <p className="mt-1 text-xs text-muted-foreground/60">
        Describe your diagram and I'll generate it for you
      </p>
    </div>
  );
}

const UserBubble = memo<{ message: { id: string; parts: Array<{ type: string; text?: string }> } }>(
  ({ message }) => (
    <div className="max-w-[85%] rounded-2xl rounded-br-sm bg-primary px-3 py-2 text-primary-foreground">
      {message.parts.map((part, i) =>
        part.type === "text" ? (
          <p key={`${message.id}-${i}`} className="text-sm whitespace-pre-wrap">
            {part.text}
          </p>
        ) : null,
      )}
    </div>
  ),
);
UserBubble.displayName = "UserBubble";

const SUGGESTION_PATTERN = /^(Note|Consider|Suggestion|Tip):/m;

/** Wrap lines matching suggestion patterns with styled container */
function renderWithSuggestions(text: string, messageId: string, index: number) {
  if (!SUGGESTION_PATTERN.test(text)) {
    return (
      <MemoizedMarkdown
        key={`${messageId}-${index}`}
        content={text}
        id={`copilot-${messageId}-${index}`}
      />
    );
  }

  const lines = text.split("\n");
  const segments: Array<{ isSuggestion: boolean; content: string }> = [];
  let currentSegment = { isSuggestion: false, content: "" };

  for (const line of lines) {
    const isSuggestionLine = /^(Note|Consider|Suggestion|Tip):/.test(line.trim());
    if (isSuggestionLine !== currentSegment.isSuggestion && currentSegment.content) {
      segments.push({ ...currentSegment });
      currentSegment = { isSuggestion: isSuggestionLine, content: line };
    } else {
      currentSegment.content += (currentSegment.content ? "\n" : "") + line;
    }
  }
  if (currentSegment.content) segments.push(currentSegment);

  return segments.map((seg, si) =>
    seg.isSuggestion ? (
      <div
        key={`${messageId}-${index}-sug-${si}`}
        className="my-1.5 flex items-start gap-2 rounded-md border-l-2 border-[var(--ai-accent)] bg-muted/30 px-3 py-2"
      >
        <Icons.Lightbulb className="mt-0.5 size-3.5 shrink-0 text-[var(--ai-accent)]" />
        <div className="min-w-0 flex-1">
          <MemoizedMarkdown
            content={seg.content}
            id={`copilot-${messageId}-${index}-sug-${si}`}
          />
        </div>
      </div>
    ) : (
      <MemoizedMarkdown
        key={`${messageId}-${index}-txt-${si}`}
        content={seg.content}
        id={`copilot-${messageId}-${index}-txt-${si}`}
      />
    ),
  );
}

const AssistantBubble = memo<{
  message: { id: string; parts: Array<{ type: string; text?: string; state?: string }> };
  isStreaming: boolean;
  diagramId: string;
  proposalStatus: ProposalStatus;
  lastProposalOutcome: "accepted" | "rejected" | null;
  changeSummary: string;
}>(({ message, isStreaming, diagramId, proposalStatus, lastProposalOutcome, changeSummary }) => {
  const hasToolResult = message.parts.some(
    (p) => isGenerateDiagramTool(p) && p.state === "output-available",
  );

  const handleAccept = useCallback(() => {
    acceptCurrentProposal(diagramId);
  }, [diagramId]);

  const handleReject = useCallback(() => {
    rejectCurrentProposal();
  }, []);

  return (
    <div className="max-w-[95%]">
      <Prose className="text-sm">
        {message.parts.map((part, i) =>
          part.type === "text" && part.text
            ? renderWithSuggestions(part.text, message.id, i)
            : null,
        )}
        {isStreaming && message.parts.length === 0 && (
          <span className="inline-block size-2 animate-pulse rounded-full bg-muted-foreground/50" />
        )}
      </Prose>
      {hasToolResult && proposalStatus === "pending" && (
        <div className="mt-2 flex items-center gap-2">
          <span className="text-xs text-muted-foreground">{changeSummary}</span>
          <Button size="sm" onClick={handleAccept}>
            <Icons.Check className="mr-1 size-3" /> Accept
          </Button>
          <Button size="sm" variant="ghost" onClick={handleReject}>
            <Icons.X className="mr-1 size-3" /> Reject
          </Button>
        </div>
      )}
      {hasToolResult && proposalStatus !== "pending" && lastProposalOutcome === "rejected" && (
        <div className="mt-1 flex items-center gap-1.5 text-xs text-muted-foreground">
          <Icons.X className="size-3 text-red-500" />
          Changes discarded
        </div>
      )}
      {hasToolResult && proposalStatus !== "pending" && lastProposalOutcome !== "rejected" && (
        <div className="mt-1 flex items-center gap-1.5 text-xs text-muted-foreground">
          <Icons.Check className="size-3 text-green-500" />
          Diagram updated
        </div>
      )}
    </div>
  );
});
AssistantBubble.displayName = "AssistantBubble";