frontend/components/reasoning-visualizer.tsx

"use client";

import { useEffect, useState } from "react";

type ReasoningStep = {
  step: string;
  status: "pending" | "running" | "completed" | "error";
  message?: string;
  details?: Record<string, any>;
  timestamp?: number;
};

type ReasoningVisualizerProps = {
  reasoningTrace?: Array<Record<string, any>>;
  isActive?: boolean;
  onComplete?: () => void;
};

const STEP_ICONS: Record<string, string> = {
  request_received: "📥",
  admin_rules_check: "🛡️",
  intent_detection: "🧠",
  rag_prefetch: "📚",
  tool_scoring: "📊",
  tool_selection: "🎯",
  tool_execution: "⚙️",
  llm_response: "💬",
  result_merger: "🔀",
  parallel_execution: "⚡",
  error: "❌",
};

const STEP_LABELS: Record<string, string> = {
  request_received: "Request Received",
  admin_rules_check: "Checking Admin Rules",
  intent_detection: "Detecting Intent",
  rag_prefetch: "Pre-fetching RAG Results",
  tool_scoring: "Scoring Tools",
  tool_selection: "Selecting Tools",
  tool_execution: "Executing Tools",
  llm_response: "Generating Response",
  result_merger: "Merging Results",
  parallel_execution: "Parallel Execution",
  error: "Error",
};

