import { useState, useRef, useCallback, useEffect, useContext } from "react";
import { useJsonForms } from "@jsonforms/react";
import { Group } from "@visx/group";
import { hierarchy, Tree } from "@visx/hierarchy";
import { Zoom } from "@visx/zoom";
import { localPoint } from "@visx/event";
import { DataContext } from "../../context/Context";
import { Toolbar } from "./Toolbar";
import { NodeContainer } from "./NodeContainer";
import {
  ChartProps,
  TreeNode,
  DEFAULT_MARGIN,
  TOOLBAR_HEIGHT,
  MIN_ZOOM,
  MAX_ZOOM,
} from "./Interface";
import {
  mergeChildren,
  getLinkComponent,
  getLevelStyle,
  extractLevelStyles,
  countVisibleLeaves,
  computeDynamicHeight,
  collectExpandedNodeIds,
} from "./Utils";

function SvgDefs() {
  return (
    <defs>
      <filter id="node-hover-shadow" x="-20%" y="-30%" width="140%" height="160%">
        <feDropShadow dx="0" dy="6" stdDeviation="8" floodColor="#000000" floodOpacity="0.11" />
      </filter>
      <filter id="toggle-hover-shadow" x="-60%" y="-60%" width="220%" height="220%">
        <feDropShadow dx="0" dy="2" stdDeviation="3" floodColor="#000000" floodOpacity="0.16" />
      </filter>
    </defs>
  );
}

const ANIMATION_STYLES = `
  @keyframes nodeSlideIn {
    0%   { opacity: 0; transform: translateY(-20px); }
    100% { opacity: 1; transform: translateY(0px); }
  }
  @keyframes linkDraw {
    0%   { stroke-dashoffset: 1; }
    100% { stroke-dashoffset: 0; }
  }
  @keyframes toggleSpin {
    from { transform: rotate(0deg); }
    to   { transform: rotate(360deg); }
  }
`;

