or1-design / dfgraph / frontend / src / main.ts

import cytoscape from "cytoscape";
import elk from "cytoscape-elk";
import svg from "cytoscape-svg";
import type { GraphUpdate, GraphNode, GraphEdge, GraphRegion } from "@common/types";
import { stylesheet } from "./style";
import { logicalLayout, physicalLayout } from "@common/layout";
import { exportSvg, exportPng, copyPng } from "@common/export";

cytoscape.use(elk);
cytoscape.use(svg);

const ROUTING_CATEGORY = "routing";

const cy = cytoscape({
  container: document.getElementById("graph"),
  style: stylesheet,
  elements: [],
  maxZoom: 2.0,
  minZoom: 0.1,
});

function buildLabel(node: GraphNode): string {
  if (node.label) return node.label;
  if (node.const !== null) {
    return `${node.opcode}\n${node.const}`;
  }
  return node.opcode;
}

function buildPhysicalLabel(node: GraphNode): string {
  if (node.label) return node.label;
  const lines: string[] = [];
  if (node.const !== null) {
    lines.push(`${node.opcode}=${node.const}`);
  } else {
    lines.push(node.opcode);
  }
  if (node.iram_offset !== null) lines.push(`iram:${node.iram_offset}`);
  if (node.ctx !== null) lines.push(`ctx:${node.ctx}`);
  return lines.join("\n");
}

function buildElements(update: GraphUpdate): cytoscape.ElementDefinition[] {
  const elements: cytoscape.ElementDefinition[] = [];
  const regionParents = new Map<string, string>();

  for (const region of update.regions) {
    if (region.kind === "function") {
      elements.push({
        data: { id: region.tag, label: region.tag },
      });
      for (const nodeId of region.node_ids) {
        regionParents.set(nodeId, region.tag);
      }
    }
  }

  for (const node of update.nodes) {
    const el: cytoscape.ElementDefinition = {
      data: {
        id: node.id,
        label: buildLabel(node),
        colour: node.colour,
        category: node.category,
        pe: node.pe,
        iram_offset: node.iram_offset,
        ctx: node.ctx,
      },
      classes: node.has_error ? "error" : undefined,
    };
    const parent = regionParents.get(node.id);
    if (parent) {
      el.data.parent = parent;
    }
    elements.push(el);
  }

  for (const edge of update.edges) {
    const sourceNode = update.nodes.find((n) => n.id === edge.source);
    let sourceLabel: string | undefined;
    if (sourceNode && sourceNode.category === ROUTING_CATEGORY && edge.source_port) {
      sourceLabel = edge.source_port === "L" ? "T" : "F";
    }

    const edgeClasses = [
      edge.has_error ? "error" : undefined,
      edge.synthetic ? "synthetic" : undefined,
    ].filter(Boolean).join(" ") || undefined;

    elements.push({
      data: {
        id: `${edge.source}->${edge.target}:${edge.port}`,
        source: edge.source,
        target: edge.target,
        targetLabel: edge.port,
        sourceLabel: sourceLabel ?? "",
      },
      classes: edgeClasses,
    });
  }

  return elements;
}

