flowfi-network/components/visualizations/SankeyDiagram.tsx

'use client'

import { useMemo } from 'react'
import {
  sankey as d3Sankey,
  sankeyLinkHorizontal,
  SankeyNode,
  SankeyLink,
} from 'd3-sankey'
import { FlowData } from '@/lib/types'
import AnimatedPipe from './AnimatedPipe'
import FlowParticles from './FlowParticles'
import WaterDistortion from './WaterDistortion'

interface SankeyDiagramProps {
  data: FlowData
  width?: number
  height?: number
  showLabels?: boolean
  speed?: 'slow' | 'normal' | 'fast'
  showParticles?: boolean
  showDistortion?: boolean
  maxParticles?: number
  className?: string
}

interface SNode {
  id: string
  label?: string
  color?: string
}

interface SLink {
  source: string
  target: string
  value: number
  color?: string
}

export default function SankeyDiagram({
  data,
  width = 800,
  height = 400,
  showLabels = false,
  speed = 'normal',
  showParticles = true,
  showDistortion = false,
  maxParticles = 6,
  className = '',
}: SankeyDiagramProps) {
  const { nodes, links } = useMemo(() => {
    const sankeyNodes = data.nodes.map(n => ({ ...n }))
    const sankeyLinks = data.links.map(l => ({
      source: l.source,
      target: l.target,
      value: l.value,
      color: l.color,
    }))

    const generator = (d3Sankey as any)()
      .nodeId((d: any) => d.id)
      .nodeWidth(12)
      .nodePadding(16)
      .extent([[1, 1], [width - 1, height - 1]])

    const result = generator({
      nodes: sankeyNodes,
      links: sankeyLinks,
    })

    return {
      nodes: result.nodes as any[],
      links: result.links as any[],
    }
  }, [data, width, height])

  const linkPath = sankeyLinkHorizontal()
  const filterId = `distortion-${Math.random().toString(36).slice(2, 7)}`

  let particleCount = 0

  return (
    <svg
      viewBox={`0 0 ${width} ${height}`}
      className={`w-full h-auto ${className}`}
      preserveAspectRatio="xMidYMid meet"
    >
      <defs>
        {showDistortion && <WaterDistortion id={filterId} />}
        {/* Define paths for particles */}
        {links.map((link, i) => {
          const d = linkPath(link as any)
          return d ? (
            <path key={`path-${i}`} id={`flow-path-${i}`} d={d} fill="none" />
          ) : null
        })}
      </defs>

      <g
        className={showDistortion ? 'water-distortion' : ''}
        filter={showDistortion ? `url(#${filterId})` : undefined}
      >
        {/* Links */}
        {links.map((link, i) => {
          const d = linkPath(link as any)
          if (!d) return null
          const linkWidth = Math.max(2, (link as any).width || 4)
          return (
            <AnimatedPipe
              key={`link-${i}`}
              d={d}
              width={linkWidth}
              color={link.color || '#2dd4bf'}
              speed={speed}
              opacity={0.6}
            />
          )
        })}

        {/* Particles */}
        {showParticles && links.map((link, i) => {
          if (particleCount >= maxParticles) return null
          const pCount = Math.min(2, maxParticles - particleCount)
          particleCount += pCount
          return (
            <FlowParticles
              key={`particles-${i}`}
              pathId={`flow-path-${i}`}
              count={pCount}
              color={link.color?.replace(/[0-9a-f]{2}$/i, 'ff') || '#2dd4bf'}
              size={2}
              duration={speed === 'fast' ? 2 : speed === 'slow' ? 5 : 3}
            />
          )
        })}

        {/* Nodes */}
        {nodes.map((node, i) => (
          <g key={`node-${i}`}>
            <rect
              x={node.x0}
              y={node.y0}
              width={(node.x1 ?? 0) - (node.x0 ?? 0)}
              height={(node.y1 ?? 0) - (node.y0 ?? 0)}
              fill={node.color || '#2dd4bf'}
              rx={3}
              opacity={0.8}
            />
            {showLabels && node.label && (
              <text
                x={(node.x0 ?? 0) < width / 2
                  ? (node.x1 ?? 0) + 8
                  : (node.x0 ?? 0) - 8
                }
                y={((node.y0 ?? 0) + (node.y1 ?? 0)) / 2}
                textAnchor={(node.x0 ?? 0) < width / 2 ? 'start' : 'end'}
                dominantBaseline="middle"
                fill="#e0f2f1"
                fontSize={11}
                fontFamily="'Outfit', system-ui, sans-serif"
                fontWeight={300}
                opacity={0.65}
              >
                {node.label}
              </text>
            )}
          </g>
        ))}
      </g>
    </svg>
  )
}