Highlight Country Map

"use client";

import React, { useCallback, useEffect, useMemo, useRef } from "react";
import { cn } from "@/lib/cn";
import {
  type CountryFeature,
  computeExtrusionOffsets,
  interpolateHexColor,
  MULTI_VIBRANT_SWATCHES,
  projectLonLat,
  WORLD_COUNTRIES_DATA,
  WORLD_MAP_THEMES,
  type WorldMapProjectionType,
  type WorldMapThemeId,
} from "./world-map-depth-engine";

const BASE_W = 1100;
const BASE_H = 620;
const DEFAULT_ZOOM = 1.0;
const MIN_ZOOM = 1.0;
const MAX_ZOOM = 8.0;

export interface WorldMapDepthProps {
  selectedCountries?: string[];
  onCountryClick?: (countryId: string) => void;
  onHoverCountry?: (country: CountryFeature | null) => void;
  depthPx?: number;
  extrudeAngleDeg?: number;
  pitchDeg?: number;
  projection?: WorldMapProjectionType;
  theme?: WorldMapThemeId;
  hoverColor?: string;
  showGraticule?: boolean;
  showLabels?: boolean;
  showConnectingArcs?: boolean;
  interactive?: boolean;
  zoom?: number;
  onZoomChange?: (zoom: number) => void;
  className?: string;
}

interface CachedCountryGeometry {
  country: CountryFeature;
  path2D: Path2D;
  centroidPx: [number, number];
}

export const WorldMapDepth = React.forwardRef<
  HTMLDivElement,
  WorldMapDepthProps
