app/frontend/src/components/Flow/BrainView.tsx

import { useQuery } from "@tanstack/react-query"
import { useNavigate } from "@tanstack/react-router"
import {
  type Edge,
  type Node,
  ReactFlow,
  ReactFlowProvider,
  useReactFlow,
} from "@xyflow/react"
import "@xyflow/react/dist/style.css"
import {
  forceCollide,
  forceLink,
  forceManyBody,
  forceSimulation,
  forceX,
  forceY,
  type SimulationNodeDatum,
} from "d3-force"
import { motion } from "motion/react"
import { useEffect, useMemo } from "react"

import type { BrainGraph } from "@/client"
import { scaleIn } from "@/lib/motion"
import { BrainEdge, type BrainEdgeData } from "./BrainEdge"
import { BrainNode, type BrainNodeData } from "./BrainNode"
import "./flow.css"
import { graphQueryOptions } from "./queries"
import { useFlowSocket } from "./useFlowSocket"

// Never past 1: the labels are `text-xs`, and a graph small enough to fit
// twice over should not render its text at twice the size of the sidebar.
const FIT = { padding: 0.2, maxZoom: 1 }

const nodeTypes = { brain: BrainNode }
const edgeTypes = { brain: BrainEdge }

/** Circle diameter, from how many neurons this one is wired to. */
const SIZE_MIN = 32
const SIZE_MAX = 96
/**
 * Room around a circle for the label it shows on hover. Reserved in the layout
 * rather than in the node box, so revealing a label moves nothing.
 */
const LABEL_ROOM = 26

type Placed = SimulationNodeDatum & { id: string; size: number; room: number }

/**
 * Where the neurons sit: a force layout run to rest once and then frozen.
 *
 * A live simulation would keep nudging nodes under the pointer while someone is
 * panning, and a graph that never stops moving is unreadable. Starting from a
 * fixed ellipse rather than d3's own random phyllotaxis also means the same
 * flows lay out the same way twice.
 *
 * `forceX`/`forceY` rather than `forceCenter`: centering only translates the
 * whole thing, so unconnected flows — which is most of them — would push each
 * other apart forever with nothing pulling back.
 */
function build(graph: BrainGraph): { nodes: Node[]; edges: Edge[] } {
  const degree = new Map<string, number>()
  for (const edge of graph.edges ?? []) {
    degree.set(edge.source, (degree.get(edge.source) ?? 0) + 1)
    degree.set(edge.target, (degree.get(edge.target) ?? 0) + 1)
  }

  const placed: Placed[] = (graph.nodes ?? []).map((node, index, all) => {
    const size = Math.min(SIZE_MAX, SIZE_MIN + 8 * (degree.get(node.id) ?? 0))
    const angle = (index / all.length) * 2 * Math.PI
    const radius = 60 + all.length * 10
    return {
      id: node.id,
      size,
      room: size / 2 + LABEL_ROOM,
      // An ellipse, not a circle: the band across Home is far wider than it is
      // tall, so the layout is meant to come out that way too. Seeding the
      // shape it should settle into beats making the forces below fight a
      // circle into one.
      x: Math.cos(angle) * radius * 1.8,
      y: Math.sin(angle) * radius * 0.55,
    }
  })
  const links = (graph.edges ?? []).map((edge) => ({
    source: edge.source,
    target: edge.target,
  }))

  forceSimulation(placed)
    .force(
      "link",
      forceLink(links)
        .id((node) => (node as Placed).id)
        .distance(130),
    )
    .force("charge", forceManyBody().strength(-320).distanceMax(500))
    .force(
      "collide",
      forceCollide<Placed>((node) => node.room),
    )
    // Centring pulls nearly three times as hard vertically as horizontally, so
    // the graph compresses into a band roughly twice as wide as it is tall
    // instead of the square an even pull settles into. Coming out of the
    // simulation rather than out of a transform keeps the circles round and the
    // edges the length the link force asked for.
    .force("x", forceX(0).strength(0.05))
    .force("y", forceY(0).strength(0.14))
    .stop()
    .tick(300)

  // Both of an edge's handles sit at a neuron's centre, so the edge is what has
  // to trim itself back to the rims; the radius is only known here.
  const radius = new Map(placed.map((node) => [node.id, node.size / 2]))

  return {
    nodes: placed.map((node, index) => {
      const source = (graph.nodes ?? [])[index]
      return {
        id: node.id,
        type: "brain",
        // React Flow places by the top-left corner; the layout means centres.
        position: {
          x: (node.x ?? 0) - node.size / 2,
          y: (node.y ?? 0) - node.size / 2,
        },
        data: {
          label: source.label,
          kind: source.kind,
          members: source.members ?? [],
          flows: source.flows ?? [],
          size: node.size,
        } satisfies BrainNodeData,
      }
    }),
    edges: (graph.edges ?? []).map((edge) => ({
      id: `${edge.source}->${edge.target}`,
      type: "brain",
      source: edge.source,
      target: edge.target,
      data: {
        messages: edge.messages ?? [],
        sourceRadius: radius.get(edge.source) ?? 0,
        targetRadius: radius.get(edge.target) ?? 0,
      } satisfies BrainEdgeData,
    })),
  }
}

