import {
  useRef,
  useEffect,
  useCallback,
  useState,
  forwardRef,
  useImperativeHandle,
} from "react";
import type {
  GridCell,
  GridToken,
  Tool,
  ServerMessage,
  ClientMessage,
} from "../types";
import { useWebSocket } from "../hooks/useWebSocket";
import TokenDialog from "./TokenDialog";
import "./Grid.css";

const DEFAULT_ZOOM = 40;
const MIN_ZOOM = 8;
const MAX_ZOOM = 160;
const ZOOM_STEP = 1.12;

const BG_COLOR = "#111827";
const GRID_COLOR = "rgba(255,255,255,0.07)";
const GRID_COLOR_MAJOR = "rgba(255,255,255,0.16)";

/** BFS stops at this many cells; region is considered unbounded → paint only the clicked cell. */
const MAX_FLOOD_CELLS = 2500;

/** World units per second for WASD keyboard panning. */
const WASD_PAN_SPEED = 12;

interface Camera {
  offsetX: number;
  offsetY: number;
  zoom: number;
}

interface Props {
  mapId: string;
  tool: Tool;
  paintColor: string;
  tokenColor: string;
  onColorsLoaded: (colors: string[]) => void;
}

export interface GridHandle {
  sendColorUpdate: (colors: string[]) => void;
}

function cellKey(x: number, y: number): string {
  return `${x},${y}`;
}

function canvasToCell(
  cx: number,
  cy: number,
  cam: Camera,
): { x: number; y: number } {
  return {
    x: Math.floor(cx / cam.zoom + cam.offsetX),
    y: Math.floor(cy / cam.zoom + cam.offsetY),
  };
}

function cellToCanvas(
  cellX: number,
  cellY: number,
  cam: Camera,
): { x: number; y: number } {
  return {
    x: (cellX - cam.offsetX) * cam.zoom,
    y: (cellY - cam.offsetY) * cam.zoom,
  };
}

function drawToken(
  ctx: CanvasRenderingContext2D,
  cellX: number,
  cellY: number,
  label: string,
  color: string,
  cam: Camera,
) {
  const zoom = cam.zoom;
  const { x: px, y: py } = cellToCanvas(cellX, cellY, cam);
  const cx = px + zoom / 2;
  const cy = py + zoom / 2;
  const r = zoom * 0.38;

  ctx.shadowColor = "rgba(0,0,0,0.6)";
  ctx.shadowBlur = 5;

  ctx.beginPath();
  ctx.arc(cx, cy, r, 0, Math.PI * 2);
  ctx.fillStyle = color;
  ctx.fill();

  ctx.strokeStyle = "rgba(255,255,255,0.55)";
  ctx.lineWidth = Math.max(1, zoom * 0.04);
  ctx.stroke();

  ctx.shadowBlur = 0;

  if (zoom >= 14) {
    const words = label.trim().split(/\s+/);
    const initials =
      words.length >= 2
        ? (words[0][0] + words[1][0]).toUpperCase()
        : label.slice(0, 2).toUpperCase();
    ctx.fillStyle = "#ffffff";
    ctx.font = `bold ${Math.round(zoom * 0.3)}px system-ui, sans-serif`;
    ctx.textAlign = "center";
    ctx.textBaseline = "middle";
    ctx.fillText(initials, cx, cy);
  }
}

function clamp(v: number, lo: number, hi: number) {
  return Math.max(lo, Math.min(hi, v));
}

/**
 * BFS flood fill from (startX, startY) through uncolored cells.
 * Returns the list of cells to fill, or null if the region is unbounded
 * (search exceeded MAX_FLOOD_CELLS).
 */