export function DrawHierarchyChart({
  path,
  width: totalWidth,
  height: totalHeight,
  data,
  linkType,
  stepPercent,
  nodeWidth,
  nodeHeight,
  styleConfig,
  margin = DEFAULT_MARGIN,
  renderers,
  cells,
  enabled,
  nodeElements,
  islazyLoad,
  lazyLoadFunction
}: ChartProps) {
  const ctx = useJsonForms();
  const canvasBgColor: string = styleConfig.canvasBgColor ?? "#ffffff";
  const linkColor: string = styleConfig?.linkColor;
  const { serviceProvider } = useContext(DataContext);

  const [localTree, setLocalTree] = useState<TreeNode>(data);

  const [expandedIds, setExpandedIds] = useState<Set<number | string>>(
    () => collectExpandedNodeIds(data)
  );

  const [nodeHeights, setNodeHeights] = useState<Record<number | string, number>>({});

  const [loadingNodeIds, setLoadingNodeIds] = useState<Set<number | string>>(new Set());
  const [recentlyExpandedId, setRecentlyExpandedId] = useState<number | string | null>(null);
  const expandAnimTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  const isDraggingRef = useRef(false);
  const dragStartPosRef = useRef({ x: 0, y: 0 });
  const lastPinchDistRef = useRef<number | null>(null);

  const levelStyles = extractLevelStyles(styleConfig);

  useEffect(() => {
    setLocalTree(data);
    setExpandedIds(
      collectExpandedNodeIds(data)
    );
    setNodeHeights({});
  }, [data, islazyLoad]);

  const handleHeightChange = useCallback((id: number | string, height: number) => {
    setNodeHeights((prev) => {
      if (prev[id] === height) return prev;
      return { ...prev, [id]: height };
    });
  }, []);

  const triggerExpandAnimation = useCallback((targetId: number | string) => {
    if (expandAnimTimerRef.current) clearTimeout(expandAnimTimerRef.current);
    setRecentlyExpandedId(targetId);
    expandAnimTimerRef.current = setTimeout(() => setRecentlyExpandedId(null), 700);
  }, []);

  const handleToggle = useCallback(
    async (targetId: number | string, nodeData: TreeNode) => {
      if (isDraggingRef.current) return;

      const isCurrentlyExpanded = expandedIds.has(targetId);
      const childrenAlreadyLoaded =
        Array.isArray(nodeData.children) && nodeData.children.length > 0;

      const shouldFetch =
        islazyLoad && !isCurrentlyExpanded && !childrenAlreadyLoaded;

      if (shouldFetch) {
        setLoadingNodeIds((prev) => new Set(prev).add(targetId));
        try {
          const fetchedChildren: TreeNode[] = await serviceProvider(
            ctx,
            { onClick: lazyLoadFunction || "onNodeExpandChange" },
            {
              event: { _reactName: "onClick" },
              path,
              paramValue: { path, expandedNodeId: targetId },
            }
          );

          if (Array.isArray(fetchedChildren)) {
            setLocalTree((prev) => mergeChildren(prev, targetId, fetchedChildren));
            setExpandedIds((prev) => {
              const next = new Set(prev);
              next.add(targetId);
              return next;
            });
            triggerExpandAnimation(targetId);
          }
        } catch (err) {
          console.error("HierarchyChart lazy load error:", err);
        } finally {
          setLoadingNodeIds((prev) => {
            const next = new Set(prev);
            next.delete(targetId);
            return next;
          });
        }
        return;
      }

      setExpandedIds((prev) => {
        const next = new Set(prev);
        if (next.has(targetId)) {
          next.delete(targetId);
        } else {
          next.add(targetId);
          triggerExpandAnimation(targetId);
        }
        return next;
      });
    },
    [expandedIds, islazyLoad, serviceProvider, ctx, lazyLoadFunction, triggerExpandAnimation, path]
  );


  const innerWidth = totalWidth - margin.left - margin.right;
  const canvasHeight = totalHeight - TOOLBAR_HEIGHT;
  const NODE_MIN_PADDING = 24 + nodeWidth * 0.25;
  const origin = { x: 0, y: 60 };

  const sizeWidth = Math.max(
    innerWidth,
    countVisibleLeaves(localTree, expandedIds) * (nodeWidth + NODE_MIN_PADDING)
  );
  const sizeHeight = computeDynamicHeight(localTree, expandedIds, nodeHeight, nodeHeights);

  const LinkComponent = getLinkComponent(linkType);
  const hasElements = nodeElements.length > 0;

  if (totalWidth < 10) return null;

  return (
    <>
      <style>{ANIMATION_STYLES}</style>

      <Zoom<SVGSVGElement>
        width={totalWidth}
        height={canvasHeight}
        scaleXMin={MIN_ZOOM}
        scaleXMax={MAX_ZOOM}
        scaleYMin={MIN_ZOOM}
        scaleYMax={MAX_ZOOM}
      >
        {(zoom) => {
          return (
            <div style={{ position: "relative", width: totalWidth, height: totalHeight }}>
              <Toolbar
                zoom={zoom}
                totalWidth={totalWidth}
              />
              <svg
                width={totalWidth}
                height={canvasHeight}
                ref={zoom.containerRef}
                style={{
                  cursor: zoom.isDragging ? "grabbing" : "grab",
                  touchAction: "none",
                  display: "block",
                }}
                onMouseMove={(e) => {
                  if (zoom.isDragging) {
                    const dx = Math.abs(e.clientX - dragStartPosRef.current.x);
                    const dy = Math.abs(e.clientY - dragStartPosRef.current.y);
                    if (dx > 3 || dy > 3) isDraggingRef.current = true;
                    zoom.dragMove(e);
                  }
                }}
                onMouseUp={() => {
                  zoom.dragEnd();
                  setTimeout(() => { isDraggingRef.current = false; }, 0);
                }}
                onMouseLeave={() => {
                  if (zoom.isDragging) zoom.dragEnd();
                }}
                onTouchStart={(e) => {
                  if (e.touches.length === 2) {
                    lastPinchDistRef.current = Math.hypot(
                      e.touches[1].clientX - e.touches[0].clientX,
                      e.touches[1].clientY - e.touches[0].clientY
                    );
                  } else {
                    lastPinchDistRef.current = null;
                    zoom.dragStart(e);
                  }
                }}
                onTouchMove={(e) => {
                  if (e.touches.length === 2) {
                    const dist = Math.hypot(
                      e.touches[1].clientX - e.touches[0].clientX,
                      e.touches[1].clientY - e.touches[0].clientY
                    );
                    if (lastPinchDistRef.current !== null) {
                      const scaleFactor = dist / lastPinchDistRef.current;
                      const svgRect = zoom.containerRef.current?.getBoundingClientRect();
                      const midClientX = (e.touches[0].clientX + e.touches[1].clientX) / 2;
                      const midClientY = (e.touches[0].clientY + e.touches[1].clientY) / 2;
                      const point = {
                        x: svgRect ? midClientX - svgRect.left : midClientX,
                        y: svgRect ? midClientY - svgRect.top : midClientY,
                      };
                      zoom.scale({ scaleX: scaleFactor, scaleY: scaleFactor, point });
                    }
                    lastPinchDistRef.current = dist;
                  } else {
                    lastPinchDistRef.current = null;
                    zoom.dragMove(e);
                  }
                }}
                onTouchEnd={() => {
                  lastPinchDistRef.current = null;
                  zoom.dragEnd();
                }}
                onWheel={(e) => {
                  const point = localPoint(e) || { x: 0, y: 0 };
                  const sf = e.deltaY < 0 ? 1.1 : 0.9;
                  zoom.scale({ scaleX: sf, scaleY: sf, point });
                }}
              >
                <SvgDefs />
                <rect
                  width={totalWidth}
                  height={canvasHeight}
                  fill={canvasBgColor}
                  onMouseDown={(e) => {
                    dragStartPosRef.current = { x: e.clientX, y: e.clientY };
                    isDraggingRef.current = false;
                    zoom.dragStart(e);
                  }}
                />
                <Group transform={zoom.toString()}>
                  <Group top={margin.top} left={margin.left}>
                    <Tree
                      root={hierarchy(localTree, (d) =>
                        expandedIds.has(d.id) ? d.children : null
                      )}
                      size={[sizeWidth, sizeHeight]}
                      separation={(a, b) => {
                        return a.parent === b.parent ? 1 : 1.5;
                      }}
                    >
                      {(tree) => {
                      const links = tree.links();
                      return (
                        <Group top={origin.y} left={origin.x}>
                          <defs>
                            {links.map((link, i) => {
                              const gradientId = `gradient-${i}`;
                              return (
                                <linearGradient
                                  key={gradientId}
                                  id={gradientId}
                                  gradientUnits="userSpaceOnUse"
                                  x1={link.source.x}
                                  y1={link.source.y}
                                  x2={link.target.x}
                                  y2={link.target.y}
                                >
                                  <stop offset="0%" stopColor={linkColor || "#949090"} stopOpacity={0.9} />
                                  <stop offset="50%" stopColor={linkColor || "#90a4ae"} stopOpacity={0.6} />
                                  <stop offset="100%" stopColor={linkColor || "#e0e0e0"} stopOpacity={0.3} />
                                </linearGradient>
                              );
                            })}
                          </defs>

                          {links.map((link, i) => {
                            const gradientId = `gradient-${i}`;
                            return (
                              <LinkComponent
                                key={i}
                                data={link}
                                percent={stepPercent}
                                stroke={`url(#${gradientId})`}
                                strokeWidth="3"
                                fill="none"
                              />
                            );
                          })}

                          {tree.descendants().map((node) => {
                            const top = node.y;
                            const left = node.x;
                            const nodeId = node.data.id;

                            const storedH = nodeHeights[nodeId];
                            const currentHeight =
                              storedH && storedH > 0 ? storedH : nodeHeight;

                            const hasChildren =
                              node.data.hasChildren === true ||
                              (Array.isArray(node.data.children) &&
                                node.data.children.length > 0);

                            const isCollapsed = !expandedIds.has(nodeId);
                            const depth = node.depth;
                            const levelStyle = getLevelStyle(levelStyles, depth);
                            const isHighlighted = node.data.isHighlighted === true;
                            const isNodeLoading = loadingNodeIds.has(nodeId);

                            const isChildOfExpanded =
                              recentlyExpandedId !== null &&
                              node.parent?.data?.id === recentlyExpandedId;

                            const childIndex = isChildOfExpanded
                              ? (node.parent?.children?.findIndex(
                                  (c) => c.data.id === nodeId
                                ) ?? 0)
                              : 0;

                            return (
                              <Group top={top} left={left} key={nodeId}>
                                <g
                                  style={
                                    isChildOfExpanded
                                      ? {
                                          animation: `nodeSlideIn 0.45s cubic-bezier(0.22,1,0.36,1) both`,
                                          animationDelay: `${childIndex * 55}ms`,
                                        }
                                      : undefined
                                  }
                                >
                                  {hasElements ? (
                                    <NodeContainer
                                      nodeData={node.data}
                                      renderers={renderers}
                                      cells={cells}
                                      enabled={enabled}
                                      nodeElements={nodeElements}
                                      nodeWidth={nodeWidth}
                                      minHeight={nodeHeight}
                                      hasChildren={hasChildren}
                                      isCollapsed={isCollapsed}
                                      onToggle={() =>
                                        handleToggle(nodeId, node.data)
                                      }
                                      onHeightChange={(h) =>
                                        handleHeightChange(nodeId, h)
                                      }
                                      currentHeight={currentHeight}
                                      levelStyle={levelStyle}
                                      depth={depth}
                                      isHighlighted={isHighlighted}
                                      isNodeLoading={isNodeLoading}
                                      styleConfig={styleConfig}
                                    />
                                  ) : (
                                    <rect
                                      height={nodeHeight}
                                      width={nodeWidth}
                                      y={-nodeHeight / 2}
                                      x={-nodeWidth / 2}
                                      fill={levelStyle.backgroundColor ?? "#ffffff"}
                                      stroke={
                                        isHighlighted
                                          ? "#1BA1A6"
                                          : levelStyle.borderColor ?? "#aaaaaa"
                                      }
                                      strokeWidth={isHighlighted ? 3 : levelStyle.borderWidth ?? 1}
                                      strokeDasharray={levelStyle.borderDash ?? "0"}
                                      strokeOpacity={levelStyle.borderOpacity ?? 1}
                                      rx={levelStyle.borderRadius ?? 6}
                                      style={{
                                        cursor: hasChildren ? "pointer" : "default",
                                      }}
                                      onClick={() =>
                                        hasChildren && handleToggle(nodeId, node.data)
                                      }
                                    />
                                  )}
                                </g>
                              </Group>
                            );
                          })}
                        </Group>
                      );
                    }}
                    </Tree>
                  </Group>
                </Group>
              </svg>
            </div>
          );
        }}
      </Zoom>
    </>
  );
}