function BrainCanvas() {
  const navigate = useNavigate()
  const { fitView } = useReactFlow()
  useFlowSocket()
  // The socket invalidates every "flows" key on a rebuild, so publishing
  // anywhere re-fetches this and the layout runs again.
  const { data } = useQuery(graphQueryOptions())

  const { nodes, edges } = useMemo(
    () => build(data ?? { nodes: [], edges: [] }),
    [data],
  )

  // A rebuild lays the whole graph out afresh, so the viewport someone was
  // looking through no longer frames anything. Only when the set of neurons
  // actually changed: a value arriving must not move the canvas.
  const shape = nodes.map((node) => node.id).join(" ")
  useEffect(() => {
    if (!shape) return
    // After the new nodes have been measured, or the fit is of the old ones.
    // Instant, never animated: an animated fit travels from React Flow's
    // default viewport to the content, which on first load is the whole graph
    // visibly sliding in from the corner.
    const frame = requestAnimationFrame(() => fitView(FIT))
    return () => cancelAnimationFrame(frame)
  }, [shape, fitView])

  return (
    // Held back until there is a layout, so the graph grows in from the centre
    // of the band it is already fitted to rather than appearing mid-pan.
    <motion.div
      variants={scaleIn}
      initial="hidden"
      animate={nodes.length ? "visible" : "hidden"}
      className="h-full w-full"
    >
      <ReactFlow
        nodes={nodes}
        edges={edges}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        nodesDraggable={false}
        nodesConnectable={false}
        elementsSelectable={false}
        proOptions={{ hideAttribution: true }}
        fitView
        fitViewOptions={FIT}
        minZoom={0.1}
        maxZoom={2}
        // A band inside a page that scrolls. React Flow claims the wheel by
        // default, so leaving these on would trap the page under the pointer;
        // the fit above keeps the whole graph in view without panning anyway.
        zoomOnScroll={false}
        zoomOnDoubleClick={false}
        panOnDrag={false}
        preventScrolling={false}
        onNodeClick={(_event, node) => {
          const [flow] = (node.data as BrainNodeData).flows
          if (flow)
            navigate({ to: "/flows/$flowName", params: { flowName: flow } })
        }}
        className="brain-flat h-full w-full"
      />
    </motion.div>
  )
}

/**
 * Every flow at once, merged on what each node talks to.
 *
 * Sits flat on whatever surface the page paints — no card, no canvas of its
 * own — because it is the overview at the top of Home rather than a screen you
 * navigate to. Read-only by design: what a neuron stands for lives in the flow
 * it came from, and clicking one goes there.
 */
export function BrainView() {
  return (
    <div className="h-72 w-full sm:h-96">
      <ReactFlowProvider>
        <BrainCanvas />
      </ReactFlowProvider>
    </div>
  )
}