function floodFill(
  startX: number,
  startY: number,
  cells: Map<string, GridCell>,
): Array<{ x: number; y: number }> | null {
  const visited = new Set<string>();
  const queue: Array<{ x: number; y: number }> = [{ x: startX, y: startY }];
  const found: Array<{ x: number; y: number }> = [];

  visited.add(cellKey(startX, startY));

  const dirs = [
    { dx: 1, dy: 0 },
    { dx: -1, dy: 0 },
    { dx: 0, dy: 1 },
    { dx: 0, dy: -1 },
  ];

  while (queue.length > 0) {
    const { x, y } = queue.shift()!;
    found.push({ x, y });

    if (found.length > MAX_FLOOD_CELLS) return null; // unbounded

    for (const { dx, dy } of dirs) {
      const nx = x + dx;
      const ny = y + dy;
      const key = cellKey(nx, ny);
      if (!visited.has(key) && !cells.has(key)) {
        visited.add(key);
        queue.push({ x: nx, y: ny });
      }
    }
  }

  return found;
}

/**
 * Clamp the camera so at least one colored cell or token remains visible in
 * the current viewport. No-op when the map has no content yet.
 */
function clampCameraToContent(
  cam: Camera,
  cells: Map<string, GridCell>,
  tokens: Map<string, GridToken>,
  canvasW: number,
  canvasH: number,
) {
  if (cells.size === 0 && tokens.size === 0) return;

  const viewLeft = cam.offsetX;
  const viewRight = cam.offsetX + canvasW / cam.zoom;
  const viewTop = cam.offsetY;
  const viewBottom = cam.offsetY + canvasH / cam.zoom;

  // Quick visibility check
  let anyVisible = false;

  cellLoop: for (const cell of cells.values()) {
    if (
      cell.x + 1 > viewLeft &&
      cell.x < viewRight &&
      cell.y + 1 > viewTop &&
      cell.y < viewBottom
    ) {
      anyVisible = true;
      break cellLoop;
    }
  }

  if (!anyVisible) {
    for (const tok of tokens.values()) {
      if (
        tok.x + 1 > viewLeft &&
        tok.x < viewRight &&
        tok.y + 1 > viewTop &&
        tok.y < viewBottom
      ) {
        anyVisible = true;
        break;
      }
    }
  }

  if (anyVisible) return;

  // Find the bounding box of all content
  let minX = Infinity,
    maxX = -Infinity;
  let minY = Infinity,
    maxY = -Infinity;

  for (const c of cells.values()) {
    if (c.x < minX) minX = c.x;
    if (c.x > maxX) maxX = c.x;
    if (c.y < minY) minY = c.y;
    if (c.y > maxY) maxY = c.y;
  }
  for (const t of tokens.values()) {
    if (t.x < minX) minX = t.x;
    if (t.x > maxX) maxX = t.x;
    if (t.y < minY) minY = t.y;
    if (t.y > maxY) maxY = t.y;
  }

  const viewW = canvasW / cam.zoom;
  const viewH = canvasH / cam.zoom;

  // X: bring the nearest content edge to the nearest viewport edge
  if (maxX + 1 <= viewLeft) {
    // All content is to the left — show rightmost cell at the left edge
    cam.offsetX = maxX;
  } else if (minX >= viewRight) {
    // All content is to the right — show leftmost cell at the right edge
    cam.offsetX = minX - viewW + 1;
  }

  // Y: same logic
  if (maxY + 1 <= viewTop) {
    cam.offsetY = maxY;
  } else if (minY >= viewBottom) {
    cam.offsetY = minY - viewH + 1;
  }
}