function buildPhysicalElements(update: GraphUpdate): cytoscape.ElementDefinition[] {
  const elements: cytoscape.ElementDefinition[] = [];

  // Build set of edge targets for seed const detection
  const edgeTargets = new Set<string>();
  for (const edge of update.edges) {
    edgeTargets.add(edge.target);
  }

  const seedConstIds = new Set<string>();
  const nodePeMap = new Map<string, number | null>();
  const peIds = new Set<number>();

  for (const node of update.nodes) {
    nodePeMap.set(node.id, node.pe);
    if (node.pe !== null) peIds.add(node.pe);

    if (node.category === "config" && node.opcode === "const" && !edgeTargets.has(node.id)) {
      seedConstIds.add(node.id);
    }
  }

  // Create PE cluster parent nodes
  for (const peId of peIds) {
    elements.push({
      data: { id: `pe-${peId}`, label: `PE ${peId}` },
      classes: "pe-cluster",
    });
  }

  // Create operation nodes
  for (const node of update.nodes) {
    const isSeedConst = seedConstIds.has(node.id);
    const el: cytoscape.ElementDefinition = {
      data: {
        id: node.id,
        label: buildPhysicalLabel(node),
        colour: node.colour,
        category: node.category,
        pe: node.pe,
        iram_offset: node.iram_offset,
        ctx: node.ctx,
      },
      classes:
        [isSeedConst ? "seed-const" : undefined, node.has_error ? "error" : undefined].filter(Boolean).join(" ") ||
        undefined,
    };
    // Synthetic SM nodes float freely (not parented to any PE cluster)
    if (!node.synthetic && !isSeedConst && node.pe !== null) {
      el.data.parent = `pe-${node.pe}`;
    }
    elements.push(el);
  }

  // Scan cross-PE edges (excluding seed const sources) to find PE pairs
  type PePair = { srcPe: number; tgtPe: number };
  const pePairKey = (src: number, tgt: number) => `${src}->${tgt}`;
  const crossPePairs = new Map<string, PePair>();
  const busEdgeCounts = new Map<string, number>();

  for (const edge of update.edges) {
    if (seedConstIds.has(edge.source)) continue;
    const srcPe = nodePeMap.get(edge.source) ?? null;
    const tgtPe = nodePeMap.get(edge.target) ?? null;
    if (srcPe !== null && tgtPe !== null && srcPe !== tgtPe) {
      const key = pePairKey(srcPe, tgtPe);
      if (!crossPePairs.has(key)) {
        crossPePairs.set(key, { srcPe, tgtPe });
      }
      busEdgeCounts.set(key, (busEdgeCounts.get(key) ?? 0) + 1);
    }
  }

  // Create port nodes for each directed PE pair
  for (const [, pair] of crossPePairs) {
    elements.push({
      data: {
        id: `port-${pair.srcPe}-to-${pair.tgtPe}-exit`,
        parent: `pe-${pair.srcPe}`,
      },
      classes: "port-node",
    });
    elements.push({
      data: {
        id: `port-${pair.srcPe}-to-${pair.tgtPe}-entry`,
        parent: `pe-${pair.tgtPe}`,
      },
      classes: "port-node",
    });
  }

  // Create bus edges (one per PE pair)
  for (const [key, pair] of crossPePairs) {
    const count = busEdgeCounts.get(key) ?? 1;
    elements.push({
      data: {
        id: `bus-${pair.srcPe}-to-${pair.tgtPe}`,
        source: `port-${pair.srcPe}-to-${pair.tgtPe}-exit`,
        target: `port-${pair.srcPe}-to-${pair.tgtPe}-entry`,
        label: count > 1 ? `×${count}` : "",
      },
      classes: "physical bus-segment",
    });
  }

  // Create edges
  for (const edge of update.edges) {
    const srcPe = nodePeMap.get(edge.source) ?? null;
    const tgtPe = nodePeMap.get(edge.target) ?? null;
    const errorClass = edge.has_error ? " error" : "";
    const syntheticClass = edge.synthetic ? " synthetic" : "";

    if (edge.synthetic) {
      // Synthetic edges (e.g., SM request/return) — never bundled into bus
      elements.push({
        data: {
          id: `${edge.source}->${edge.target}:${edge.port}`,
          source: edge.source,
          target: edge.target,
          targetLabel: edge.port,
          sourceLabel: "",
        },
        classes: `physical synthetic${errorClass}`,
      });
    } else if (seedConstIds.has(edge.source)) {
      // Seed const: direct edge, never bundled
      elements.push({
        data: {
          id: `${edge.source}->${edge.target}:${edge.port}`,
          source: edge.source,
          target: edge.target,
          targetLabel: edge.port,
          sourceLabel: "",
        },
        classes: `seed-edge physical${errorClass}`,
      });
    } else if (srcPe !== null && tgtPe !== null && srcPe !== tgtPe) {
      // Cross-PE: split into exit segment + entry segment (bus already created)
      elements.push({
        data: {
          id: `${edge.source}->exit-${srcPe}-to-${tgtPe}:${edge.port}`,
          source: edge.source,
          target: `port-${srcPe}-to-${tgtPe}-exit`,
          sourceLabel: "",
          targetLabel: "",
        },
        classes: `physical exit-segment${errorClass}`,
      });
      elements.push({
        data: {
          id: `entry-${srcPe}-to-${tgtPe}->${edge.target}:${edge.port}`,
          source: `port-${srcPe}-to-${tgtPe}-entry`,
          target: edge.target,
          targetLabel: edge.port,
        },
        classes: `physical entry-segment${errorClass}`,
      });
    } else {
      // Intra-PE or unplaced
      elements.push({
        data: {
          id: `${edge.source}->${edge.target}:${edge.port}`,
          source: edge.source,
          target: edge.target,
          targetLabel: edge.port,
        },
        classes: `physical intra-pe${errorClass}`,
      });
    }
  }

  return elements;
}

