import React, { useContext, useMemo } from 'react'
import {
  BaseEdge,
  getBezierPath,
  type EdgeProps,
} from '@xyflow/react'
import FlowContext from '../../lib/flowContext'
import { getConnectionLabelForTarget } from '../../lib/nodeRegistry'

const EDGE_STROKE_WIDTH = 2
const DOT_MARKER_R = 1.5
const EMPTY_PATH_NODE_IDS = new Set<string>()

export function AnimatedEdge({
  id,
  source,
  sourceX,
  sourceY,
  targetX,
  targetY,
  sourcePosition,
  targetPosition,
  style,
  label: labelProp,
  interactionWidth,
  target,
}: EdgeProps) {
  const ctx = useContext(FlowContext)
  const nodes = ctx?.nodes ?? []
  const pathNodeIds = ctx?.connectionPathNodeIds ?? EMPTY_PATH_NODE_IDS
  const targetNode = useMemo(() => nodes.find((n: any) => n.id === target), [nodes, target])
  const derivedLabel = useMemo(
    () => (targetNode?.type != null ? getConnectionLabelForTarget(targetNode.type) : undefined),
    [targetNode?.type]
  )
  const label = labelProp ?? derivedLabel

  const isOnUpdatingPath = pathNodeIds.has(source) && pathNodeIds.has(target)

  const [edgePath, edgeLabelX, edgeLabelY] = getBezierPath({
    sourceX,
    sourceY,
    targetX,
    targetY,
    sourcePosition,
    targetPosition,
  })

  const startId = `animated-edge-dot-start-${id}`
  const endId = `animated-edge-dot-end-${id}`

  return (
    <>
      <defs>
        <marker
          id={startId}
          markerWidth={DOT_MARKER_R * 2}
          markerHeight={DOT_MARKER_R * 2}
          refX={DOT_MARKER_R - 1}
          refY={DOT_MARKER_R}
          orient="auto"
        >
          <circle
            r={DOT_MARKER_R}
            cx={DOT_MARKER_R}
            cy={DOT_MARKER_R}
            className="fill-primary"
            strokeWidth={2}
          />
        </marker>
        <marker
          id={endId}
          markerWidth={DOT_MARKER_R * 2}
          markerHeight={DOT_MARKER_R * 2}
          refX={DOT_MARKER_R + 1}
          refY={DOT_MARKER_R}
          orient="auto"
        >
          <circle
            r={DOT_MARKER_R}
            cx={DOT_MARKER_R}
            cy={DOT_MARKER_R}
            className="fill-primary"
            strokeWidth={2}
          />
        </marker>
      </defs>
      <BaseEdge
        path={edgePath}
        markerStart={`url(#${startId})`}
        markerEnd={`url(#${endId})`}
        style={{
          strokeWidth: EDGE_STROKE_WIDTH,
          ...style,
        }}
        className={`animated-edge-path${isOnUpdatingPath ? ' animated-edge-path--updating' : ''}`}
        interactionWidth={interactionWidth}
      />
      {label != null && (
        <g transform={`translate(${edgeLabelX}, ${edgeLabelY})`} className="nodrag nopan">
          <rect
            x={-32}
            y={-9}
            width={64}
            height={18}
            rx={4}
            ry={4}
            className="fill-background stroke-border"
            strokeWidth={1}
          />
          <text
            textAnchor="middle"
            dominantBaseline="middle"
            className="fill-foreground text-[10px] font-medium"
          >
            {label}
          </text>
        </g>
      )}
    </>
  )
}