import type { Logger } from '../common';
import {
  type Command,
  type EditorState,
  type EditorView,
  Fragment,
  type MarkType,
  Plugin,
  PluginKey,
  PMNode,
  Selection,
  Transform,
} from '../pm';
import {
  clampRange,
  createVirtualElementFromRange,
  findFirstMarkPosition,
  getMarkType,
  isTextSelection,
  type MarkScanResult,
  safeInsert,
  setupScrollAndResizeHandlers,
  type VirtualElement,
} from '../pm-utils';
import { store } from '../store';

export type Suggestion = {
  markName: string;
  trigger: string;
  synthetic: boolean;
  show: boolean;
  text: string;
  position: number;
  refresh: number;
  anchorEl: () => VirtualElement | null;
  selectedIndex: number;
};

export type SuggestionUiHandlers = Record<
  string,
  { onSelect: (selected: Suggestion) => void; optionCount?: number }
>;

type SuggestionsByMark = ReadonlyMap<
  EditorView,
  ReadonlyMap<string, Suggestion>
>;

// for better performance we scan a range around the selection
const MARK_SCAN_RANGE_PADDING = 500;

const $internalSuggestions = store.atom<SuggestionsByMark>(new Map());

const $internalSuggestionUi = store.atom<
  ReadonlyMap<EditorView, SuggestionUiHandlers>
>(new Map());

function suggestionByView(
  suggestionsByMark: SuggestionsByMark,
): ReadonlyMap<EditorView, Suggestion> {
  const next = new Map<EditorView, Suggestion>();
  for (const [view, suggestions] of suggestionsByMark) {
    const values = [...suggestions.values()];
    const active = values.at(-1);
    if (active) {
      next.set(view, active);
    }
  }
  return next;
}

function suggestionMapByMark(
  suggestions: ReadonlyMap<EditorView, Suggestion>,
): SuggestionsByMark {
  const next = new Map<EditorView, ReadonlyMap<string, Suggestion>>();
  for (const [view, suggestion] of suggestions) {
    next.set(view, new Map([[suggestion.markName, suggestion]]));
  }
  return next;
}

export const $suggestions = store.atom(
  (get) => suggestionByView(get($internalSuggestions)),
  (
    get,
    set,
    update:
      | ReadonlyMap<EditorView, Suggestion>
      | ((
          existing: ReadonlyMap<EditorView, Suggestion>,
        ) => ReadonlyMap<EditorView, Suggestion>),
  ) => {
    const existing = suggestionByView(get($internalSuggestions));
    set(
      $internalSuggestions,
      suggestionMapByMark(
        typeof update === 'function' ? update(existing) : update,
      ),
    );
  },
);

export const $suggestion = store.atom((get) => {
  for (const suggestion of suggestionByView(
    get($internalSuggestions),
  ).values()) {
    return suggestion;
  }
  return undefined;
});

export const $suggestionUi = store.atom(
  (get) => get($internalSuggestionUi),
  (
    get,
    set,
    update:
      | ReadonlyMap<EditorView, SuggestionUiHandlers>
      | ((
          existing: ReadonlyMap<EditorView, SuggestionUiHandlers>,
        ) => ReadonlyMap<EditorView, SuggestionUiHandlers>),
  ) => {
    set(
      $internalSuggestionUi,
      typeof update === 'function'
        ? update(get($internalSuggestionUi))
        : update,
    );
  },
);

function getSuggestionForView(
  state: EditorState,
  view: EditorView | undefined,
  markName?: string,
) {
  if (!view) {
    const suggestion = store.get(state, $suggestion);
    return !markName || suggestion?.markName === markName
      ? suggestion
      : undefined;
  }
  const suggestions = store.get(state, $internalSuggestions).get(view);
  if (!suggestions) {
    return undefined;
  }
  if (markName) {
    return suggestions.get(markName);
  }
  return [...suggestions.values()].at(-1);
}

function setSuggestionForView(
  state: EditorState,
  view: EditorView,
  markName: string,
  suggestion: Suggestion | undefined,
) {
  const current = store.get(state, $internalSuggestions);
  const currentForView = current.get(view);
  if (!suggestion && !currentForView?.has(markName)) {
    return;
  }
  const next = new Map(current);
  if (suggestion) {
    const nextForView = new Map(currentForView ?? new Map());
    nextForView.delete(markName);
    nextForView.set(markName, suggestion);
    next.set(view, nextForView);
  } else {
    const nextForView = new Map(currentForView);
    nextForView.delete(markName);
    if (nextForView.size) {
      next.set(view, nextForView);
    } else {
      next.delete(view);
    }
  }
  store.set(state, $internalSuggestions, next);
}

