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 } from "lucide-react";
import { trpc } from "@/lib/trpc";
import { useSubscriptionWithErrorHandling } from "@/hooks";
import {
  type ParsedMessage,
  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 containerRef = useRef<HTMLDivElement>(null);
  // Accumulate raw JSONL: initial query data + live subscription chunks
  const rawBufferRef = useRef<string>('');

  // 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);
        rawBufferRef.current += chunk;
        setMessages(parseAgentOutput(rawBufferRef.current));
      },
      onError: (error) => {
        console.error('Agent output subscription error:', error);
      },
      autoReconnect: true,
      maxReconnectAttempts: 3,
    }
  );

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

  // Reset output when agent changes
  useEffect(() => {
    rawBufferRef.current = '';
    setMessages([]);
    setFollow(true);
  }, [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;

  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>

      {/* Output content */}
      <div
        ref={containerRef}
        onScroll={handleScroll}
        className="flex-1 overflow-y-auto 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">
            {messages.map((message, index) => (
              <div key={index} className={getMessageStyling(message.type)}>
                {message.type === 'system' && (
                  <div className="flex items-center gap-2">
                    <Badge variant="secondary" className="text-xs bg-terminal-border text-terminal-system">System</Badge>
                    <span className="text-xs text-terminal-muted">{message.content}</span>
                  </div>
                )}

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

                {message.type === 'tool_call' && (
                  <div className="border-l-2 border-terminal-tool pl-3 py-1">
                    <Badge variant="default" className="mb-1 text-xs">
                      {message.meta?.toolName}
                    </Badge>
                    <div className="font-mono text-xs text-terminal-muted whitespace-pre-wrap">
                      {message.content}
                    </div>
                  </div>
                )}

                {message.type === 'tool_result' && (
                  <div className="border-l-2 border-terminal-result pl-3 py-1 bg-white/[0.02]">
                    <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">
                      {message.content}
                    </div>
                  </div>
                )}

                {message.type === 'error' && (
                  <div className="border-l-2 border-terminal-error pl-3 py-1 bg-terminal-error/10">
                    <Badge variant="destructive" className="mb-1 text-xs">
                      Error
                    </Badge>
                    <div className="font-mono text-xs text-terminal-error whitespace-pre-wrap">
                      {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>
                      )}
                    </div>
                  </div>
                )}
              </div>
            ))}
          </div>
        )}
      </div>
    </div>
  );
}