Aegis/frontend/src/components/heatmap/AdvancedHeatmap.tsx

import { useMemo, useRef } from "react";
import { useVirtualizer } from "@tanstack/react-virtual";
import type { HeatmapTechnique } from "../../api/heatmap";
import HeatmapCell from "./HeatmapCell";

// MITRE ATT&CK Enterprise tactics in canonical order
const TACTIC_ORDER = [
  "reconnaissance",
  "resource-development",
  "initial-access",
  "execution",
  "persistence",
  "privilege-escalation",
  "defense-evasion",
  "credential-access",
  "discovery",
  "lateral-movement",
  "collection",
  "command-and-control",
  "exfiltration",
  "impact",
];

const formatTacticName = (tactic: string): string =>
  tactic
    .split("-")
    .map((w) => w.charAt(0).toUpperCase() + w.slice(1))
    .join(" ");

interface AdvancedHeatmapProps {
  techniques: HeatmapTechnique[];
  onCellClick: (techniqueId: string) => void;
  zoom: "compact" | "normal" | "expanded";
}

/** Virtualised tactic column — renders only visible rows. */
function TacticColumn({
  tactic,
  techniques,
  zoom,
  onCellClick,
}: {
  tactic: string;
  techniques: HeatmapTechnique[];
  zoom: "compact" | "normal" | "expanded";
  onCellClick: (techniqueId: string) => void;
}) {
  const parentRef = useRef<HTMLDivElement>(null);

  const rowHeight = zoom === "compact" ? 28 : zoom === "normal" ? 40 : 60;

  const rowVirtualizer = useVirtualizer({
    count: techniques.length,
    getScrollElement: () => parentRef.current,
    estimateSize: () => rowHeight,
    overscan: 10,
  });

  const columnWidth =
    zoom === "compact" ? "w-32" : zoom === "normal" ? "w-44" : "w-56";

  return (
    <div className={`${columnWidth} flex-shrink-0`}>
      {/* Tactic header */}
      <div className="mb-2 rounded-lg bg-gray-800 px-2 py-2">
        <h3 className="text-center text-xs font-semibold text-cyan-400">
          {formatTacticName(tactic)}
        </h3>
        <p className="mt-0.5 text-center text-[10px] text-gray-500">
          {techniques.length} techniques
        </p>
      </div>

      {/* Virtualised list */}
      <div
        ref={parentRef}
        className="overflow-y-auto scrollbar-thin scrollbar-thumb-gray-700 scrollbar-track-gray-900"
        style={{ maxHeight: "calc(100vh - 320px)" }}
      >
        <div
          style={{
            height: `${rowVirtualizer.getTotalSize()}px`,
            width: "100%",
            position: "relative",
          }}
        >
          {rowVirtualizer.getVirtualItems().map((virtualRow) => {
            const tech = techniques[virtualRow.index];
            return (
              <div
                key={tech.techniqueID + tactic}
                style={{
                  position: "absolute",
                  top: 0,
                  left: 0,
                  width: "100%",
                  height: `${virtualRow.size}px`,
                  transform: `translateY(${virtualRow.start}px)`,
                  padding: "2px 0",
                }}
              >
                <HeatmapCell
                  technique={tech}
                  size={zoom}
                  onClick={onCellClick}
                />
              </div>
            );
          })}
        </div>
      </div>
    </div>
  );
}

export default function AdvancedHeatmap({
  techniques,
  onCellClick,
  zoom,
}: AdvancedHeatmapProps) {
  // Group techniques by tactic
  const groupedByTactic = useMemo(() => {
    const groups: Record<string, HeatmapTechnique[]> = {};

    for (const tech of techniques) {
      // Normalize tactic names
      const tacticRaw = tech.tactic || "unknown";
      const tacticNormalized = tacticRaw
        .trim()
        .toLowerCase()
        .replace(/\s+/g, "-")
        .replace(/_/g, "-");

      if (!groups[tacticNormalized]) {
        groups[tacticNormalized] = [];
      }
      groups[tacticNormalized].push(tech);
    }

    // Sort techniques within each tactic by techniqueID
    for (const tactic of Object.keys(groups)) {
      groups[tactic].sort((a, b) =>
        a.techniqueID.localeCompare(b.techniqueID),
      );
    }

    return groups;
  }, [techniques]);

  // Get ordered tactics
  const orderedTactics = useMemo(() => {
    const tacticSet = new Set(Object.keys(groupedByTactic));
    const ordered = TACTIC_ORDER.filter((t) => tacticSet.has(t));
    const remaining = Array.from(tacticSet).filter(
      (t) => !TACTIC_ORDER.includes(t),
    );
    return [...ordered, ...remaining];
  }, [groupedByTactic]);

  if (techniques.length === 0) {
    return (
      <div className="rounded-xl border border-gray-800 bg-gray-900 p-8 text-center">
        <p className="text-gray-400">No techniques found for the selected layer</p>
      </div>
    );
  }

  return (
    <div className="overflow-x-auto rounded-xl border border-gray-800 bg-gray-900">
      <div className="min-w-max p-3">
        <div className="flex gap-2">
          {orderedTactics.map((tactic) => (
            <TacticColumn
              key={tactic}
              tactic={tactic}
              techniques={groupedByTactic[tactic] || []}
              zoom={zoom}
              onCellClick={onCellClick}
            />
          ))}
        </div>
      </div>
    </div>
  );
}