export function updateSuggestionForView(
  state: EditorState,
  view: EditorView,
  markName: string,
  update: (suggestion: Suggestion) => Suggestion,
) {
  const suggestion = getSuggestionForView(state, view, markName);
  if (!suggestion) {
    return false;
  }
  setSuggestionForView(state, view, markName, update(suggestion));
  return true;
}

// TODO current if the query is /sdsd /s dasdlkasd and the user moves out the second subquery is selected
export function pluginSuggestion({
  markName,
  trigger,
  logger,
  endOnWhitespace = false,
  endOnQuery = false,
}: {
  markName: string;
  trigger: string;
  logger?: Logger;
  endOnWhitespace?: boolean;
  endOnQuery?: boolean;
}) {
  return new Plugin({
    key: new PluginKey(`suggestion-${markName}`),
    view: (view) => {
      const abortController = new AbortController();

      setupScrollAndResizeHandlers((refresh_counter) => {
        const suggestion = getSuggestionForView(view.state, view, markName);
        if (suggestion) {
          setSuggestionForView(view.state, view, markName, {
            ...suggestion,
            refresh: refresh_counter,
          });
        }
      }, abortController.signal);

      return {
        destroy: () => {
          abortController.abort();
          setSuggestionForView(view.state, view, markName, undefined);
        },
        update: (view, lastState) => {
          const { state } = view;
          if (lastState === state) {
            return;
          }

          const markType = getMarkType(state.schema, markName);

          if (
            lastState.doc.eq(state.doc) &&
            state.selection.eq(lastState.selection)
          ) {
            logger?.debug(
              'suggestion:selection and doc not changed',
              state.selection.toJSON(),
              lastState.selection.toJSON(),
            );
            return;
          }

          const wasMarkActive = isMarkActiveAtSelection(lastState, markType);
          const isMarkActive = isMarkActiveAtSelection(state, markType);

          // Remove the mark if the user has moved to a different place
          if (!isMarkActive && wasMarkActive) {
            logger?.debug('suggestion:mark was active mark removed');
            removeSuggestMark({
              markName,
              selection: lastState.selection,
            })(state, view.dispatch, view);
            return;
          }

          // Clear the mark if the trigger is deleted but the marked text remains
          if (isMarkActive && !doesQueryHaveTrigger(state, markType, trigger)) {
            logger?.debug('suggestion:trigger deleted but marked text remains');
            removeSuggestMark({
              markName,
              selection: state.selection,
            })(state, view.dispatch, view);
            return;
          }

          const suggestion = getSuggestionForView(state, view, markName);
          if (isMarkActive) {
            const [scanRangeStart, scanRangeEnd] = clampRange(
              state.selection.from - MARK_SCAN_RANGE_PADDING,
              state.selection.to + MARK_SCAN_RANGE_PADDING,
              state.doc.content.size,
            );
            const result = querySuggestionMarkTextContent(
              markName,
              state,
              scanRangeStart,
              scanRangeEnd,
            );
            if (!result) {
              return;
            }
            logger?.debug('querytext', result?.text);

            if (endOnQuery && result.text !== trigger) {
              logger?.debug('suggestion:text after trigger ended the query');
              removeSuggestMark({
                markName,
                selection: state.selection,
                trigger,
              })(state, view.dispatch, view);
              return;
            }

            // Whitespace ends the query for providers that opt in (the
            // slash menu): the typed text stays in the document, the menu
            // closes. Providers like wiki links keep spaces in queries.
            if (
              endOnWhitespace &&
              /\s/.test((result.text ?? '').slice(trigger.length))
            ) {
              logger?.debug('suggestion:whitespace ended the query');
              removeSuggestMark({
                markName,
                selection: state.selection,
                trigger,
              })(state, view.dispatch, view);
              return;
            }

            if (
              suggestion?.markName === markName &&
              suggestion.text === result.text &&
              suggestion.position === result.start
            ) {
              return;
            }

            const activeMark = state.doc
              .nodeAt(result.start)
              ?.marks.find((mark) => mark.type === markType);
            setSuggestionForView(state, view, markName, {
              selectedIndex:
                suggestion?.markName === markName &&
                suggestion.text === result.text
                  ? suggestion.selectedIndex
                  : 0,
              markName,
              trigger,
              synthetic: activeMark?.attrs.synthetic === true,
              show: true,
              text: result.text ?? '',
              position: result.start,
              refresh: suggestion?.refresh ?? 0,
              anchorEl: () => {
                const [start, end] = clampRange(
                  result.start,
                  result.end,
                  view.state.doc.content.size,
                );
                return createVirtualElementFromRange(view, start, end);
              },
            });
          } else {
            if (suggestion !== undefined) {
              setSuggestionForView(state, view, markName, undefined);
            }
          }
        },
      };
    },
  });
}

