import { Extension } from "@tiptap/core";
import { Plugin, PluginKey, type EditorState, type Transaction } from "@tiptap/pm/state";
import { Decoration, DecorationSet } from "@tiptap/pm/view";

export type BlockSelectionState = {
  anchorIndex: number;
  headIndex: number;
} | null;

export const blockSelectionKey = new PluginKey<BlockSelectionState>(
  "blockSelection",
);

function selectedRange(
  state: BlockSelectionState,
): { from: number; to: number } | null {
  if (!state) return null;
  return {
    from: Math.min(state.anchorIndex, state.headIndex),
    to: Math.max(state.anchorIndex, state.headIndex),
  };
}

/** Returns doc positions spanning the selected block range. */
export function getBlockRange(
  editorState: EditorState,
  sel: BlockSelectionState,
): { fromPos: number; toPos: number } | null {
  if (!sel) return null;
  const range = selectedRange(sel)!;
  const doc = editorState.doc;
  let fromPos = 0;
  let toPos = 0;
  let idx = 0;
  doc.forEach((node, offset) => {
    if (idx === range.from) fromPos = offset;
    if (idx === range.to) toPos = offset + node.nodeSize;
    idx++;
  });
  return { fromPos, toPos };
}

function isPrintable(e: KeyboardEvent): boolean {
  if (e.ctrlKey || e.metaKey || e.altKey) return false;
  return e.key.length === 1;
}

export const BlockSelectionExtension = Extension.create({
  name: "blockSelection",

  addProseMirrorPlugins() {
    return [
      new Plugin<BlockSelectionState>({
        key: blockSelectionKey,

        state: {
          init(): BlockSelectionState {
            return null;
          },
          apply(tr: Transaction, value: BlockSelectionState): BlockSelectionState {
            const meta = tr.getMeta(blockSelectionKey);
            if (meta !== undefined) return meta;
            // Doc changed while selection active → clear (positions stale)
            if (value && tr.docChanged) return null;
            // User set a new text selection (not from our plugin) → clear
            if (value && tr.selectionSet && !tr.getMeta("blockSelectionInternal")) {
              return null;
            }
            return value;
          },
        },

        props: {
          decorations(state: EditorState): DecorationSet {
            const sel = blockSelectionKey.getState(state);
            const range = selectedRange(sel);
            if (!range) return DecorationSet.empty;

            const decorations: Decoration[] = [];
            let idx = 0;
            state.doc.forEach((node, pos) => {
              if (idx >= range.from && idx <= range.to) {
                decorations.push(
                  Decoration.node(pos, pos + node.nodeSize, {
                    class: "block-selected",
                  }),
                );
              }
              idx++;
            });
            return DecorationSet.create(state.doc, decorations);
          },

          attributes(state: EditorState): Record<string, string> | null {
            const sel = blockSelectionKey.getState(state);
            if (sel) return { class: "has-block-selection" };
            return null;
          },

          handleKeyDown(view, event) {
            const sel = blockSelectionKey.getState(view.state);
            if (!sel) return false;

            const childCount = view.state.doc.childCount;

            if (event.key === "ArrowDown" && event.shiftKey) {
              event.preventDefault();
              const newHead = Math.min(sel.headIndex + 1, childCount - 1);
              const tr = view.state.tr.setMeta(blockSelectionKey, {
                anchorIndex: sel.anchorIndex,
                headIndex: newHead,
              });
              tr.setMeta("blockSelectionInternal", true);
              view.dispatch(tr);
              return true;
            }

            if (event.key === "ArrowUp" && event.shiftKey) {
              event.preventDefault();
              const newHead = Math.max(sel.headIndex - 1, 0);
              const tr = view.state.tr.setMeta(blockSelectionKey, {
                anchorIndex: sel.anchorIndex,
                headIndex: newHead,
              });
              tr.setMeta("blockSelectionInternal", true);
              view.dispatch(tr);
              return true;
            }

            if (event.key === "Escape") {
              event.preventDefault();
              view.dispatch(
                view.state.tr.setMeta(blockSelectionKey, null),
              );
              return true;
            }

            if (event.key === "Backspace" || event.key === "Delete") {
              event.preventDefault();
              const range = selectedRange(sel);
              if (!range) return true;
              const blockRange = getBlockRange(view.state, sel);
              if (!blockRange) return true;
              const tr = view.state.tr.delete(blockRange.fromPos, blockRange.toPos);
              tr.setMeta(blockSelectionKey, null);
              view.dispatch(tr);
              return true;
            }

            if (isPrintable(event)) {
              // Delete selected blocks, clear selection, let PM handle char insertion
              const blockRange = getBlockRange(view.state, sel);
              if (blockRange) {
                const tr = view.state.tr.delete(blockRange.fromPos, blockRange.toPos);
                tr.setMeta(blockSelectionKey, null);
                view.dispatch(tr);
              }
              return false;
            }

            // Modifier-only keys (Shift, Ctrl, Alt, Meta) — ignore
            if (["Shift", "Control", "Alt", "Meta"].includes(event.key)) {
              return false;
            }

            // Any other key — clear selection and pass through
            view.dispatch(
              view.state.tr.setMeta(blockSelectionKey, null),
            );
            return false;
          },

          handleClick(view) {
            const sel = blockSelectionKey.getState(view.state);
            if (sel) {
              view.dispatch(
                view.state.tr.setMeta(blockSelectionKey, null),
              );
            }
            return false;
          },
        },
      }),
    ];
  },
});
