1  // Copyright (c) 2026 VPL Solutions. All rights reserved.
  2  // Licensed under the MIT License. See LICENSE for details.
  3  //
  4  // DagTrace — Interactive animated Planner → Retriever → Synthesizer visualization
  5  // ADR-011: Makes the aiPolaris DAG execution tangible and auditable in real time.
  6  //
  7  // Design:
  8  //   - Three node cards animate in as each stage completes (not shown upfront)
  9  //   - Connecting arrows draw left-to-right as stages complete
 10  //   - Clicking a node expands a drawer showing that stage's output
 11  //   - Nodes pulse amber while running, green on done, red on error
 12  //   - Error nodes show an actionable recovery hint
 13  //
 14  // Trace data is memory-only — never written to localStorage. ADR-011.
 15  
 16  import { useState } from 'react';
 17  import {
 18    BrainCircuit, Search, PenLine, ChevronDown, ChevronUp,
 19    AlertCircle, FileText, Hash, Clock, CheckCircle2, XCircle,
 20    Loader2, ArrowRight,
 21  } from 'lucide-react';
 22  import type { DagNodeState, DagStage, PolarisCitation, PolarisChunk } from '../../api/aipolaris';
 23  
 24  // ── Stage config ──────────────────────────────────────────────────────────────
 25  
 26  interface StageConfig {
 27    icon: React.ComponentType<{ className?: string }>;
 28    accent: string;          // Tailwind color token base
 29    ringColor: string;
 30    bgIdle: string;
 31    bgRunning: string;
 32    bgDone: string;
 33    bgError: string;
 34    borderRunning: string;
 35    borderDone: string;
 36    borderError: string;
 37    textRunning: string;
 38    textDone: string;
 39  }
 40  
 41  const STAGE_CONFIG: Record<DagStage, StageConfig> = {
 42    planner: {
 43      icon: BrainCircuit,
 44      accent: 'purple',
 45      ringColor:     'ring-purple-400 dark:ring-purple-500',
 46      bgIdle:        'bg-gray-50 dark:bg-gray-800/40',
 47      bgRunning:     'bg-purple-50 dark:bg-purple-900/20',
 48      bgDone:        'bg-white dark:bg-white/[0.03]',
 49      bgError:       'bg-red-50 dark:bg-red-900/20',
 50      borderRunning: 'border-purple-300 dark:border-purple-600',
 51      borderDone:    'border-purple-200 dark:border-purple-800',
 52      borderError:   'border-red-300 dark:border-red-700',
 53      textRunning:   'text-purple-600 dark:text-purple-400',
 54      textDone:      'text-purple-700 dark:text-purple-300',
 55    },
 56    retriever: {
 57      icon: Search,
 58      accent: 'blue',
 59      ringColor:     'ring-blue-400 dark:ring-blue-500',
 60      bgIdle:        'bg-gray-50 dark:bg-gray-800/40',
 61      bgRunning:     'bg-blue-50 dark:bg-blue-900/20',
 62      bgDone:        'bg-white dark:bg-white/[0.03]',
 63      bgError:       'bg-red-50 dark:bg-red-900/20',
 64      borderRunning: 'border-blue-300 dark:border-blue-600',
 65      borderDone:    'border-blue-200 dark:border-blue-800',
 66      borderError:   'border-red-300 dark:border-red-700',
 67      textRunning:   'text-blue-600 dark:text-blue-400',
 68      textDone:      'text-blue-700 dark:text-blue-300',
 69    },
 70    synthesizer: {
 71      icon: PenLine,
 72      accent: 'emerald',
 73      ringColor:     'ring-emerald-400 dark:ring-emerald-500',
 74      bgIdle:        'bg-gray-50 dark:bg-gray-800/40',
 75      bgRunning:     'bg-emerald-50 dark:bg-emerald-900/20',
 76      bgDone:        'bg-white dark:bg-white/[0.03]',
 77      bgError:       'bg-red-50 dark:bg-red-900/20',
 78      borderRunning: 'border-emerald-300 dark:border-emerald-600',
 79      borderDone:    'border-emerald-200 dark:border-emerald-800',
 80      borderError:   'border-red-300 dark:border-red-700',
 81      textRunning:   'text-emerald-600 dark:text-emerald-400',
 82      textDone:      'text-emerald-700 dark:text-emerald-300',
 83    },
 84  };
 85  
 86  // ── Drawer content per stage ──────────────────────────────────────────────────
 87  
 88  function PlannerDrawer({ subTasks }: { subTasks?: string[] }) {
 89    if (!subTasks?.length) return (
 90      <p className="text-xs text-gray-400 dark:text-gray-500 italic">No sub-tasks recorded.</p>
 91    );
 92    return (
 93      <div className="space-y-1.5">
 94        <p className="text-[10px] font-semibold uppercase tracking-wider text-gray-400 dark:text-gray-500 mb-2">
 95          Decomposed into {subTasks.length} sub-task{subTasks.length !== 1 ? 's' : ''}
 96        </p>
 97        {subTasks.map((task, i) => (
 98          <div key={i} className="flex items-start gap-2">
 99            <span className="shrink-0 mt-0.5 w-4 h-4 rounded-full bg-purple-100 dark:bg-purple-900/40 text-purple-600 dark:text-purple-400 text-[9px] font-bold flex items-center justify-center">
100              {i + 1}
101            </span>
102            <span className="text-xs text-gray-700 dark:text-gray-300 leading-snug">{task}</span>
103          </div>
104        ))}
105      </div>
106    );
107  }
108  
109  function RetrieverDrawer({ chunks }: { chunks?: PolarisChunk[] }) {
110    if (!chunks?.length) return (
111      <p className="text-xs text-gray-400 dark:text-gray-500 italic">No chunks retrieved.</p>
112    );
113    return (
114      <div className="space-y-2.5">
115        <p className="text-[10px] font-semibold uppercase tracking-wider text-gray-400 dark:text-gray-500 mb-2">
116          {chunks.length} chunk{chunks.length !== 1 ? 's' : ''} retrieved
117        </p>
118        {chunks.map((chunk, i) => (
119          <div key={i} className="rounded-lg border border-gray-100 dark:border-white/5 bg-gray-50 dark:bg-white/[0.02] p-2.5 space-y-1.5">
120            <div className="flex items-start justify-between gap-2">
121              <div className="flex items-center gap-1.5 min-w-0">
122                <FileText className="w-3 h-3 text-blue-400 shrink-0" />
123                <span className="text-xs font-medium text-gray-800 dark:text-gray-200 truncate">{chunk.title}</span>
124              </div>
125              {/* Reranker score bar */}
126              <div className="flex items-center gap-1.5 shrink-0">
127                <div className="w-16 h-1.5 rounded-full bg-gray-200 dark:bg-gray-700 overflow-hidden">
128                  <div
129                    className="h-full rounded-full bg-blue-400 dark:bg-blue-500 transition-all duration-500"
130                    style={{ width: `${Math.round(chunk.reranker_score * 100)}%` }}
131                  />
132                </div>
133                <span className="text-[10px] text-gray-400 dark:text-gray-500 font-mono w-7 text-right">
134                  {(chunk.reranker_score * 100).toFixed(0)}%
135                </span>
136              </div>
137            </div>
138            <p className="text-[11px] text-gray-500 dark:text-gray-400 leading-snug line-clamp-2">
139              {chunk.content}
140            </p>
141            <div className="flex items-center gap-1 text-[10px] text-gray-400 dark:text-gray-500">
142              <Hash className="w-2.5 h-2.5" />
143              <span className="font-mono truncate">{chunk.source}</span>
144            </div>
145          </div>
146        ))}
147      </div>
148    );
149  }
150  
151  function SynthesizerDrawer({ citations }: { citations?: PolarisCitation[] }) {
152    if (!citations?.length) return (
153      <p className="text-xs text-gray-400 dark:text-gray-500 italic">No citations available.</p>
154    );
155    return (
156      <div className="space-y-2">
157        <p className="text-[10px] font-semibold uppercase tracking-wider text-gray-400 dark:text-gray-500 mb-2">
158          {citations.length} citation{citations.length !== 1 ? 's' : ''} used
159        </p>
160        {citations.map((citation, i) => (
161          <div key={i} className="rounded-lg border border-gray-100 dark:border-white/5 bg-gray-50 dark:bg-white/[0.02] p-2.5 space-y-1">
162            <div className="flex items-center gap-1.5">
163              <CheckCircle2 className="w-3 h-3 text-emerald-400 shrink-0" />
164              <span className="text-xs font-medium text-gray-800 dark:text-gray-200">{citation.title}</span>
165            </div>
166            <p className="text-[11px] text-gray-500 dark:text-gray-400 leading-snug italic line-clamp-3 pl-4">
167              "{citation.excerpt}"
168            </p>
169            <div className="flex items-center gap-1 text-[10px] text-gray-400 dark:text-gray-500 pl-4">
170              <Hash className="w-2.5 h-2.5" />
171              <span className="font-mono truncate">{citation.source}</span>
172            </div>
173          </div>
174        ))}
175      </div>
176    );
177  }
178  
179  // ── Node card ─────────────────────────────────────────────────────────────────
180  
181  interface NodeCardProps {
182    node: DagNodeState;
183    isExpanded: boolean;
184    onToggle: () => void;
185  }
186  
187  function NodeCard({ node, isExpanded, onToggle }: NodeCardProps) {
188    const cfg = STAGE_CONFIG[node.stage];
189    const Icon = cfg.icon;
190    const isClickable = node.status === 'done' || node.status === 'error';
191  
192    const bg = node.status === 'idle'    ? cfg.bgIdle
193             : node.status === 'running' ? cfg.bgRunning
194             : node.status === 'done'    ? cfg.bgDone
195                                         : cfg.bgError;
196  
197    const border = node.status === 'running' ? cfg.borderRunning
198                 : node.status === 'done'    ? cfg.borderDone
199                 : node.status === 'error'   ? cfg.borderError
200                                             : 'border-gray-100 dark:border-white/5';
201  
202    return (
203      <div className={[
204        'flex flex-col rounded-xl border transition-all duration-300 overflow-hidden',
205        bg, border,
206        node.status === 'running' ? 'shadow-sm ring-1 ring-offset-0 ' + cfg.ringColor : '',
207        isClickable ? 'cursor-pointer' : 'cursor-default',
208      ].join(' ')}>
209        {/* Card header */}
210        <div
211          className="flex items-center gap-2.5 px-3 py-2.5"
212          onClick={isClickable ? onToggle : undefined}
213          role={isClickable ? 'button' : undefined}
214          tabIndex={isClickable ? 0 : undefined}
215          onKeyDown={isClickable ? (e) => { if (e.key === 'Enter' || e.key === ' ') { e.preventDefault(); onToggle(); } } : undefined}
216          aria-expanded={isClickable ? isExpanded : undefined}
217          aria-label={isClickable ? `${node.stage.charAt(0).toUpperCase() + node.stage.slice(1)} stage — ${node.status}. ${isExpanded ? 'Collapse' : 'Expand'} details` : undefined}
218        >
219          {/* Status icon */}
220          <div className={[
221            'flex items-center justify-center w-7 h-7 rounded-lg transition-all duration-200',
222            node.status === 'idle'    ? 'bg-gray-100 dark:bg-gray-700/50' :
223            node.status === 'running' ? `bg-${cfg.accent}-100 dark:bg-${cfg.accent}-900/30` :
224            node.status === 'done'    ? `bg-${cfg.accent}-50 dark:bg-${cfg.accent}-900/20` :
225                                        'bg-red-100 dark:bg-red-900/30',
226          ].join(' ')}>
227            {node.status === 'running' ? (
228              <Loader2 className={`w-3.5 h-3.5 animate-spin ${cfg.textRunning}`} />
229            ) : node.status === 'done' ? (
230              <Icon className={`w-3.5 h-3.5 ${cfg.textDone}`} />
231            ) : node.status === 'error' ? (
232              <XCircle className="w-3.5 h-3.5 text-red-500" />
233            ) : (
234              <Icon className="w-3.5 h-3.5 text-gray-300 dark:text-gray-500" />
235            )}
236          </div>
237  
238          {/* Label + sublabel */}
239          <div className="flex-1 min-w-0">
240            <div className="flex items-center gap-1.5">
241              <span className={[
242                'text-xs font-semibold',
243                node.status === 'idle'  ? 'text-gray-400 dark:text-gray-500' :
244                node.status === 'error' ? 'text-red-600 dark:text-red-400' :
245                                          cfg.textDone,
246              ].join(' ')}>
247                {node.label}
248              </span>
249              {node.status === 'running' && (
250                <span className={`text-[10px] ${cfg.textRunning} animate-pulse`}>
251                  {node.sublabel}
252                </span>
253              )}
254            </div>
255            {node.status === 'done' && (
256              <div className="flex items-center gap-2 mt-0.5">
257                <span className="text-[10px] text-gray-400 dark:text-gray-500">{node.sublabel}</span>
258                {node.latency_ms !== null && (
259                  <span className="flex items-center gap-0.5 text-[10px] text-gray-300 dark:text-gray-500">
260                    <Clock className="w-2.5 h-2.5" />
261                    {Math.round(node.latency_ms)}ms
262                  </span>
263                )}
264              </div>
265            )}
266          </div>
267  
268          {/* Expand chevron */}
269          {isClickable && (
270            <div className="shrink-0 text-gray-300 dark:text-gray-500">
271              {isExpanded
272                ? <ChevronUp className="w-3.5 h-3.5" />
273                : <ChevronDown className="w-3.5 h-3.5" />
274              }
275            </div>
276          )}
277        </div>
278  
279        {/* Error hint */}
280        {node.status === 'error' && node.errorHint && (
281          <div className="flex items-start gap-2 px-3 pb-2.5 text-xs text-red-500 dark:text-red-400">
282            <AlertCircle className="w-3.5 h-3.5 mt-0.5 shrink-0" />
283            <span>{node.errorHint}</span>
284          </div>
285        )}
286  
287        {/* Expandable drawer */}
288        {isExpanded && node.status === 'done' && (
289          <div className="border-t border-gray-100 dark:border-white/5 px-3 py-3 animate-in slide-in-from-top-1 duration-200">
290            {node.stage === 'planner'     && <PlannerDrawer subTasks={node.subTasks} />}
291            {node.stage === 'retriever'   && <RetrieverDrawer chunks={node.chunks} />}
292            {node.stage === 'synthesizer' && <SynthesizerDrawer citations={node.citations} />}
293          </div>
294        )}
295      </div>
296    );
297  }
298  
299  // ── Connector arrow ───────────────────────────────────────────────────────────
300  
301  function ConnectorArrow({ active }: { active: boolean }) {
302    return (
303      <div className="flex items-center justify-center self-start mt-[22px] mx-1 shrink-0">
304        <div className={[
305          'h-px w-5 transition-all duration-500',
306          active ? 'bg-gray-300 dark:bg-gray-600' : 'bg-gray-150 dark:bg-gray-800',
307        ].join(' ')} />
308        <ArrowRight className={[
309          'w-3 h-3 -ml-1 transition-colors duration-500',
310          active ? 'text-gray-300 dark:text-gray-500' : 'text-gray-150 dark:text-gray-800',
311        ].join(' ')} />
312      </div>
313    );
314  }
315  
316  // ── Main DagTrace component ───────────────────────────────────────────────────
317  
318  interface DagTraceProps {
319    nodes: DagNodeState[];
320    traceId?: string;
321  }
322  
323  export function DagTrace({ nodes, traceId }: DagTraceProps) {
324    const [expandedStage, setExpandedStage] = useState<DagStage | null>(null);
325  
326    // Only render once at least one node is no longer idle
327    const hasStarted = nodes.some(n => n.status !== 'idle');
328    if (!hasStarted) return null;
329  
330    const toggleExpand = (stage: DagStage) => {
331      setExpandedStage(prev => prev === stage ? null : stage);
332    };
333  
334    const connectorActive = (idx: number) =>
335      nodes[idx].status === 'done' || nodes[idx].status === 'error';
336  
337    return (
338      <div className="rounded-xl border border-gray-100 dark:border-white/[0.06] bg-gray-50/50 dark:bg-white/[0.01] p-3 space-y-2">
339        {/* Header */}
340        <div className="flex items-center justify-between">
341          <div className="flex items-center gap-1.5">
342            <div className="w-1.5 h-1.5 rounded-full bg-purple-400 dark:bg-purple-500 animate-pulse" />
343            <span className="text-[10px] font-semibold uppercase tracking-wider text-gray-400 dark:text-gray-500">
344              Execution trace
345            </span>
346          </div>
347          {traceId && (
348            <span className="font-mono text-[10px] text-gray-300 dark:text-gray-700">
349              {traceId.slice(0, 8)}
350            </span>
351          )}
352        </div>
353  
354        {/* DAG pipeline */}
355        <div className="flex items-start">
356          {nodes.map((node, idx) => (
357            <div key={node.stage} className="flex items-start flex-1 min-w-0">
358              <div className="flex-1 min-w-0">
359                <NodeCard
360                  node={node}
361                  isExpanded={expandedStage === node.stage}
362                  onToggle={() => toggleExpand(node.stage)}
363                />
364              </div>
365              {idx < nodes.length - 1 && (
366                <ConnectorArrow active={connectorActive(idx)} />
367              )}
368            </div>
369          ))}
370        </div>
371  
372        {/* Footer hint */}
373        <p className="text-[10px] text-gray-300 dark:text-gray-700 text-center">
374          Click a completed node to inspect its output
375        </p>
376      </div>
377    );
378  }