Codewalkers/apps/web/src/components/AgentOutputViewer.tsx

import { useEffect, useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { ArrowDown, Pause, Play, AlertCircle, Square, ChevronRight, ChevronDown, CheckCircle2, Loader2, Circle } from "lucide-react";
import { trpc } from "@/lib/trpc";
import { useSubscriptionWithErrorHandling } from "@/hooks";
import {
  type ParsedMessage,
  type TimestampedChunk,
  getMessageStyling,
  parseAgentOutput,
} from "@/lib/parse-agent-output";

interface AgentOutputViewerProps {
  agentId: string;
  agentName?: string;
  status?: string;
  onStop?: (id: string) => void;
}

export function AgentOutputViewer({ agentId, agentName, status, onStop }: AgentOutputViewerProps) {
  const [messages, setMessages] = useState<ParsedMessage[]>([]);
  const [follow, setFollow] = useState(true);
  const [expandedResults, setExpandedResults] = useState<Set<number>>(new Set());
  const containerRef = useRef<HTMLDivElement>(null);
  // Accumulate timestamped chunks: initial query data + live subscription chunks
  const chunksRef = useRef<TimestampedChunk[]>([]);

  // Load initial/historical output
  const outputQuery = trpc.getAgentOutput.useQuery(
    { id: agentId },
    {
      refetchOnWindowFocus: false,
    }
  );

  // Subscribe to live output with error handling
  const subscription = useSubscriptionWithErrorHandling(
    () => trpc.onAgentOutput.useSubscription({ agentId }),
    {
      onData: (event: any) => {
        // TrackedEnvelope shape: { id, data: { agentId, data: string } }
        const raw = event?.data?.data ?? event?.data;
        const chunk = typeof raw === 'string' ? raw : JSON.stringify(raw);
        chunksRef.current = [...chunksRef.current, { content: chunk, createdAt: new Date().toISOString() }];
        setMessages(parseAgentOutput(chunksRef.current));
      },
      onError: (error) => {
        console.error('Agent output subscription error:', error);
      },
      autoReconnect: true,
      maxReconnectAttempts: 3,
    }
  );

  // Set initial output when query loads
  useEffect(() => {
    if (outputQuery.data) {
      chunksRef.current = outputQuery.data;
      setMessages(parseAgentOutput(outputQuery.data));
    }
  }, [outputQuery.data]);

  // Reset output when agent changes
  useEffect(() => {
    chunksRef.current = [];
    setMessages([]);
    setFollow(true);
    setExpandedResults(new Set());
  }, [agentId]);

  // Auto-scroll to bottom when following
  useEffect(() => {
    if (follow && containerRef.current) {
      containerRef.current.scrollTop = containerRef.current.scrollHeight;
    }
  }, [messages, follow]);

  // Handle scroll to detect user scrolling up
  function handleScroll() {
    if (!containerRef.current) return;
    const { scrollTop, scrollHeight, clientHeight } = containerRef.current;
    const isAtBottom = scrollHeight - scrollTop - clientHeight < 50;
    if (!isAtBottom && follow) {
      setFollow(false);
    }
  }

  // Jump to bottom
  function scrollToBottom() {
    if (containerRef.current) {
      containerRef.current.scrollTop = containerRef.current.scrollHeight;
      setFollow(true);
    }
  }

  const isLoading = outputQuery.isLoading;
  const hasOutput = messages.length > 0;

  const lastTodoWrite = [...messages]
    .reverse()
    .find(m => m.type === "tool_call" && m.meta?.toolName === "TodoWrite");

  const currentTodos = lastTodoWrite
    ? (lastTodoWrite.meta?.toolInput as { todos: Array<{ content: string; status: string; activeForm: string }> } | undefined)?.todos ?? []
    : [];

  return (
    <div className="flex flex-col h-full rounded-lg border border-terminal-border overflow-hidden">
      {/* Header */}
      <div className="flex items-center justify-between border-b border-terminal-border bg-terminal px-4 py-2">
        <div className="flex items-center gap-2">
          <span className="text-sm font-medium text-terminal-fg font-mono">
            {agentName ? `Output: ${agentName}` : "Agent Output"}
          </span>
          {subscription.error && (
            <div className="flex items-center gap-1 text-terminal-error" title={subscription.error.message}>
              <AlertCircle className="h-3 w-3" />
              <span className="text-xs">Connection error</span>
            </div>
          )}
          {subscription.isConnecting && (
            <span className="text-xs text-terminal-warning">Connecting...</span>
          )}
        </div>
        <div className="flex items-center gap-2">
          {onStop && (status === "running" || status === "waiting_for_input") && (
            <Button
              variant="destructive"
              size="sm"
              onClick={() => onStop(agentId)}
              className="h-7"
            >
              <Square className="mr-1 h-3 w-3" />
              Stop
            </Button>
          )}
          <Button
            variant="ghost"
            size="sm"
            onClick={() => setFollow(!follow)}
            className="h-7 text-terminal-muted hover:text-terminal-fg hover:bg-white/5"
          >
            {follow ? (
              <>
                <Pause className="mr-1 h-3 w-3" />
                Following
              </>
            ) : (
              <>
                <Play className="mr-1 h-3 w-3" />
                Paused
              </>
            )}
          </Button>
          {!follow && (
            <Button
              variant="ghost"
              size="sm"
              onClick={scrollToBottom}
              className="h-7 text-terminal-muted hover:text-terminal-fg hover:bg-white/5"
            >
              <ArrowDown className="mr-1 h-3 w-3" />
              Jump to bottom
            </Button>
          )}
        </div>
      </div>

      {/* Todo strip */}
      {currentTodos.length > 0 && (
        <div className="bg-terminal border-b border-terminal-border px-4 py-2">
          <div className="text-[10px] text-terminal-muted/60 font-mono uppercase tracking-widest mb-1">TASKS</div>
          {currentTodos.slice(0, 5).map((todo, i) => (
            <div key={i} className="flex items-center gap-2 font-mono text-xs">
              {todo.status === "completed" && <CheckCircle2 className="h-3 w-3 text-terminal-muted/60" />}
              {todo.status === "in_progress" && <Loader2 className="h-3 w-3 text-terminal-fg animate-spin" />}
              {todo.status === "pending" && <Circle className="h-3 w-3 text-terminal-muted/40" />}
              <span className={
                todo.status === "completed"
                  ? "text-terminal-muted/60 line-through"
                  : todo.status === "in_progress"
                  ? "text-terminal-fg"
                  : "text-terminal-muted"
              }>{todo.content}</span>
            </div>
          ))}
          {currentTodos.length > 5 && (
            <div className="font-mono text-xs text-terminal-muted/60">+ {currentTodos.length - 5} more</div>
          )}
        </div>
      )}

      {/* Output content */}
      <div
        ref={containerRef}
        onScroll={handleScroll}
        className="flex-1 overflow-y-auto overflow-x-hidden bg-terminal p-4"
      >
        {isLoading ? (
          <div className="text-terminal-muted text-sm">Loading output...</div>
        ) : !hasOutput ? (
          <div className="text-terminal-muted text-sm">No output yet...</div>
        ) : (
          <div className="space-y-2 min-w-0">
            {messages.map((message, index) => (
              <div key={index} className={getMessageStyling(message.type)}>
                {message.type === 'system' && (
                  <div className="text-terminal-muted/60 text-xs font-mono">
                    {message.content}
                  </div>
                )}

                {message.type === 'text' && (
                  <>
                    <Timestamp date={message.timestamp} />
                    <div className="font-mono text-sm whitespace-pre-wrap break-words text-terminal-fg">
                      {message.content}
                    </div>
                  </>
                )}

                {message.type === 'tool_call' && (
                  <div className="border-l-2 border-terminal-tool pl-3 py-1 min-w-0">
                    <div className="flex items-center gap-2 mb-1">
                      <Badge variant="default" className="text-xs">
                        {message.meta?.toolName}
                      </Badge>
                      <Timestamp date={message.timestamp} />
                    </div>
                    <div className="font-mono text-xs text-terminal-muted whitespace-pre-wrap break-words">
                      {message.content}
                    </div>
                  </div>
                )}

                {message.type === 'tool_result' && (
                  <div
                    className="border-l-2 border-terminal-result pl-3 py-1 bg-white/[0.02] min-w-0 cursor-pointer"
                    onClick={() => setExpandedResults(prev => {
                      const next = new Set(prev);
                      if (next.has(index)) next.delete(index); else next.add(index);
                      return next;
                    })}
                  >
                    {expandedResults.has(index) ? (
                      <>
                        <ChevronDown className="h-4 w-4 text-terminal-muted inline-block shrink-0 mr-1" />
                        <Badge variant="outline" className="mb-1 text-xs text-terminal-result border-terminal-result">
                          Result
                        </Badge>
                        <div className="font-mono text-xs text-terminal-muted whitespace-pre-wrap break-words">
                          {message.content}
                        </div>
                      </>
                    ) : (
                      <>
                        <ChevronRight className="h-4 w-4 text-terminal-muted inline-block shrink-0 mr-1" />
                        <span className="text-terminal-muted/60 text-xs font-mono truncate">
                          {message.meta?.toolName === "Task"
                            ? `${(message.meta.toolInput as any)?.subagent_type ?? "Subagent"} result`
                            : message.content.substring(0, 80)}
                        </span>
                      </>
                    )}
                  </div>
                )}

                {message.type === 'error' && (
                  <div className="border-l-2 border-terminal-error pl-3 py-1 bg-terminal-error/10 min-w-0">
                    <Badge variant="destructive" className="mb-1 text-xs">
                      Error
                    </Badge>
                    <div className="font-mono text-xs text-terminal-error whitespace-pre-wrap break-words">
                      {message.content}
                    </div>
                  </div>
                )}

                {message.type === 'session_end' && (
                  <div className="border-t border-terminal-border pt-2 mt-4">
                    <div className="flex items-center gap-2">
                      <Badge variant={message.meta?.isError ? "destructive" : "default"} className="text-xs">
                        {message.content}
                      </Badge>
                      {message.meta?.cost && (
                        <span className="text-xs text-terminal-muted">${message.meta.cost.toFixed(4)}</span>
                      )}
                      {message.meta?.duration && (
                        <span className="text-xs text-terminal-muted">{(message.meta.duration / 1000).toFixed(1)}s</span>
                      )}
                      <Timestamp date={message.timestamp} />
                    </div>
                  </div>
                )}
              </div>
            ))}
          </div>
        )}
      </div>
    </div>
  );
}

function formatTime(date: Date): string {
  return date.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit', second: '2-digit', hour12: false });
}

function Timestamp({ date }: { date?: Date }) {
  if (!date) return null;
  return (
    <span className="shrink-0 text-[10px] text-terminal-muted/60 font-mono tabular-nums">
      {formatTime(date)}
    </span>
  );
}