function isMarkActiveAtSelection(state: EditorState, markType: MarkType) {
  const { from, to } = state.selection;
  if (!isTextSelection(state.selection)) {
    return false;
  }
  const clampedFrom = Math.max(0, from - 1);
  return state.doc.rangeHasMark(clampedFrom, to, markType);
}

function doesQueryHaveTrigger(
  state: EditorState,
  markType: MarkType,
  trigger: string,
) {
  const { nodeBefore } = state.selection.$from;
  if (!nodeBefore) {
    return false;
  }
  const suggestMark = markType.isInSet(nodeBefore.marks || []);
  if (!suggestMark) {
    return false;
  }

  const textContent = nodeBefore.textContent || '';
  return textContent.includes(trigger);
}

/**
 * Programmatically start a suggestion at the current cursor, producing the
 * same document state as if the user had typed the trigger text — the
 * suggestion plugin picks the mark up on its next update and shows the UI.
 * The mark is flagged `synthetic` so Escape removes the trigger text too,
 * instead of leaving a stray character the user never typed.
 */
export function openSuggestion({
  markName,
  trigger,
}: {
  markName: string;
  trigger: string;
}): Command {
  return (state, dispatch, view) => {
    const { schema, selection } = state;
    if (!selection.empty || !isTextSelection(selection)) {
      return false;
    }
    const linkMarkType = schema.marks.link;
    if (
      selection.$from.parent.type.spec.code ||
      selection.$from
        .marks()
        .some((mark) => mark.type.spec.code || mark.type === linkMarkType)
    ) {
      return false;
    }

    const mark = schema.mark(markName, { trigger, synthetic: true });
    const marks = selection.$from.marks();
    const tr = state.tr
      .replaceWith(
        selection.from,
        selection.to,
        schema.text(trigger, [mark, ...marks]),
      )
      .addStoredMark(mark)
      .scrollIntoView();

    if (dispatch) {
      if (view && typeof view.focus === 'function') {
        view.focus();
      }
      dispatch(tr);
    }
    return true;
  };
}

export function removeSuggestMark({
  markName,
  selection,
  trigger,
}: {
  markName: string;
  selection: Selection;
  /**
   * The provider's trigger text. Supply it when the query is ending because
   * the user kept typing: a synthetic trigger is then removed on its own and
   * everything typed after it stays. Leave it off to abandon the suggestion
   * outright (Escape, deactivation), which drops the whole synthetic range.
   */
  trigger?: string;
}): Command {
  return (state, dispatch, _view) => {
    const { schema } = state;
    const markType = getMarkType(schema, markName);

    if (!isTextSelection(selection)) {
      return false;
    }

    const queryMark = findFirstMarkPosition(
      markType,
      state.doc,
      selection.from - 1,
      selection.to,
    );

    if (_view) {
      setSuggestionForView(state, _view, markName, undefined);
    }

    if (
      !queryMark &&
      state.storedMarks &&
      markType.isInSet(state.storedMarks)
    ) {
      dispatch?.(
        state.tr.removeStoredMark(markType).setMeta('addToHistory', false),
      );
      return true;
    }

    if (!queryMark) {
      return false;
    }

    // A synthetic trigger (opened programmatically, e.g. the "+" block
    // button) was never typed: deactivating it removes the text too, so a
    // "/" the user never wrote cannot linger in the document. Typed
    // triggers keep their text and only lose the mark.
    const activeMark = state.doc
      .nodeAt(queryMark.start)
      ?.marks.find((mark) => mark.type === markType);
    if (activeMark?.attrs.synthetic) {
      // The mark is inclusive, so once the user types the range covers those
      // characters too. Deleting all of it would eat what they just typed, so
      // when a trigger is supplied only the trigger text goes and the rest is
      // merely unmarked.
      const triggerEnd =
        trigger === undefined
          ? queryMark.end
          : Math.min(queryMark.start + trigger.length, queryMark.end);
      const tr = state.tr.delete(queryMark.start, triggerEnd);
      if (triggerEnd < queryMark.end) {
        tr.removeMark(
          queryMark.start,
          queryMark.end - (triggerEnd - queryMark.start),
          markType,
        );
      }
      dispatch?.(tr.removeStoredMark(markType).setMeta('addToHistory', false));
      return true;
    }

    dispatch?.(
      state.tr
        .removeMark(queryMark.start, queryMark.end, markType)
        .removeStoredMark(markType)
        .setMeta('addToHistory', false),
    );
    return true;
  };
}

