import React, { useEffect, useMemo, useRef } from "react";
import * as d3 from "d3";
import { sankey, sankeyLinkHorizontal } from "d3-sankey";

/**
 * dataRows: [{ source: "PC1", target: "1.2.3.4:443", value: 12 }, ...]
 */
export default function SankeyViz({ dataRows, height = 520 }) {
  const ref = useRef(null);

  const graph = useMemo(() => {
    const rows = Array.isArray(dataRows) ? dataRows : [];
    const nodesMap = new Map();
    const nodes = [];
    const links = [];

    function getNode(name) {
      if (!nodesMap.has(name)) {
        const id = nodes.length;
        nodesMap.set(name, id);
        nodes.push({ name });
      }
      return nodesMap.get(name);
    }

    for (const r of rows) {
      if (!r?.source || !r?.target) continue;
      const s = getNode(String(r.source));
      const t = getNode(String(r.target));
      const v = Number(r.value ?? 1);
      links.push({ source: s, target: t, value: isFinite(v) ? v : 1 });
    }

    return { nodes, links };
  }, [dataRows]);

  useEffect(() => {
    const el = ref.current;
    if (!el) return;

    el.innerHTML = "";

    const width = el.clientWidth || 900;
    const svg = d3
      .select(el)
      .append("svg")
      .attr("width", "100%")
      .attr("viewBox", `0 0 ${width} ${height}`)
      .style("background", "rgba(255,255,255,0.02)")
      .style("border", "1px solid rgba(255,255,255,0.08)")
      .style("borderRadius", "14px");

    if (!graph.nodes.length || !graph.links.length) {
      svg
        .append("text")
        .attr("x", width / 2)
        .attr("y", height / 2)
        .attr("text-anchor", "middle")
        .style("fill", "rgba(255,255,255,0.7)")
        .style("font-size", "14px")
        .text("No data for this selection.");
      return;
    }

    const sk = sankey()
      .nodeWidth(18)
      .nodePadding(14)
      .extent([
        [20, 20],
        [width - 20, height - 20],
      ]);

    const { nodes, links } = sk({
      nodes: graph.nodes.map((d) => ({ ...d })),
      links: graph.links.map((d) => ({ ...d })),
    });

    // Links
    svg
      .append("g")
      .attr("fill", "none")
      .selectAll("path")
      .data(links)
      .join("path")
      .attr("d", sankeyLinkHorizontal())
      .attr("stroke", "rgba(80,160,255,0.55)")
      .attr("stroke-width", (d) => Math.max(1, d.width))
      .attr("opacity", 0.9);

    // Nodes
    const node = svg
      .append("g")
      .selectAll("g")
      .data(nodes)
      .join("g");

    node
      .append("rect")
      .attr("x", (d) => d.x0)
      .attr("y", (d) => d.y0)
      .attr("height", (d) => Math.max(1, d.y1 - d.y0))
      .attr("width", (d) => d.x1 - d.x0)
      .attr("rx", 6)
      .attr("fill", "rgba(255,255,255,0.18)")
      .attr("stroke", "rgba(255,255,255,0.15)");

    node
      .append("text")
      .attr("x", (d) => (d.x0 < width / 2 ? d.x1 + 8 : d.x0 - 8))
      .attr("y", (d) => (d.y0 + d.y1) / 2)
      .attr("dy", "0.35em")
      .attr("text-anchor", (d) => (d.x0 < width / 2 ? "start" : "end"))
      .style("fill", "rgba(255,255,255,0.85)")
      .style("font-size", "12px")
      .text((d) => d.name);
  }, [graph, height]);

  return <div ref={ref} style={{ width: "100%" }} />;
}