src/components/MemoryGraph.tsx

/**
 * MemoryGraph — Typographic Constellation
 *
 * Implements proper D3 Enter/Update/Exit to prevent visual reloading,
 * and optimized forceX/forceY gravity to keep labels centered and clickable.
 */

import { useEffect, useRef, useState } from 'react';
import * as d3 from 'd3';
import type { GraphData, GraphNode } from '../api/client';

const TYPE_COLORS: Record<string, string> = {
  episodic: 'var(--color-episodic)',
  semantic: 'var(--color-semantic)',
  procedural: 'var(--color-procedural)',
};

interface MemoryGraphProps {
  data: GraphData;
  width?: number;
  height?: number;
  onNodeClick?: (node: GraphNode) => void;
  onForget?: (nodeId: string) => void;
}

interface SimNode extends GraphNode, d3.SimulationNodeDatum {}
interface SimEdge extends d3.SimulationLinkDatum<SimNode> {
  similarity: number;
}

export default function MemoryGraph({
  data,
  width = 600,
  height = 500,
  onNodeClick,
  onForget,
}: MemoryGraphProps) {
  const svgRef = useRef<SVGSVGElement>(null);
  const containerRef = useRef<HTMLDivElement>(null);
  
  const simulationRef = useRef<d3.Simulation<SimNode, SimEdge> | null>(null);
  const zoomGroupRef = useRef<d3.Selection<SVGGElement, unknown, null, undefined> | null>(null);
  
  const [selectedNode, setSelectedNode] = useState<GraphNode | null>(null);
  const [dimensions, setDimensions] = useState({ width, height });

  // 1. Auto-resize observer
  useEffect(() => {
    if (!containerRef.current) return;
    const observer = new ResizeObserver((entries) => {
      for (const entry of entries) {
        setDimensions({
          width: entry.contentRect.width,
          height: entry.contentRect.height,
        });
      }
    });
    observer.observe(containerRef.current);
    return () => observer.disconnect();
  }, []);

  // 2. Initialize SVG and Simulation ONCE
  useEffect(() => {
    if (!svgRef.current) return;
    const svg = d3.select(svgRef.current);
    svg.selectAll('*').remove();

    const g = svg.append('g');
    zoomGroupRef.current = g;

    // Zoom behavior
    const zoom = d3.zoom<SVGSVGElement, unknown>()
      .scaleExtent([0.3, 4])
      .on('zoom', (event) => {
        g.attr('transform', event.transform);
      });
    svg.call(zoom);

    // Init Simulation
    simulationRef.current = d3.forceSimulation<SimNode>()
      .force('link', d3.forceLink<SimNode, SimEdge>().id((d) => d.id).distance((d) => 120 * (1 - d.similarity)))
      .force('charge', d3.forceManyBody().strength(-150))
      // Gentle gravity towards center to keep things bunched
      .force('x', d3.forceX().strength(0.08))
      .force('y', d3.forceY().strength(0.08))
      // Prevent text overlap
      .force('collide', d3.forceCollide<SimNode>().radius((d) => 30 + (d.importance * 15)));

    return () => {
      simulationRef.current?.stop();
    };
  }, []);

  // 3. Update dimensions of forces
  useEffect(() => {
    if (!simulationRef.current || dimensions.width === 0) return;
    const { width: w, height: h } = dimensions;
    
    simulationRef.current.force('x', d3.forceX(w / 2).strength(0.08));
    simulationRef.current.force('y', d3.forceY(h / 2).strength(0.08));
    simulationRef.current.alpha(0.3).restart();
  }, [dimensions]);

  // 4. Enter / Update / Exit Data
  useEffect(() => {
    if (!simulationRef.current || !zoomGroupRef.current) return;
    const g = zoomGroupRef.current;

    // We must pass D3 our data references, but keep them stable
    // Deep clone data to avoid mutating React state
    const nodes: SimNode[] = data.nodes.map((n) => ({ ...n }));
    const nodeMap = new Map(nodes.map((n) => [n.id, n]));
    
    const links: SimEdge[] = data.edges
      .filter((e) => nodeMap.has(e.source) && nodeMap.has(e.target))
      .map((e) => ({
        source: e.source,
        target: e.target,
        similarity: e.similarity,
      }));

    // Identify old nodes to preserve their x/y coordinates
    const oldNodesMap = new Map(simulationRef.current.nodes().map(n => [n.id, n]));
    for (const n of nodes) {
      if (oldNodesMap.has(n.id)) {
        const old = oldNodesMap.get(n.id)!;
        n.x = old.x;
        n.y = old.y;
        n.vx = old.vx;
        n.vy = old.vy;
      }
    }

    const simulation = simulationRef.current;
    simulation.nodes(nodes);
    
    const linkForce = simulation.force<d3.ForceLink<SimNode, SimEdge>>('link');
    if (linkForce) linkForce.links(links);

    // --- Edges ---
    let linkElements = g.selectAll<SVGLineElement, SimEdge>('line.graph-link')
      .data(links, (d) => `${(d.source as any).id || d.source}-${(d.target as any).id || d.target}`);

    linkElements.exit().remove();

    const linkEnter = linkElements.enter()
      .append('line')
      .attr('class', 'graph-link')
      .attr('stroke', 'var(--color-border)')
      .attr('stroke-opacity', 0.6)
      .attr('stroke-width', 1);

    linkElements = linkEnter.merge(linkElements);

    // --- Nodes ---
    let nodeElements = g.selectAll<SVGGElement, SimNode>('g.graph-node')
      .data(nodes, (d) => d.id);

    nodeElements.exit()
      .transition().duration(400)
      .style('opacity', 0)
      .remove();

    const nodeEnter = nodeElements.enter()
      .append('g')
      .attr('class', 'graph-node')
      .style('cursor', 'pointer')
      .style('opacity', 0);

    nodeEnter.transition().duration(600).style('opacity', 1);

    // Anchor dot
    nodeEnter.append('circle')
      .attr('r', 2)
      .attr('fill', (d) => TYPE_COLORS[d.memory_type] || 'var(--color-text-muted)');

    // Typographic label
    nodeEnter.append('text')
      .text((d) => {
        const typeStr = `[${d.memory_type.substring(0,3).toUpperCase()}]`;
        const contentStr = d.content.substring(0, 20) + (d.content.length > 20 ? '...' : '');
        return `${typeStr} ${contentStr}`;
      })
      .attr('x', 6)
      .attr('y', 3)
      .style('font-family', '"JetBrains Mono", monospace')
      .style('font-size', '10px')
      .style('fill', (d) => TYPE_COLORS[d.memory_type] || 'var(--color-text-muted)')
      .style('user-select', 'none');

    nodeElements = nodeEnter.merge(nodeElements);

    // Re-bind events
    nodeElements
      .on('mouseover', function (_event, _d) {
        d3.select(this).select('text').style('fill', 'var(--color-text-primary)').style('font-weight', 'bold');
      })
      .on('mouseout', function (_event, d) {
        d3.select(this).select('text').style('fill', TYPE_COLORS[d.memory_type] || 'var(--color-text-muted)').style('font-weight', 'normal');
      })
      .on('click', (_event, d) => {
        setSelectedNode(d);
        if (onNodeClick) onNodeClick(d);
      });

    // Drag behavior
    const drag = d3.drag<SVGGElement, SimNode>()
      .on('start', (event, d) => {
        if (!event.active) simulation.alphaTarget(0.3).restart();
        d.fx = d.x;
        d.fy = d.y;
      })
      .on('drag', (event, d) => {
        d.fx = event.x;
        d.fy = event.y;
      })
      .on('end', (event, d) => {
        if (!event.active) simulation.alphaTarget(0);
        d.fx = null;
        d.fy = null;
      });

    nodeElements.call(drag);

    // Tick update
    simulation.on('tick', () => {
      linkElements
        .attr('x1', (d) => (d.source as SimNode).x!)
        .attr('y1', (d) => (d.source as SimNode).y!)
        .attr('x2', (d) => (d.target as SimNode).x!)
        .attr('y2', (d) => (d.target as SimNode).y!);

      nodeElements.attr('transform', (d) => `translate(${d.x},${d.y})`);
    });

    simulation.alpha(0.1).restart();

  }, [data, onNodeClick]);

  return (
    <div ref={containerRef} style={{ position: 'relative', width: '100%', height: '100%' }}>
      <svg
        ref={svgRef}
        width={dimensions.width}
        height={dimensions.height}
        style={{
          background: 'var(--color-surface-alt)',
        }}
      />

      {/* Selected node panel */}
      {selectedNode && (
        <div
          className="font-mono"
          style={{
            position: 'absolute',
            bottom: '40px',
            right: '40px',
            width: '300px',
            background: 'var(--color-surface)',
            padding: '24px',
            border: '1px solid var(--color-border)',
            boxShadow: '0 10px 40px rgba(0,0,0,0.5)',
          }}
        >
          <div style={{ fontSize: '10px', color: TYPE_COLORS[selectedNode.memory_type], marginBottom: '12px', letterSpacing: '0.1em' }}>
            [{selectedNode.memory_type.toUpperCase()}]
          </div>
          <p style={{ fontSize: '13px', color: 'var(--color-text-primary)', marginBottom: '16px', lineHeight: 1.5, wordWrap: 'break-word' }}>
            {selectedNode.content}
          </p>
          <div style={{ fontSize: '10px', color: 'var(--color-text-muted)', marginBottom: '16px' }}>
            WEIGHT: {selectedNode.importance.toFixed(2)}
          </div>
          <div style={{ display: 'flex', gap: '12px' }}>
            {onForget && (
              <button
                onClick={() => {
                  onForget(selectedNode.id);
                  setSelectedNode(null);
                }}
                style={{
                  color: 'var(--color-text-muted)',
                  fontSize: '11px',
                  background: 'transparent',
                  border: 'none',
                  cursor: 'pointer',
                  textDecoration: 'underline',
                }}
              >
                [ FORGET ]
              </button>
            )}
            <button
              onClick={() => setSelectedNode(null)}
              style={{
                color: 'var(--color-text-primary)',
                fontSize: '11px',
                background: 'transparent',
                border: 'none',
                cursor: 'pointer',
                textDecoration: 'underline',
              }}
            >
              [ CLOSE ]
            </button>
          </div>
        </div>
      )}
    </div>
  );
}