/**
 * Returns a copy of `doc` without any text carrying a synthetic suggestion
 * mark. The save path serializes through this so a programmatically opened
 * trigger (e.g. the "+" block button's "/") can never reach storage, even
 * when the editor unmounts or a debounced save fires while the menu is
 * still open.
 */
export function stripSyntheticSuggestionText(doc: PMNode): PMNode {
  const ranges: Array<[number, number]> = [];
  doc.descendants((node, pos) => {
    if (
      node.isText &&
      node.marks.some(
        (mark) =>
          mark.type.spec.group === 'suggestTriggerMarks' &&
          mark.attrs.synthetic,
      )
    ) {
      ranges.push([pos, pos + node.nodeSize]);
    }
    return true;
  });
  if (ranges.length === 0) {
    return doc;
  }
  const transform = new Transform(doc);
  for (const [from, to] of ranges.reverse()) {
    transform.delete(from, to);
  }
  return transform.doc;
}

/**
 * Returns the text content **within** the mark
 */
export function querySuggestionMarkTextContent(
  markName: string,
  state: EditorState,
  startPos: number,
  endPos: number,
): undefined | MarkScanResult {
  const { schema, doc } = state;
  const markType = getMarkType(schema, markName);

  const scan = findFirstMarkPosition(markType, doc, startPos, endPos);

  if (!scan) {
    return undefined;
  }

  return scan;
}

export type ReplacementContent = string | PMNode | Fragment;

export function replaceSuggestMarkWith({
  markName,
  content,
  focus = true,
}: {
  markName: string;
  content?: ReplacementContent;
  focus?: boolean;
}): Command {
  return (state, dispatch, view) => {
    const { schema } = state;
    const markType = getMarkType(schema, markName);

    const suggestion = getSuggestionForView(state, view, markName);
    if (!suggestion) {
      return false;
    }

    const suggestionEnd = suggestion.position + suggestion.text.length;
    const [scanRangeStart, scanRangeEnd] = clampRange(
      suggestion.position,
      suggestionEnd,
      state.doc.content.size,
    );

    const queryMark = findFirstMarkPosition(
      markType,
      state.doc,
      scanRangeStart,
      scanRangeEnd,
    );

    if (
      !queryMark ||
      queryMark.start !== suggestion.position ||
      queryMark.end !== suggestionEnd ||
      queryMark.text !== suggestion.text ||
      state.selection.from < queryMark.start ||
      state.selection.to > queryMark.end
    ) {
      return false;
    }

    let tr = state.tr.removeStoredMark(markType);

    // If no content is provided, simply remove the mark text
    if (!content) {
      tr = tr.delete(queryMark.start, queryMark.end);
      dispatch?.(tr);
      return true;
    }

    try {
      // Caret target in NEW-doc coordinates: nothing before queryMark.start
      // moves, so the caret belongs right after the inserted content. This
      // must not go through tr.mapping — the offset is already post-replace,
      // and mapping it again double-shifts when the inserted content is
      // longer than the replaced query text.
      let pos: number;
      if (typeof content === 'string') {
        tr = tr.replaceWith(
          queryMark.start,
          queryMark.end,
          schema.text(content),
        );
        pos = queryMark.start + content.length;
      } else if (content instanceof Fragment) {
        tr = tr.replaceWith(queryMark.start, queryMark.end, content);
        pos = queryMark.start + content.size;
      } else if (content instanceof PMNode) {
        if (content.isText) {
          tr = tr.replaceWith(queryMark.start, queryMark.end, content);
          pos = queryMark.start + content.nodeSize;
        } else if (content.isBlock) {
          tr = safeInsert(content)(tr);
          pos = tr.mapping.map(queryMark.start + 1);
        } else if (content.isInline) {
          const fragment = Fragment.fromArray([content, schema.text(' ')]);
          tr = tr.replaceWith(queryMark.start, queryMark.end, fragment);
          pos = queryMark.start + fragment.size;
        } else {
          console.warn('Unknown content type, skipping replacement');
          return false;
        }
      } else {
        console.warn('Unknown content type, skipping replacement');
        return false;
      }

      pos = Math.min(pos, tr.doc.content.size);
      tr = tr.setSelection(Selection.near(tr.doc.resolve(pos)));

      if (dispatch) {
        if (focus && view && typeof view.focus === 'function') {
          view.focus();
        }
        dispatch(tr);
      }

      return true;
    } catch (error) {
      console.error('Error replacing suggestion mark:', error);
      return false;
    }
  };
}