function getRequiredElement(id: string): HTMLElement {
  const element = document.getElementById(id);
  if (!element) throw new Error(`Required element with id "${id}" not found`);
  return element;
}

function updateErrorPanel(update: GraphUpdate): void {
  const panel = getRequiredElement("error-panel");
  const list = getRequiredElement("error-list");
  const count = getRequiredElement("error-count");
  const overlay = getRequiredElement("parse-error-overlay");

  // Handle parse error (AC5.5)
  if (update.parse_error) {
    overlay.textContent = update.parse_error;
    overlay.classList.add("visible");
    panel.classList.remove("visible");
    return;
  }
  overlay.classList.remove("visible");

  // Handle pipeline errors
  if (update.errors.length === 0) {
    panel.classList.remove("visible");
    list.innerHTML = "";
    return;
  }

  count.textContent = `${update.errors.length} error${update.errors.length > 1 ? "s" : ""}`;
  list.innerHTML = "";

  for (const error of update.errors) {
    const li = document.createElement("li");
    li.className = "error-item";

    // Create line:column span
    const lineSpan = document.createElement("span");
    lineSpan.className = "error-line";
    lineSpan.textContent = `L${error.line}:${error.column}`;
    li.appendChild(lineSpan);

    // Create category span
    const categorySpan = document.createElement("span");
    categorySpan.className = "error-category";
    categorySpan.textContent = `[${error.category}]`;
    li.appendChild(categorySpan);

    // Create message text node (safe from XSS)
    const messageNode = document.createTextNode(error.message);
    li.appendChild(messageNode);

    if (error.suggestions.length > 0) {
      for (const suggestion of error.suggestions) {
        const span = document.createElement("span");
        span.className = "error-suggestion";
        span.textContent = `→ ${suggestion}`;
        li.appendChild(span);
      }
    }

    li.addEventListener("click", () => {
      // Highlight related nodes in the graph
      cy.nodes().unselect();
      cy.nodes().forEach((node) => {
        if (node.hasClass("error")) {
          node.select();
        }
      });
    });

    list.appendChild(li);
  }

  panel.classList.add("visible");
}

type ViewMode = "logical" | "physical";
let currentView: ViewMode = "logical";
let latestUpdate: GraphUpdate | null = null;

