import { useState, useRef, useCallback } from "react";
import type { Editor } from "@tiptap/react";
import { GripVertical, Plus } from "lucide-react";
import { NodeSelection, TextSelection } from "@tiptap/pm/state";
import { Fragment, Slice, type Node as PmNode } from "@tiptap/pm/model";
import {
  blockSelectionKey,
  getBlockRange,
} from "./BlockSelectionExtension";

interface BlockDragHandleProps {
  editor: Editor | null;
  children: React.ReactNode;
}

export function BlockDragHandle({ editor, children }: BlockDragHandleProps) {
  const blockIndexRef = useRef<number | null>(null);
  const savedBlockSelRef = useRef<{
    anchorIndex: number;
    headIndex: number;
  } | null>(null);
  const blockElRef = useRef<HTMLElement | null>(null);
  const [handlePos, setHandlePos] = useState<{
    top: number;
    height: number;
  } | null>(null);

  // Track which block the mouse is over
  const onMouseMove = useCallback(
    (e: React.MouseEvent) => {
      // If hovering the handle itself, keep current position
      if ((e.target as HTMLElement).closest("[data-block-handle-row]")) return;

      const editorEl = (e.currentTarget as HTMLElement).querySelector(
        ".ProseMirror",
      );
      if (!editorEl || !editor) return;

      // Walk from event target up to a direct child of .ProseMirror
      let target = e.target as HTMLElement;
      while (
        target &&
        target !== editorEl &&
        target.parentElement !== editorEl
      ) {
        target = target.parentElement!;
      }

      if (
        target &&
        target !== editorEl &&
        target.parentElement === editorEl
      ) {
        blockElRef.current = target;
        const editorRect = editorEl.getBoundingClientRect();
        const blockRect = target.getBoundingClientRect();
        setHandlePos({
          top: blockRect.top - editorRect.top,
          height: blockRect.height,
        });

        // Track top-level block index for block selection
        try {
          const pos = editor.view.posAtDOM(target, 0);
          blockIndexRef.current = editor.view.state.doc
            .resolve(pos)
            .index(0);
        } catch {
          blockIndexRef.current = null;
        }
      }
      // Don't clear -- only onMouseLeave clears
    },
    [editor],
  );

  const onMouseLeave = useCallback(() => {
    setHandlePos(null);
    blockElRef.current = null;
    blockIndexRef.current = null;
  }, []);

  // Click on drag handle -> select block (Shift+click extends)
  const onHandleClick = useCallback(
    (e: React.MouseEvent) => {
      if (!editor) return;
      const idx = blockIndexRef.current;
      if (idx == null) return;

      // Use saved state from mousedown (PM may have cleared it due to focus change)
      const existing = savedBlockSelRef.current;

      let newSel;
      if (e.shiftKey && existing) {
        newSel = { anchorIndex: existing.anchorIndex, headIndex: idx };
      } else {
        newSel = { anchorIndex: idx, headIndex: idx };
      }

      const tr = editor.view.state.tr.setMeta(blockSelectionKey, newSel);
      tr.setMeta("blockSelectionInternal", true);
      editor.view.dispatch(tr);
      // Refocus editor so Shift+Arrow keys reach PM's handleKeyDown
      editor.view.focus();
    },
    [editor],
  );

  // Add a new empty paragraph below the hovered block
  const onHandleAdd = useCallback(() => {
    if (!editor || !blockElRef.current) return;
    const view = editor.view;
    try {
      const pos = view.posAtDOM(blockElRef.current, 0);
      const $pos = view.state.doc.resolve(pos);
      const after = $pos.after($pos.depth);
      const paragraph = view.state.schema.nodes.paragraph.create();
      const tr = view.state.tr.insert(after, paragraph);
      // Place cursor inside the new paragraph
      tr.setSelection(TextSelection.create(tr.doc, after + 1));
      view.dispatch(tr);
      view.focus();
    } catch {
      // posAtDOM can throw if the element isn't in the editor
    }
  }, [editor]);

  // Initiate ProseMirror-native drag when handle is dragged
  const onHandleDragStart = useCallback(
    (e: React.DragEvent) => {
      if (!editor || !blockElRef.current) return;

      const view = editor.view;
      const el = blockElRef.current;
      // Use saved state from mousedown (PM may have cleared it due to focus change)
      const bsel = savedBlockSelRef.current;

      try {
        // Multi-block drag: if block selection is active and hovered block is in range
        if (bsel && blockIndexRef.current != null) {
          const from = Math.min(bsel.anchorIndex, bsel.headIndex);
          const to = Math.max(bsel.anchorIndex, bsel.headIndex);

          if (
            blockIndexRef.current >= from &&
            blockIndexRef.current <= to
          ) {
            const blockRange = getBlockRange(view.state, bsel);
            if (blockRange) {
              const nodes: PmNode[] = [];
              let idx = 0;
              view.state.doc.forEach((node) => {
                if (idx >= from && idx <= to) nodes.push(node);
                idx++;
              });

              const sel = TextSelection.create(
                view.state.doc,
                blockRange.fromPos,
                blockRange.toPos,
              );
              const tr = view.state.tr.setSelection(sel);
              tr.setMeta("blockSelectionInternal", true);
              view.dispatch(tr);

              view.dragging = {
                slice: new Slice(Fragment.from(nodes), 0, 0),
                move: true,
              };

              e.dataTransfer.effectAllowed = "move";
              e.dataTransfer.setDragImage(el, 0, 0);
              e.dataTransfer.setData("application/x-pm-drag", "true");
              return;
            }
          }
        }

        // Single-block drag (existing behavior)
        const pos = view.posAtDOM(el, 0);
        const $pos = view.state.doc.resolve(pos);
        const before = $pos.before($pos.depth);

        const sel = NodeSelection.create(view.state.doc, before);
        view.dispatch(view.state.tr.setSelection(sel));

        view.dragging = { slice: sel.content(), move: true };

        e.dataTransfer.effectAllowed = "move";
        e.dataTransfer.setDragImage(el, 0, 0);
        e.dataTransfer.setData("application/x-pm-drag", "true");
      } catch {
        // posAtDOM can throw if the element isn't in the editor
      }
    },
    [editor],
  );

  return (
    <div
      className="relative"
      onMouseMove={onMouseMove}
      onMouseLeave={onMouseLeave}
    >
      {handlePos && (
        <div
          data-block-handle-row
          className="absolute left-0 flex items-start z-10"
          style={{ top: handlePos.top + 1 }}
        >
          <div
            onClick={onHandleAdd}
            onMouseDown={(e) => e.preventDefault()}
            className="flex items-center justify-center w-5 h-6 cursor-pointer rounded hover:bg-muted"
          >
            <Plus className="h-3.5 w-3.5 text-muted-foreground/60" />
          </div>
          <div
            data-drag-handle
            draggable
            onMouseDown={() => {
              if (editor) {
                savedBlockSelRef.current =
                  blockSelectionKey.getState(editor.view.state) ?? null;
              }
            }}
            onClick={onHandleClick}
            onDragStart={onHandleDragStart}
            className="flex items-center justify-center w-5 h-6 cursor-grab rounded hover:bg-muted"
          >
            <GripVertical className="h-3.5 w-3.5 text-muted-foreground/60" />
          </div>
        </div>
      )}
      {children}
    </div>
  );
}