export function ReasoningVisualizer({
  reasoningTrace = [],
  isActive = false,
  onComplete,
}: ReasoningVisualizerProps) {
  const [steps, setSteps] = useState<ReasoningStep[]>([]);
  const [currentStepIndex, setCurrentStepIndex] = useState(0);

  useEffect(() => {
    if (!reasoningTrace || reasoningTrace.length === 0) {
      setSteps([]);
      setCurrentStepIndex(0);
      return;
    }

    // Convert reasoning trace to visual steps
    const visualSteps: ReasoningStep[] = reasoningTrace.map((trace, idx) => {
      const stepName = trace.step || "unknown";
      const icon = STEP_ICONS[stepName] || "⚙️";
      const label = STEP_LABELS[stepName] || stepName.replace(/_/g, " ");

      // Build message from trace data
      let message = label;
      const details: Record<string, any> = {};

      if (stepName === "admin_rules_check") {
        const matchCount = trace.match_count || 0;
        message = matchCount > 0 
          ? `Found ${matchCount} rule violation(s)` 
          : "No violations found";
        details.matches = trace.matches || [];
      } else if (stepName === "intent_detection") {
        message = `Intent: ${trace.intent || "unknown"}`;
        details.intent = trace.intent;
      } else if (stepName === "rag_prefetch") {
        const hitCount = trace.hit_count || 0;
        message = hitCount > 0 
          ? `Found ${hitCount} relevant document(s)` 
          : "No documents found";
        details.hit_count = hitCount;
        details.latency_ms = trace.latency_ms;
      } else if (stepName === "tool_selection") {
        const decision = trace.decision;
        if (decision) {
          message = `Selected: ${decision.tool || "llm"} (${decision.action})`;
          details.decision = decision;
        }
      } else if (stepName === "tool_execution") {
        const tool = trace.tool || "unknown";
        const hitCount = trace.hit_count || 0;
        message = `${tool.toUpperCase()}: ${hitCount} result(s)`;
        details.tool = tool;
        details.hit_count = hitCount;
      } else if (stepName === "result_merger") {
        const mergedItems = trace.merged_items || 0;
        message = `Merged ${mergedItems} result(s)`;
        details.merged_items = mergedItems;
        details.sources = trace.sources || [];
      } else if (stepName === "llm_response") {
        message = "Generating final response";
        details.latency_ms = trace.latency_ms;
        details.estimated_tokens = trace.estimated_tokens;
      }

      return {
        step: stepName,
        status: idx < currentStepIndex ? "completed" : idx === currentStepIndex ? "running" : "pending",
        message,
        details,
        timestamp: Date.now(),
      };
    });

    setSteps(visualSteps);

    // Animate through steps if active
    if (isActive && visualSteps.length > 0) {
      const interval = setInterval(() => {
        setCurrentStepIndex((prev) => {
          if (prev < visualSteps.length - 1) {
            return prev + 1;
          } else {
            clearInterval(interval);
            if (onComplete) onComplete();
            return prev;
          }
        });
      }, 800); // 800ms per step

      return () => clearInterval(interval);
    } else if (!isActive && visualSteps.length > 0) {
      // Show all steps as completed if not active
      setCurrentStepIndex(visualSteps.length);
    }
  }, [reasoningTrace, isActive, currentStepIndex, onComplete]);

  if (steps.length === 0) {
    return (
      <div className="rounded-2xl border border-white/10 bg-slate-950/40 p-6">
        <p className="text-sm text-slate-400 text-center">
          No reasoning trace available. Send a message to see the agent's reasoning path.
        </p>
      </div>
    );
  }

  return (
    <div className="rounded-2xl border border-white/10 bg-slate-950/40 p-6">
      <div className="mb-4 flex items-center justify-between">
        <h3 className="text-lg font-semibold text-white">Real-Time Reasoning Path</h3>
        <span className="text-xs text-slate-400">{steps.length} steps</span>
      </div>

      <div className="space-y-3">
        {steps.map((step, idx) => {
          const isCompleted = step.status === "completed";
          const isRunning = step.status === "running";
          const isPending = step.status === "pending";

          return (
            <div
              key={idx}
              className={`relative flex items-start gap-4 rounded-xl border p-4 transition-all ${
                isRunning
                  ? "border-cyan-500/50 bg-cyan-500/10 shadow-lg shadow-cyan-500/20"
                  : isCompleted
                    ? "border-emerald-500/30 bg-emerald-500/5"
                    : "border-white/5 bg-white/5 opacity-50"
              }`}
            >
              {/* Step number and icon */}
              <div
                className={`flex h-10 w-10 shrink-0 items-center justify-center rounded-full text-lg transition-all ${
                  isRunning
                    ? "bg-cyan-500 text-white animate-pulse"
                    : isCompleted
                      ? "bg-emerald-500 text-white"
                      : "bg-slate-700 text-slate-400"
                }`}
              >
                {isRunning ? (
                  <span className="animate-spin">⏳</span>
                ) : isCompleted ? (
                  "✓"
                ) : (
                  idx + 1
                )}
              </div>

              {/* Step content */}
              <div className="flex-1 min-w-0">
                <div className="flex items-center gap-2">
                  <span className="text-lg">{STEP_ICONS[step.step] || "⚙️"}</span>
                  <h4 className="font-semibold text-white">
                    {STEP_LABELS[step.step] || step.step.replace(/_/g, " ")}
                  </h4>
                  {isRunning && (
                    <span className="ml-auto text-xs text-cyan-300 animate-pulse">
                      Running...
                    </span>
                  )}
                </div>
                <p className="mt-1 text-sm text-slate-300">{step.message}</p>

                {/* Step details */}
                {step.details && Object.keys(step.details).length > 0 && isCompleted && (
                  <div className="mt-2 space-y-1 text-xs text-slate-400">
                    {step.details.latency_ms && (
                      <span>⏱️ {step.details.latency_ms}ms</span>
                    )}
                    {step.details.hit_count !== undefined && (
                      <span>📊 {step.details.hit_count} hits</span>
                    )}
                    {step.details.estimated_tokens && (
                      <span>🔢 ~{step.details.estimated_tokens} tokens</span>
                    )}
                    {step.details.score && (
                      <span>⭐ Score: {step.details.score.toFixed(2)}</span>
                    )}
                  </div>
                )}
              </div>

              {/* Connecting line */}
              {idx < steps.length - 1 && (
                <div
                  className={`absolute left-[29px] top-[50px] h-6 w-0.5 ${
                    isCompleted ? "bg-emerald-500/50" : "bg-slate-700"
                  }`}
                />
              )}
            </div>
          );
        })}
      </div>
    </div>
  );
}