>(
  (
    {
      selectedCountries = ["USA", "BRA", "FRA", "VNM", "JPN", "AUS"],
      onCountryClick,
      onHoverCountry,
      depthPx = 8,
      extrudeAngleDeg = -125,
      pitchDeg = 14,
      projection = "natural-flat",
      theme = "monochrome-relief",
      hoverColor,
      showGraticule = false,
      showLabels = true,
      showConnectingArcs = false,
      interactive = true,
      zoom,
      onZoomChange,
      className,
    },
    ref,
  ) => {
    const canvasRef = useRef<HTMLCanvasElement | null>(null);
    const containerRef = useRef<HTMLDivElement | null>(null);

    const zoomRef = useRef<number>(zoom ?? DEFAULT_ZOOM);
    const panRef = useRef<{ x: number; y: number }>({ x: 0, y: 0 });
    const isDraggingRef = useRef<boolean>(false);
    const pointerMovedRef = useRef<boolean>(false);
    const dragStartRef = useRef<{
      x: number;
      y: number;
      panX: number;
      panY: number;
    }>({
      x: 0,
      y: 0,
      panX: 0,
      panY: 0,
    });

    useEffect(() => {
      if (zoom !== undefined) {
        zoomRef.current = Math.max(MIN_ZOOM, Math.min(MAX_ZOOM, zoom));
      }
    }, [zoom]);

    const elevationMapRef = useRef<Map<string, number>>(new Map());
    const selectedSetRef = useRef<Set<string>>(new Set(selectedCountries));
    selectedSetRef.current = useMemo(
      () => new Set(selectedCountries),
      [selectedCountries],
    );
    const selectedOrderRef = useRef<string[]>(selectedCountries);
    selectedOrderRef.current = selectedCountries;

    const hoveredIdRef = useRef<string | null>(null);

    const depthPxRef = useRef<number>(depthPx);
    depthPxRef.current = depthPx;
    const extrudeAngleRef = useRef<number>(extrudeAngleDeg);
    extrudeAngleRef.current = extrudeAngleDeg;
    const pitchDegRef = useRef<number>(pitchDeg);
    pitchDegRef.current = pitchDeg;
    const themeIdRef = useRef<WorldMapThemeId>(theme);
    themeIdRef.current = theme;
    const hoverColorRef = useRef<string | undefined>(hoverColor);
    hoverColorRef.current = hoverColor;
    const showGraticuleRef = useRef<boolean>(showGraticule);
    showGraticuleRef.current = showGraticule;
    const showLabelsRef = useRef<boolean>(showLabels);
    showLabelsRef.current = showLabels;
    const showArcsRef = useRef<boolean>(showConnectingArcs);
    showArcsRef.current = showConnectingArcs;

    useEffect(() => {
      selectedCountries.forEach((id) => {
        elevationMapRef.current.set(id, 1.0);
      });
    }, [selectedCountries]);

    const cachedGeometries = useMemo<CachedCountryGeometry[]>(() => {
      if (typeof Path2D === "undefined") return [];
      return WORLD_COUNTRIES_DATA.map((country) => {
        const p2d = new Path2D();
        for (const poly of country.polygons) {
          for (const ring of poly) {
            for (let i = 0; i < ring.length; i++) {
              const [px, py] = projectLonLat(
                ring[i][0],
                ring[i][1],
                BASE_W,
                BASE_H,
                projection,
                pitchDeg,
              );
              if (i === 0) p2d.moveTo(px, py);
              else p2d.lineTo(px, py);
            }
            p2d.closePath();
          }
        }
        const centroidPx = projectLonLat(
          country.centroid[0],
          country.centroid[1],
          BASE_W,
          BASE_H,
          projection,
          pitchDeg,
        );
        return { country, path2D: p2d, centroidPx };
      });
    }, [projection, pitchDeg]);

    const cachedGeometriesRef =
      useRef<CachedCountryGeometry[]>(cachedGeometries);
    cachedGeometriesRef.current = cachedGeometries;

    const graticulePath = useMemo<Path2D | null>(() => {
      if (typeof Path2D === "undefined") return null;
      const p2d = new Path2D();
      for (const lat of [-60, -30, 0, 30, 60]) {
        for (let lon = -180; lon <= 180; lon += 6) {
          const [x, y] = projectLonLat(
            lon,
            lat,
            BASE_W,
            BASE_H,
            projection,
            pitchDeg,
          );
          if (lon === -180) p2d.moveTo(x, y);
          else p2d.lineTo(x, y);
        }
      }
      for (const lon of [
        -180, -150, -120, -90, -60, -30, 0, 30, 60, 90, 120, 150, 180,
      ]) {
        for (let lat = -60; lat <= 80; lat += 5) {
          const [x, y] = projectLonLat(
            lon,
            lat,
            BASE_W,
            BASE_H,
            projection,
            pitchDeg,
          );
          if (lat === -60) p2d.moveTo(x, y);
          else p2d.lineTo(x, y);
        }
      }
      return p2d;
    }, [projection, pitchDeg]);

    const graticulePathRef = useRef<Path2D | null>(graticulePath);
    graticulePathRef.current = graticulePath;

    const screenToMapCoords = useCallback(
      (clientX: number, clientY: number): [number, number] | null => {
        const canvas = canvasRef.current;
        if (!canvas) return null;
        const rect = canvas.getBoundingClientRect();
        const sx = clientX - rect.left;
        const sy = clientY - rect.top;
        const W = rect.width || BASE_W;
        const H = rect.height || BASE_H;

        const fitScale = Math.min(W / BASE_W, H / BASE_H);
        const totalScale = fitScale * zoomRef.current;

        const cx = W * 0.5 + panRef.current.x;
        const cy = H * 0.5 + panRef.current.y;

        const mapX = (sx - cx) / totalScale + BASE_W * 0.5;
        const mapY = (sy - cy) / totalScale + BASE_H * 0.5;
        return [mapX, mapY];
      },
      [],
    );

    const findCountryAtPoint = useCallback(
      (mapX: number, mapY: number): string | null => {
        const canvas = canvasRef.current;
        if (!canvas) return null;
        const ctx = canvas.getContext("2d");
        if (!ctx) return null;

        const geoms = cachedGeometriesRef.current;
        const { extrudeX, extrudeY } = computeExtrusionOffsets(
          depthPxRef.current,
          extrudeAngleRef.current,
        );

        for (let i = geoms.length - 1; i >= 0; i--) {
          const g = geoms[i];
          const elev = elevationMapRef.current.get(g.country.id) || 0;
          if (elev > 0.1) {
            const testX = mapX - extrudeX * elev;
            const testY = mapY - extrudeY * elev;
            if (ctx.isPointInPath(g.path2D, testX, testY, "evenodd")) {
              return g.country.id;
            }
          }
        }

        for (let i = geoms.length - 1; i >= 0; i--) {
          const g = geoms[i];
          if (ctx.isPointInPath(g.path2D, mapX, mapY, "evenodd")) {
            return g.country.id;
          }
        }

        return null;
      },
      [],
    );

    const handlePointerDown = useCallback(
      (e: React.PointerEvent<HTMLCanvasElement>) => {
        if (!interactive) return;
        isDraggingRef.current = true;
        pointerMovedRef.current = false;
        dragStartRef.current = {
          x: e.clientX,
          y: e.clientY,
          panX: panRef.current.x,
          panY: panRef.current.y,
        };
        e.currentTarget.setPointerCapture(e.pointerId);
      },
      [interactive],
    );

    const handlePointerMove = useCallback(
      (e: React.PointerEvent<HTMLCanvasElement>) => {
        if (isDraggingRef.current) {
          const dx = e.clientX - dragStartRef.current.x;
          const dy = e.clientY - dragStartRef.current.y;
          if (Math.hypot(dx, dy) > 4) {
            pointerMovedRef.current = true;
          }
          if (pointerMovedRef.current) {
            panRef.current = {
              x: dragStartRef.current.panX + dx,
              y: dragStartRef.current.panY + dy,
            };
            return;
          }
        }

        const coords = screenToMapCoords(e.clientX, e.clientY);
        if (!coords) return;
        const hitId = findCountryAtPoint(coords[0], coords[1]);
        if (hitId !== hoveredIdRef.current) {
          hoveredIdRef.current = hitId;
          const cObj = hitId
            ? WORLD_COUNTRIES_DATA.find((c) => c.id === hitId) || null
            : null;
          onHoverCountry?.(cObj);
        }
      },
      [screenToMapCoords, findCountryAtPoint, onHoverCountry],
    );

    const handlePointerUp = useCallback(
      (e: React.PointerEvent<HTMLCanvasElement>) => {
        const wasDragging = isDraggingRef.current;
        const moved = pointerMovedRef.current;
        isDraggingRef.current = false;
        try {
          e.currentTarget.releasePointerCapture(e.pointerId);
        } catch {}

        if (wasDragging && !moved) {
          const coords = screenToMapCoords(e.clientX, e.clientY);
          if (coords) {
            const hitId = findCountryAtPoint(coords[0], coords[1]);
            if (hitId && onCountryClick) {
              onCountryClick(hitId);
            }
          }
        }
      },
      [screenToMapCoords, findCountryAtPoint, onCountryClick],
    );

    useEffect(() => {
      const canvas = canvasRef.current;
      if (!canvas || !interactive) return;

      const onWheel = (e: WheelEvent) => {
        e.preventDefault();
        const factor = e.deltaY < 0 ? 1.12 : 1 / 1.12;
        const prevZoom = zoomRef.current;
        const nextZoom = Math.max(
          MIN_ZOOM,
          Math.min(MAX_ZOOM, prevZoom * factor),
        );

        if (nextZoom === MIN_ZOOM) {
          panRef.current = { x: 0, y: 0 };
        } else {
          const rect = canvas.getBoundingClientRect();
          const mx = e.clientX - rect.left - rect.width * 0.5;
          const my = e.clientY - rect.top - rect.height * 0.5;
          const ratio = nextZoom / prevZoom;
          panRef.current = {
            x: mx - (mx - panRef.current.x) * ratio,
            y: my - (my - panRef.current.y) * ratio,
          };
        }

        zoomRef.current = nextZoom;
        onZoomChange?.(Number(nextZoom.toFixed(2)));
      };

      canvas.addEventListener("wheel", onWheel, { passive: false });
      return () => canvas.removeEventListener("wheel", onWheel);
    }, [interactive, onZoomChange]);

    useEffect(() => {
      let active = true;
      let animId = 0;
      let lastTime = performance.now();

      const renderLoop = (now: number) => {
        if (!active) return;
        const dt = Math.min((now - lastTime) / 1000, 0.05);
        lastTime = now;

        const geoms = cachedGeometriesRef.current;
        const selSet = selectedSetRef.current;
        const hovId = hoveredIdRef.current;

        const elevMap = elevationMapRef.current;
        const lerpSpeed = Math.min(1, dt * 10.5);
        for (let i = 0; i < geoms.length; i++) {
          const id = geoms[i].country.id;
          const isSelected = selSet.has(id);
          const isHovered = hovId === id;
          const target = isSelected ? 1.0 : isHovered ? 0.22 : 0.0;
          const curr = elevMap.get(id) || 0.0;
          if (Math.abs(target - curr) > 0.002) {
            elevMap.set(id, curr + (target - curr) * lerpSpeed);
          } else if (target === 0) {
            elevMap.delete(id);
          } else {
            elevMap.set(id, target);
          }
        }

        const canvas = canvasRef.current;
        if (canvas) {
          const ctx = canvas.getContext("2d");
          if (ctx) {
            const dpr =
              typeof window !== "undefined"
                ? Math.min(window.devicePixelRatio || 1, 2)
                : 1;
            const W = canvas.clientWidth || 860;
            const H = canvas.clientHeight || 520;

            if (
              canvas.width !== Math.round(W * dpr) ||
              canvas.height !== Math.round(H * dpr)
            ) {
              canvas.width = Math.round(W * dpr);
              canvas.height = Math.round(H * dpr);
            }

            const activeTheme =
              WORLD_MAP_THEMES[themeIdRef.current] ||
              WORLD_MAP_THEMES["monochrome-relief"];

            ctx.save();
            ctx.scale(dpr, dpr);

            ctx.fillStyle = activeTheme.oceanBg;
            ctx.fillRect(0, 0, W, H);

            const bgGlow = ctx.createRadialGradient(
              W * 0.5,
              H * 0.5,
              20,
              W * 0.5,
              H * 0.5,
              Math.max(W, H) * 0.6,
            );
            bgGlow.addColorStop(0, activeTheme.oceanGlow);
            bgGlow.addColorStop(1, "rgba(0, 0, 0, 0)");
            ctx.fillStyle = bgGlow;
            ctx.fillRect(0, 0, W, H);

            const fitScale = Math.min(W / BASE_W, H / BASE_H);
            const totalScale = fitScale * zoomRef.current;
            const cx = W * 0.5 + panRef.current.x;
            const cy = H * 0.5 + panRef.current.y;

            ctx.translate(cx, cy);
            ctx.scale(totalScale, totalScale);
            ctx.translate(-BASE_W * 0.5, -BASE_H * 0.5);

            if (showGraticuleRef.current && graticulePathRef.current) {
              ctx.save();
              ctx.strokeStyle = activeTheme.gridLine;
              ctx.lineWidth = 0.8 / Math.sqrt(zoomRef.current);
              ctx.stroke(graticulePathRef.current);
              ctx.restore();
            }

            ctx.save();
            ctx.lineJoin = "round";
            ctx.lineCap = "round";
            const baseBorderW = 0.75 / zoomRef.current ** 0.65;

            const elevatedList: Array<{
              geom: CachedCountryGeometry;
              elev: number;
            }> = [];

            for (let i = 0; i < geoms.length; i++) {
              const g = geoms[i];
              const id = g.country.id;
              const elev = elevMap.get(id) || 0;

              if (elev > 0.01) {
                elevatedList.push({ geom: g, elev });
              }

              const activeHoverFill =
                hoverColorRef.current || activeTheme.baseLandHoverFill;
              ctx.fillStyle =
                hovId === id ? activeHoverFill : activeTheme.baseLandFill;
              ctx.fill(g.path2D, "evenodd");
              ctx.strokeStyle = activeTheme.baseLandStroke;
              ctx.lineWidth = baseBorderW;
              ctx.stroke(g.path2D);
            }
            ctx.restore();

            if (elevatedList.length > 0) {
              const { extrudeX, extrudeY, shadowX, shadowY } =
                computeExtrusionOffsets(
                  depthPxRef.current,
                  extrudeAngleRef.current,
                );

              elevatedList.sort((a, b) => {
                const projA =
                  a.geom.centroidPx[0] * extrudeX +
                  a.geom.centroidPx[1] * extrudeY;
                const projB =
                  b.geom.centroidPx[0] * extrudeX +
                  b.geom.centroidPx[1] * extrudeY;
                return projA - projB;
              });

              ctx.save();
              for (let i = 0; i < elevatedList.length; i++) {
                const { geom, elev } = elevatedList[i];
                ctx.save();
                ctx.translate(shadowX * elev, shadowY * elev);
                ctx.shadowColor = activeTheme.shadowColor;
                ctx.shadowBlur = 16 * elev;
                ctx.shadowOffsetX = shadowX * 0.35 * elev;
                ctx.shadowOffsetY = shadowY * 0.35 * elev;
                ctx.fillStyle = activeTheme.shadowColor;
                ctx.fill(geom.path2D, "evenodd");
                ctx.restore();
              }
              ctx.restore();

              const orderList = selectedOrderRef.current;
              for (let i = 0; i < elevatedList.length; i++) {
                const { geom, elev } = elevatedList[i];
                const id = geom.country.id;
                const orderIdx = Math.max(0, orderList.indexOf(id));

                const isCountryHovered = hovId === id;
                const customHov = hoverColorRef.current;

                const swatch =
                  isCountryHovered && customHov
                    ? {
                        wallDark: "#020617",
                        wallLight: customHov,
                        topStart: customHov,
                        topEnd: customHov,
                        rim: "#ffffff",
                      }
                    : themeIdRef.current === "multi-vibrant"
                      ? MULTI_VIBRANT_SWATCHES[
                          orderIdx % MULTI_VIBRANT_SWATCHES.length
                        ]
                      : {
                          wallDark: activeTheme.extrudeWallDark,
                          wallLight: activeTheme.extrudeWallLight,
                          topStart: activeTheme.topFaceStart,
                          topEnd: activeTheme.topFaceEnd,
                          rim: activeTheme.topRimStroke,
                        };

                const totalShift = Math.hypot(extrudeX * elev, extrudeY * elev);
                const numSlices = Math.max(
                  6,
                  Math.min(24, Math.ceil(totalShift * 1.15)),
                );

                ctx.save();
                ctx.lineJoin = "round";
                for (let s = 1; s <= numSlices; s++) {
                  const frac = s / numSlices;
                  const sx = extrudeX * elev * frac;
                  const sy = extrudeY * elev * frac;
                  const wallColor = interpolateHexColor(
                    swatch.wallDark,
                    swatch.wallLight,
                    frac * 0.9,
                  );

                  ctx.save();
                  ctx.translate(sx, sy);
                  ctx.fillStyle = wallColor;
                  ctx.fill(geom.path2D, "evenodd");
                  ctx.strokeStyle = wallColor;
                  ctx.lineWidth = 1.1 / zoomRef.current ** 0.5;
                  ctx.stroke(geom.path2D);
                  ctx.restore();
                }
                ctx.restore();

                ctx.save();
                ctx.translate(extrudeX * elev, extrudeY * elev);

                const [cxPt, cyPt] = geom.centroidPx;
                const grad = ctx.createLinearGradient(
                  cxPt - 60,
                  cyPt - 60,
                  cxPt + 60,
                  cyPt + 60,
                );
                grad.addColorStop(0, swatch.topStart);
                grad.addColorStop(1, swatch.topEnd);

                ctx.fillStyle = grad;
                ctx.shadowColor = swatch.topStart;
                ctx.shadowBlur = 10 * elev;
                ctx.fill(geom.path2D, "evenodd");

                ctx.shadowBlur = 0;
                ctx.strokeStyle = swatch.rim;
                ctx.lineWidth =
                  (hovId === id ? 1.8 : 1.25) / zoomRef.current ** 0.55;
                ctx.stroke(geom.path2D);
                ctx.restore();
              }

              const activeSelectedGeoms = elevatedList.filter((item) =>
                selSet.has(item.geom.country.id),
              );
              if (showArcsRef.current && activeSelectedGeoms.length > 1) {
                ctx.save();
                for (let i = 0; i < activeSelectedGeoms.length - 1; i++) {
                  const a = activeSelectedGeoms[i];
                  const b = activeSelectedGeoms[i + 1];
                  const x1 = a.geom.centroidPx[0] + extrudeX * a.elev;
                  const y1 = a.geom.centroidPx[1] + extrudeY * a.elev;
                  const x2 = b.geom.centroidPx[0] + extrudeX * b.elev;
                  const y2 = b.geom.centroidPx[1] + extrudeY * b.elev;
                  const midX = (x1 + x2) * 0.5;
                  const midY =
                    (y1 + y2) * 0.5 - Math.hypot(x2 - x1, y2 - y1) * 0.22;

                  ctx.beginPath();
                  ctx.moveTo(x1, y1);
                  ctx.quadraticCurveTo(midX, midY, x2, y2);
                  ctx.strokeStyle = activeTheme.labelAccent;
                  ctx.lineWidth = 1.4 / zoomRef.current ** 0.5;
                  ctx.setLineDash([4, 4]);
                  ctx.stroke();
                }
                ctx.restore();
              }

              if (showLabelsRef.current) {
                ctx.save();
                const invScale = 1 / zoomRef.current ** 0.65;

                for (let i = 0; i < activeSelectedGeoms.length; i++) {
                  const { geom, elev } = activeSelectedGeoms[i];
                  const tx = geom.centroidPx[0] + extrudeX * elev;
                  const ty = geom.centroidPx[1] + extrudeY * elev;

                  ctx.beginPath();
                  ctx.arc(tx, ty, 2.8 * invScale, 0, Math.PI * 2);
                  ctx.fillStyle = "#ffffff";
                  ctx.fill();

                  const labelText = `${geom.country.name}`;
                  const fontSize = Math.max(7, 10.5 * invScale);
                  ctx.font = `700 ${fontSize}px system-ui, -apple-system, sans-serif`;
                  const textMetrics = ctx.measureText(labelText);
                  const padX = 5.5 * invScale;
                  const pillW = textMetrics.width + padX * 2;
                  const pillH = 16 * invScale;
                  const pillX = tx - pillW * 0.5;
                  const pillY = ty - pillH - 5 * invScale;

                  ctx.fillStyle = activeTheme.labelBg;
                  ctx.beginPath();
                  ctx.rect(pillX, pillY, pillW, pillH);
                  ctx.fill();
                  ctx.strokeStyle = activeTheme.topRimStroke;
                  ctx.lineWidth = 0.7 * invScale;
                  ctx.stroke();

                  ctx.fillStyle = activeTheme.labelText;
                  ctx.textAlign = "center";
                  ctx.textBaseline = "middle";
                  ctx.fillText(labelText, tx, pillY + pillH * 0.52);
                }
                ctx.restore();
              }
            }

            ctx.restore();
          }
        }

        animId = requestAnimationFrame(renderLoop);
      };

      animId = requestAnimationFrame(renderLoop);
      return () => {
        active = false;
        if (animId) cancelAnimationFrame(animId);
      };
    }, []);

    return (
      <div
        ref={(node) => {
          containerRef.current = node;
          if (typeof ref === "function") ref(node);
          else if (ref)
            (ref as React.MutableRefObject<HTMLDivElement | null>).current =
              node;
        }}
        className={cn(
          "relative w-full h-full overflow-hidden select-none",
          className,
        )}
      >
        <canvas
          ref={canvasRef}
          onPointerDown={handlePointerDown}
          onPointerMove={handlePointerMove}
          onPointerUp={handlePointerUp}
          onPointerLeave={() => {
            hoveredIdRef.current = null;
            onHoverCountry?.(null);
          }}
          className={cn(
            "w-full h-full block",
            interactive
              ? "cursor-grab active:cursor-grabbing"
              : "pointer-events-none",
          )}
          style={{ touchAction: "none" }}
        />
      </div>
    );
  },
);

WorldMapDepth.displayName = "WorldMapDepth";

export default WorldMapDepth;