function routeEdges(selector?: string): void {
  const nodes = cy.nodes().not(":parent").not(".port-node");
  const nodePositions: Array<{ x: number; y: number }> = [];
  nodes.forEach((n) => {
    nodePositions.push({ x: n.position("x"), y: n.position("y") });
  });

  let leftCount = 0;
  let rightCount = 0;

  const edges = selector ? cy.edges(selector) : cy.edges();
  edges.forEach((edge) => {
    const sy = edge.source().position("y");
    const ty = edge.target().position("y");
    const sx = edge.source().position("x");
    const tx = edge.target().position("x");
    const span = Math.abs(ty - sy);

    if (span < 80) {
      edge.style({ "curve-style": "bezier" });
      return;
    }

    const blocked = nodePositions.some((p) => {
      return p.y > Math.min(sy, ty) + 25 && p.y < Math.max(sy, ty) - 25;
    });

    if (!blocked) {
      edge.style({ "curve-style": "bezier" });
      return;
    }

    const avgX = (sx + tx) / 2;
    const centerX = (Math.min(...nodePositions.map((p) => p.x)) + Math.max(...nodePositions.map((p) => p.x))) / 2;
    const goLeft = avgX >= centerX;

    const baseOffset = 30;
    const stagger = goLeft ? 12 * leftCount : 12 * rightCount;
    const offset = (goLeft ? -1 : 1) * (baseOffset + stagger);

    if (goLeft) leftCount++;
    else rightCount++;

    edge.style({
      "curve-style": "unbundled-bezier",
      "control-point-distances": [offset * 0.6, offset, offset * 0.6],
      "control-point-weights": [0.2, 0.5, 0.8],
    });
  });
}

function renderLogical(update: GraphUpdate): void {
  cy.batch(() => {
    cy.elements().remove();
    cy.add(buildElements(update));
  });
  const layout = cy.layout(logicalLayout());
  layout.on("layoutstop", () => {
    routeEdges();
    cy.fit(undefined, 40);
  });
  layout.run();
}

function renderPhysical(update: GraphUpdate): void {
  cy.batch(() => {
    cy.elements().remove();
    cy.add(buildPhysicalElements(update));
  });
  const layout = cy.layout(physicalLayout());
  layout.on("layoutstop", () => {
    routeEdges(".intra-pe, .seed-edge");
    cy.fit(undefined, 40);
  });
  layout.run();
}

function renderUpdate(update: GraphUpdate): void {
  latestUpdate = update;
  if (currentView === "physical" && update.stage !== "allocate") {
    currentView = "logical";
    const toggleBtn = document.getElementById("view-toggle");
    if (toggleBtn) {
      toggleBtn.textContent = "Physical View";
    }
  }
  if (currentView === "logical") {
    renderLogical(update);
  } else {
    renderPhysical(update);
  }
  updateErrorPanel(update);
}

function setupToggleButton(): void {
  const toggleBtn = document.getElementById("view-toggle");
  if (toggleBtn) {
    toggleBtn.addEventListener("click", () => {
      if (!latestUpdate) return;
      if (currentView === "logical") {
        if (latestUpdate.stage !== "allocate") {
          // Physical view not available
          return;
        }
        currentView = "physical";
        toggleBtn.textContent = "Logical View";
        renderPhysical(latestUpdate);
      } else {
        currentView = "logical";
        toggleBtn.textContent = "Physical View";
        renderLogical(latestUpdate);
      }
    });
  }
}

function setupExportButtons(): void {
  document.getElementById("export-svg")?.addEventListener("click", () => exportSvg(cy));
  document.getElementById("export-png")?.addEventListener("click", () => exportPng(cy));
  document.getElementById("copy-png")?.addEventListener("click", () => copyPng(cy));
}

function connect(): void {
  const protocol = location.protocol === "https:" ? "wss:" : "ws:";
  const ws = new WebSocket(`${protocol}//${location.host}/ws`);

  ws.onmessage = (event: MessageEvent) => {
    const update: GraphUpdate = JSON.parse(event.data);
    if (update.type === "graph_update") {
      renderUpdate(update);
    }
  };

  ws.onerror = (event: Event) => {
    console.error("WebSocket error:", event);
  };

  ws.onclose = () => {
    setTimeout(connect, 2000);
  };
}

setupToggleButton();
setupExportButtons();
connect();