trustos/frontend/src/components/AttackPathVisualizer.tsx

"use client";
import { useRef, useEffect, useState } from "react";

export interface AttackNode {
  id: string;
  label: string;
  type: "attacker" | "entry_point" | "pivot" | "target";
  risk_level: "none" | "low" | "medium" | "high" | "critical";
}

export interface AttackEdge {
  source: string;
  target: string;
}

interface AttackPathVisualizerProps {
  nodes: AttackNode[];
  edges: AttackEdge[];
  narrative?: string;
}

const getNodeColors = (type: string) => {
  const colors: Record<string, { bg: string; border: string; text: string }> = {
    attacker: {
      bg: "bg-red-950/60",
      border: "border-red-700/60",
      text: "text-red-300",
    },
    entry_point: {
      bg: "bg-orange-950/60",
      border: "border-orange-700/60",
      text: "text-orange-300",
    },
    pivot: {
      bg: "bg-amber-950/60",
      border: "border-amber-700/60",
      text: "text-amber-300",
    },
    target: {
      bg: "bg-blue-950/60",
      border: "border-blue-700/60",
      text: "text-blue-300",
    },
  };
  return colors[type] || { bg: "bg-vault-dark", border: "border-vault-border", text: "text-vault-muted" };
};

const getRiskColor = (level: string) => {
  const colors: Record<string, string> = {
    none: "#6b7280",
    low: "#3b82f6",
    medium: "#f59e0b",
    high: "#ef4444",
    critical: "#dc2626",
  };
  return colors[level] || "#6b7280";
};

export default function AttackPathVisualizer({
  nodes,
  edges,
  narrative,
}: AttackPathVisualizerProps) {
  const svgRef = useRef<SVGSVGElement>(null);
  const [hoveredNode, setHoveredNode] = useState<string | null>(null);

  if (!nodes || nodes.length === 0) {
    return (
      <div className="text-vault-muted text-sm">
        No attack path data available
      </div>
    );
  }

  // Calculate positions for linear layout (left to right)
  const padding = 40;
  const nodeWidth = 120;
  const nodeHeight = 80;
  const svgWidth = nodes.length * (nodeWidth + 60) + padding * 2;
  const svgHeight = 200;

  const nodePositions: Record<string, { x: number; y: number }> = {};
  nodes.forEach((node, index) => {
    nodePositions[node.id] = {
      x: padding + index * (nodeWidth + 60),
      y: svgHeight / 2 - nodeHeight / 2,
    };
  });

  return (
    <div className="w-full">
      {narrative && (
        <div className="bg-vault-dark/40 border border-vault-sapphireDim/30 rounded-lg p-3 mb-4">
          <p className="text-xs text-vault-sapphireLight font-medium mb-1">
            Attack Narrative
          </p>
          <p className="text-vault-subtle text-sm leading-relaxed">{narrative}</p>
        </div>
      )}

      <div className="bg-vault-dark/20 border border-vault-border/20 rounded-lg p-2 overflow-x-auto">
        <svg
          ref={svgRef}
          width={svgWidth}
          height={svgHeight}
          className="min-w-full"
          style={{ display: "block" }}
        >
          {/* Edges/Arrows */}
          <defs>
            <marker
              id="arrowhead"
              markerWidth="10"
              markerHeight="10"
              refX="9"
              refY="3"
              orient="auto"
            >
              <polygon points="0 0, 10 3, 0 6" fill="#3b82d4" />
            </marker>
          </defs>

          {edges.map((edge, idx) => {
            const source = nodePositions[edge.source];
            const target = nodePositions[edge.target];

            if (!source || !target) return null;

            const x1 = source.x + nodeWidth / 2;
            const y1 = source.y + nodeHeight / 2;
            const x2 = target.x - nodeWidth / 2;
            const y2 = target.y + nodeHeight / 2;

            return (
              <line
                key={`edge-${idx}`}
                x1={x1}
                y1={y1}
                x2={x2}
                y2={y2}
                stroke="#3b82d4"
                strokeWidth="2"
                markerEnd="url(#arrowhead)"
                opacity={hoveredNode ? 0.3 : 0.6}
                style={{ transition: "opacity 0.2s" }}
              />
            );
          })}

          {/* Nodes */}
          {nodes.map((node) => {
            const pos = nodePositions[node.id];
            if (!pos) return null;

            const colors = getNodeColors(node.type);
            const isHovered = hoveredNode === node.id;

            return (
              <g key={node.id}>
                {/* Node background */}
                <rect
                  x={pos.x - nodeWidth / 2}
                  y={pos.y - nodeHeight / 2}
                  width={nodeWidth}
                  height={nodeHeight}
                  rx="6"
                  className={`${colors.bg} ${colors.border}`}
                  stroke={getRiskColor(node.risk_level)}
                  strokeWidth={isHovered ? "2" : "1"}
                  opacity={hoveredNode && !isHovered ? 0.4 : 1}
                  style={{
                    cursor: "pointer",
                    transition: "opacity 0.2s, stroke-width 0.2s",
                  }}
                  onMouseEnter={() => setHoveredNode(node.id)}
                  onMouseLeave={() => setHoveredNode(null)}
                />

                {/* Node label */}
                <text
                  x={pos.x}
                  y={pos.y - 12}
                  textAnchor="middle"
                  className={colors.text}
                  fontSize="12"
                  fontWeight="600"
                  pointerEvents="none"
                >
                  {node.label}
                </text>

                {/* Node type */}
                <text
                  x={pos.x}
                  y={pos.y + 8}
                  textAnchor="middle"
                  fill="#94a3b8"
                  fontSize="10"
                  pointerEvents="none"
                  fontStyle="italic"
                >
                  {node.type.replace(/_/g, " ")}
                </text>

                {/* Risk level indicator */}
                <circle
                  cx={pos.x + nodeWidth / 2 - 8}
                  cy={pos.y - nodeHeight / 2 + 8}
                  r="5"
                  fill={getRiskColor(node.risk_level)}
                  opacity="0.8"
                  pointerEvents="none"
                />
              </g>
            );
          })}
        </svg>
      </div>

      {/* Legend */}
      <div className="grid grid-cols-2 md:grid-cols-4 gap-3 mt-4 text-xs">
        {["attacker", "entry_point", "pivot", "target"].map((type) => {
          const colors = getNodeColors(type);
          return (
            <div key={type} className="flex items-center gap-2">
              <div
                className={`w-3 h-3 rounded ${colors.bg} border ${colors.border}`}
              />
              <span className="text-vault-muted">{type.replace(/_/g, " ")}</span>
            </div>
          );
        })}
      </div>

      {/* Risk level legend */}
      <div className="mt-3 pt-3 border-t border-vault-border/20">
        <p className="text-xs text-vault-muted font-medium mb-2">Risk Level</p>
        <div className="flex flex-wrap gap-3">
          {["none", "low", "medium", "high", "critical"].map((level) => (
            <div key={level} className="flex items-center gap-2">
              <div
                className="w-3 h-3 rounded-full"
                style={{ backgroundColor: getRiskColor(level) }}
              />
              <span className="text-vault-muted text-xs capitalize">{level}</span>
            </div>
          ))}
        </div>
      </div>
    </div>
  );
}