const Grid = forwardRef<GridHandle, Props>(function Grid(
  { mapId, tool, paintColor, tokenColor, onColorsLoaded },
  ref,
) {
  const containerRef = useRef<HTMLDivElement>(null);
  const canvasRef = useRef<HTMLCanvasElement>(null);

  const cameraRef = useRef<Camera>({
    offsetX: -2,
    offsetY: -2,
    zoom: DEFAULT_ZOOM,
  });

  const cellsRef = useRef<Map<string, GridCell>>(new Map());
  const tokensRef = useRef<Map<string, GridToken>>(new Map());

  const [tick, setTick] = useState(0);
  const redraw = useCallback(() => setTick((n) => n + 1), []);

  // ---- Mouse interaction state (refs to avoid stale closures) ----
  const isPanning = useRef(false);
  const panStart = useRef<{
    mx: number;
    my: number;
    ox: number;
    oy: number;
  } | null>(null);
  const isDrawing = useRef(false);
  const isErasing = useRef(false);
  const lastPainted = useRef<string | null>(null);
  const isDragging = useRef(false);
  const dragTokenId = useRef<string | null>(null);
  const dragCellPos = useRef<{ x: number; y: number } | null>(null);

  // ---- WASD state ----
  const keysHeld = useRef<Set<string>>(new Set());
  const rafId = useRef<number | null>(null);
  const lastFrameTime = useRef<number | null>(null);

  // ---- Stable send ref so handlers never go stale ----
  const sendRef = useRef<(msg: ClientMessage) => void>(() => {});

  const [cursor, setCursor] = useState<string>("default");
  const [dialogPos, setDialogPos] = useState<{ x: number; y: number } | null>(
    null,
  );

  useImperativeHandle(ref, () => ({
    sendColorUpdate(colors: string[]) {
      sendRef.current({ type: "update_colors", colors });
    },
  }));

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

    const resize = () => {
      canvas.width = container.clientWidth;
      canvas.height = container.clientHeight;
      redraw();
    };

    resize();
    const observer = new ResizeObserver(resize);
    observer.observe(container);
    return () => observer.disconnect();
  }, [redraw]);

  // Keep a stable ref to the callback so handleMessage doesn't re-create
  const onColorsLoadedRef = useRef(onColorsLoaded);
  useEffect(() => {
    onColorsLoadedRef.current = onColorsLoaded;
  }, [onColorsLoaded]);

  const handleMessage = useCallback(
    (msg: ServerMessage) => {
      switch (msg.type) {
        case "state": {
          cellsRef.current.clear();
          tokensRef.current.clear();
          msg.cells.forEach((c) => cellsRef.current.set(cellKey(c.x, c.y), c));
          msg.tokens.forEach((t) => tokensRef.current.set(t.id, t));
          onColorsLoadedRef.current(msg.colors);
          redraw();
          break;
        }
        case "cell_painted": {
          const key = cellKey(msg.x, msg.y);
          cellsRef.current.set(key, {
            map_id: mapId,
            x: msg.x,
            y: msg.y,
            color: msg.color,
          });
          redraw();
          break;
        }
        case "cells_batch_painted": {
          msg.cells.forEach((c) => {
            const key = cellKey(c.x, c.y);
            cellsRef.current.set(key, {
              map_id: mapId,
              x: c.x,
              y: c.y,
              color: c.color,
            });
          });
          redraw();
          break;
        }
        case "cell_erased": {
          cellsRef.current.delete(cellKey(msg.x, msg.y));
          redraw();
          break;
        }
        case "token_added": {
          tokensRef.current.set(msg.id, {
            id: msg.id,
            map_id: mapId,
            x: msg.x,
            y: msg.y,
            label: msg.label,
            color: msg.color,
          });
          redraw();
          break;
        }
        case "token_moved": {
          const tok = tokensRef.current.get(msg.id);
          if (tok) {
            tokensRef.current.set(msg.id, { ...tok, x: msg.x, y: msg.y });
            redraw();
          }
          break;
        }
        case "token_deleted": {
          tokensRef.current.delete(msg.id);
          redraw();
          break;
        }
        case "colors_updated": {
          onColorsLoadedRef.current(msg.colors);
          break;
        }
        case "error":
          console.error("[Grid WS]", msg.message);
          break;
      }
    },
    [mapId, redraw],
  );

  const { send } = useWebSocket(mapId, handleMessage);
  useEffect(() => {
    sendRef.current = send;
  }, [send]);

  useEffect(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;
    const ctx = canvas.getContext("2d");
    if (!ctx) return;

    const W = canvas.width;
    const H = canvas.height;
    const cam = cameraRef.current;
    const { offsetX, offsetY, zoom } = cam;

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

    const startCX = Math.floor(offsetX) - 1;
    const endCX = Math.ceil(offsetX + W / zoom) + 1;
    const startCY = Math.floor(offsetY) - 1;
    const endCY = Math.ceil(offsetY + H / zoom) + 1;

    // Painted cells
    cellsRef.current.forEach((cell) => {
      if (
        cell.x < startCX ||
        cell.x > endCX ||
        cell.y < startCY ||
        cell.y > endCY
      )
        return;
      const { x: px, y: py } = cellToCanvas(cell.x, cell.y, cam);
      ctx.fillStyle = cell.color;
      ctx.fillRect(px, py, zoom, zoom);
    });

    // Grid lines
    if (zoom >= 5) {
      ctx.lineWidth = 1;
      for (let x = startCX; x <= endCX; x++) {
        ctx.strokeStyle = x % 5 === 0 ? GRID_COLOR_MAJOR : GRID_COLOR;
        const px = (x - offsetX) * zoom + 0.5;
        ctx.beginPath();
        ctx.moveTo(px, 0);
        ctx.lineTo(px, H);
        ctx.stroke();
      }
      for (let y = startCY; y <= endCY; y++) {
        ctx.strokeStyle = y % 5 === 0 ? GRID_COLOR_MAJOR : GRID_COLOR;
        const py = (y - offsetY) * zoom + 0.5;
        ctx.beginPath();
        ctx.moveTo(0, py);
        ctx.lineTo(W, py);
        ctx.stroke();
      }
    }

    // Tokens (skip the one being dragged)
    tokensRef.current.forEach((token) => {
      if (isDragging.current && dragTokenId.current === token.id) return;
      drawToken(ctx, token.x, token.y, token.label, token.color, cam);
    });

    // Drag ghost
    if (isDragging.current && dragCellPos.current && dragTokenId.current) {
      const tok = tokensRef.current.get(dragTokenId.current);
      if (tok) {
        ctx.globalAlpha = 0.6;
        drawToken(
          ctx,
          dragCellPos.current.x,
          dragCellPos.current.y,
          tok.label,
          tok.color,
          cam,
        );
        ctx.globalAlpha = 1;
      }
    }
  }, [tick]);

  function applyClampAndRedraw() {
    const canvas = canvasRef.current;
    if (canvas) {
      clampCameraToContent(
        cameraRef.current,
        cellsRef.current,
        tokensRef.current,
        canvas.width,
        canvas.height,
      );
    }
    redraw();
  }

  function applyZoom(canvasX: number, canvasY: number, factor: number) {
    const cam = cameraRef.current;
    const worldX = canvasX / cam.zoom + cam.offsetX;
    const worldY = canvasY / cam.zoom + cam.offsetY;
    const newZoom = clamp(cam.zoom * factor, MIN_ZOOM, MAX_ZOOM);
    cam.offsetX = worldX - canvasX / newZoom;
    cam.offsetY = worldY - canvasY / newZoom;
    cam.zoom = newZoom;
    applyClampAndRedraw();
  }

  useEffect(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;
    const onWheel = (e: WheelEvent) => {
      e.preventDefault();
      const rect = canvas.getBoundingClientRect();
      const cx = e.clientX - rect.left;
      const cy = e.clientY - rect.top;
      const factor = e.deltaY < 0 ? ZOOM_STEP : 1 / ZOOM_STEP;
      applyZoom(cx, cy, factor);
    };
    canvas.addEventListener("wheel", onWheel, { passive: false });
    return () => canvas.removeEventListener("wheel", onWheel);
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  useEffect(() => {
    function rafTick(timestamp: number) {
      const keys = keysHeld.current;
      if (keys.size === 0) {
        lastFrameTime.current = null;
        rafId.current = null;
        return;
      }

      const dt =
        lastFrameTime.current !== null
          ? (timestamp - lastFrameTime.current) / 1000
          : 0;
      lastFrameTime.current = timestamp;

      const cam = cameraRef.current;
      const speed = WASD_PAN_SPEED;

      if (keys.has("a")) cam.offsetX -= speed * dt;
      if (keys.has("d")) cam.offsetX += speed * dt;
      if (keys.has("w")) cam.offsetY -= speed * dt;
      if (keys.has("s")) cam.offsetY += speed * dt;

      const canvas = canvasRef.current;
      if (canvas) {
        clampCameraToContent(
          cam,
          cellsRef.current,
          tokensRef.current,
          canvas.width,
          canvas.height,
        );
      }
      setTick((n) => n + 1);

      rafId.current = requestAnimationFrame(rafTick);
    }

    function onKeyDown(e: KeyboardEvent) {
      if (
        e.target instanceof HTMLInputElement ||
        e.target instanceof HTMLTextAreaElement
      )
        return;
      // Don't intercept WASD when modifier keys are held (e.g. Shift+keys are for tool shortcuts)
      if (e.shiftKey || e.ctrlKey || e.metaKey || e.altKey) return;
      const key = e.key.toLowerCase();
      if (["w", "a", "s", "d"].includes(key)) {
        e.preventDefault();
        keysHeld.current.add(key);
        if (rafId.current === null) {
          lastFrameTime.current = null;
          rafId.current = requestAnimationFrame(rafTick);
        }
      }
    }

    function onKeyUp(e: KeyboardEvent) {
      keysHeld.current.delete(e.key.toLowerCase());
    }

    window.addEventListener("keydown", onKeyDown);
    window.addEventListener("keyup", onKeyUp);
    return () => {
      window.removeEventListener("keydown", onKeyDown);
      window.removeEventListener("keyup", onKeyUp);
      if (rafId.current !== null) {
        cancelAnimationFrame(rafId.current);
        rafId.current = null;
      }
    };
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  function getCanvasPoint(e: React.MouseEvent) {
    const rect = canvasRef.current!.getBoundingClientRect();
    return { x: e.clientX - rect.left, y: e.clientY - rect.top };
  }

  function tokenAtCell(x: number, y: number): GridToken | null {
    for (const t of tokensRef.current.values()) {
      if (t.x === x && t.y === y) return t;
    }
    return null;
  }

  function handleMouseDown(e: React.MouseEvent) {
    e.preventDefault();
    const { x: mx, y: my } = getCanvasPoint(e);
    const cell = canvasToCell(mx, my, cameraRef.current);

    // ---- Pan tool ----
    if (tool === "pan" && e.button === 0) {
      isPanning.current = true;
      panStart.current = {
        mx,
        my,
        ox: cameraRef.current.offsetX,
        oy: cameraRef.current.offsetY,
      };
      setCursor("grabbing");
      return;
    }

    // ---- Zoom tool ----
    if (tool === "zoom") {
      if (e.button === 0) applyZoom(mx, my, ZOOM_STEP * ZOOM_STEP);
      else if (e.button === 2) applyZoom(mx, my, 1 / (ZOOM_STEP * ZOOM_STEP));
      return;
    }

    // ---- Draw tool ----
    if (tool === "draw") {
      if (e.button === 0) {
        if (e.shiftKey) {
          // Shift+click → flood fill uncolored region
          const key = cellKey(cell.x, cell.y);
          if (!cellsRef.current.has(key)) {
            const region = floodFill(cell.x, cell.y, cellsRef.current);
            if (region === null || region.length === 1) {
              // Unbounded or trivially single cell → paint one cell
              sendRef.current({
                type: "paint_cell",
                x: cell.x,
                y: cell.y,
                color: paintColor,
              });
            } else {
              // Bounded enclosed region → batch paint
              sendRef.current({
                type: "paint_cells",
                cells: region.map(({ x, y }) => ({ x, y, color: paintColor })),
              });
            }
          }
        } else {
          isDrawing.current = true;
          lastPainted.current = cellKey(cell.x, cell.y);
          sendRef.current({
            type: "paint_cell",
            x: cell.x,
            y: cell.y,
            color: paintColor,
          });
        }
      } else if (e.button === 2) {
        isErasing.current = true;
        const key = cellKey(cell.x, cell.y);
        lastPainted.current = key;
        if (cellsRef.current.has(key)) {
          sendRef.current({ type: "erase_cell", x: cell.x, y: cell.y });
        }
      }
      return;
    }

    // ---- Token tool ----
    if (tool === "token") {
      if (e.button === 2) {
        const tok = tokenAtCell(cell.x, cell.y);
        if (tok) sendRef.current({ type: "delete_token", id: tok.id });
        return;
      }
      if (e.button === 0) {
        const tok = tokenAtCell(cell.x, cell.y);
        if (tok) {
          isDragging.current = true;
          dragTokenId.current = tok.id;
          dragCellPos.current = { x: cell.x, y: cell.y };
          redraw();
        } else {
          setDialogPos({ x: cell.x, y: cell.y });
        }
      }
    }
  }

  function handleMouseMove(e: React.MouseEvent) {
    const { x: mx, y: my } = getCanvasPoint(e);

    // Pan
    if (isPanning.current && panStart.current) {
      const cam = cameraRef.current;
      const dx = (mx - panStart.current.mx) / cam.zoom;
      const dy = (my - panStart.current.my) / cam.zoom;
      cam.offsetX = panStart.current.ox - dx;
      cam.offsetY = panStart.current.oy - dy;
      applyClampAndRedraw();
      return;
    }

    // Draw / erase stroke
    if (isDrawing.current || isErasing.current) {
      const cell = canvasToCell(mx, my, cameraRef.current);
      const key = cellKey(cell.x, cell.y);
      if (lastPainted.current !== key) {
        lastPainted.current = key;
        if (isDrawing.current) {
          sendRef.current({
            type: "paint_cell",
            x: cell.x,
            y: cell.y,
            color: paintColor,
          });
        } else if (isErasing.current && cellsRef.current.has(key)) {
          sendRef.current({ type: "erase_cell", x: cell.x, y: cell.y });
        }
      }
      return;
    }

    // Token drag
    if (isDragging.current && dragCellPos.current) {
      const cell = canvasToCell(mx, my, cameraRef.current);
      if (
        dragCellPos.current.x !== cell.x ||
        dragCellPos.current.y !== cell.y
      ) {
        dragCellPos.current = { x: cell.x, y: cell.y };
        redraw();
      }
    }
  }

  function handleMouseUp(_e: React.MouseEvent) {
    if (isPanning.current) {
      isPanning.current = false;
      panStart.current = null;
      setCursor("grab");
      return;
    }

    if (isDrawing.current || isErasing.current) {
      isDrawing.current = false;
      isErasing.current = false;
      lastPainted.current = null;
      return;
    }

    if (isDragging.current && dragTokenId.current && dragCellPos.current) {
      const tok = tokensRef.current.get(dragTokenId.current);
      if (
        tok &&
        (tok.x !== dragCellPos.current.x || tok.y !== dragCellPos.current.y)
      ) {
        sendRef.current({
          type: "move_token",
          id: dragTokenId.current,
          x: dragCellPos.current.x,
          y: dragCellPos.current.y,
        });
      }
      isDragging.current = false;
      dragTokenId.current = null;
      dragCellPos.current = null;
      redraw();
    }
  }

  function handleMouseLeave() {
    isPanning.current = false;
    panStart.current = null;
    isDrawing.current = false;
    isErasing.current = false;
    lastPainted.current = null;
    if (isDragging.current) {
      isDragging.current = false;
      dragTokenId.current = null;
      dragCellPos.current = null;
      redraw();
    }
  }

  // Sync cursor CSS to active tool
  useEffect(() => {
    switch (tool) {
      case "pan":
        setCursor("grab");
        break;
      case "zoom":
        setCursor("zoom-in");
        break;
      case "draw":
        setCursor("crosshair");
        break;
      case "token":
        setCursor("crosshair");
        break;
    }
  }, [tool]);

  function handleAddToken(label: string, color: string) {
    if (!dialogPos) return;
    sendRef.current({
      type: "add_token",
      x: dialogPos.x,
      y: dialogPos.y,
      label,
      color,
    });
    setDialogPos(null);
  }

  return (
    <div ref={containerRef} className="grid-container">
      <canvas
        ref={canvasRef}
        className="grid-canvas"
        style={{ cursor }}
        onMouseDown={handleMouseDown}
        onMouseMove={handleMouseMove}
        onMouseUp={handleMouseUp}
        onMouseLeave={handleMouseLeave}
        onContextMenu={(e) => e.preventDefault()}
      />
      {dialogPos && (
        <TokenDialog
          defaultColor={tokenColor}
          onConfirm={handleAddToken}
          onCancel={() => setDialogPos(null)}
        />
      )}
    </div>
  );
});

export default Grid;