//
// Copyright 2023 DXOS.org
//

import { invertedEffects } from '@codemirror/commands';
import { getChunks } from '@codemirror/merge';
import {
  type ChangeDesc,
  EditorSelection,
  type Extension,
  StateEffect,
  StateField,
  type Text,
} from '@codemirror/state';
import {
  type Command,
  Decoration,
  EditorView,
  type Rect,
  RectangleMarker,
  ViewPlugin,
  keymap,
  layer,
} from '@codemirror/view';
import sortBy from 'lodash.sortby';

import { type CleanupFn, debounce } from '@dxos/async';
import { log } from '@dxos/log';
import { isNonNullable } from '@dxos/util';

import { type Comment, type Range } from '../../types/index.ts';
import { Cursor, singleValueFacet, wrapWithCatch } from '../../util/index.ts';
import { documentId } from '../state/selection.ts';

//
// Options
//

export type CommentsOptions = {
  /**
   * Document id.
   */
  id: string;
  /**
   * Key shortcut to create a new thread.
   */
  key?: string;
  /**
   * Get the current comments.
   */
  getComments?: () => Comment[];
  /**
   * Subscribe to external comment updates.
   */
  subscribe?: (sink: () => void) => CleanupFn;
  /**
   * The branch under review; new comments are tagged with it (undefined = main/unbranched). Scopes
   * a comment to the branch being reviewed so its visibility tracks the active diff.
   */
  reviewBranch?: string;
  /**
   * Prohibit creating new comments (e.g. on a per-user suggestion branch). Existing comments still
   * render; the create command becomes a no-op.
   */
  readonly?: boolean;
  /**
   * Called to create a new thread and return the thread id.
   */
  onCreate?: (params: { cursor: string; from: number; location?: Rect | null; branch?: string }) => void;
  /**
   * Selection cut/deleted.
   */
  onDelete?: (params: { id: string }) => void;
  /**
   * Called when a comment is moved.
   */
  onUpdate?: (params: { id: string; cursor: string }) => void;
  /**
   * Called to notify which thread is currently closest to the cursor.
   */
  onSelect?: (state: CommentsState) => void;
  /**
   * Called when a comment is deliberately activated (clicked). Distinct from `onSelect` (passive
   * cursor-proximity): a click is an explicit user intent, e.g. to reveal the thread in a companion.
   * The editor-side selection/highlight is already applied by the time this fires.
   */
  onActivate?: (id: string) => void;
};

/**
 * Comment threads.
 * 1). Updates the EditorModel to store relative selections for a set of comments threads.
 *     Since the selections are relative, they do not need to be updated when the document is edited.
 * 2). Implements a StateField to track absolute selections corresponding to the comments (i.e., when the document is edited).
 * 3). Creates decoration marks to apply classes to each selection.
 * 4). Tracks the current cursor position to:
 *     a). Update the decoration to show if the cursor is within a current selection.
 *     b). Calls a handler to indicate which is the closest selection (e.g., to update the thread sidebar).
 * 5). Optionally, implements a hoverTooltip to show hints when creating a selection range.
 */
export const comments = (options: CommentsOptions): Extension => {
  const { key: shortcut = "meta-'" } = options;

  const handleSelect = debounce((state: CommentsState) => options.onSelect?.(state), 200);

  return [
    optionsFacet.of(options),
    options.id ? documentId.of(options.id) : undefined,
    commentsState,
    commentsDecorations,
    commentsHighlightLayer,
    handleCommentClick,
    styles,

    //
    // External comment synchronization.
    //
    ViewPlugin.fromClass(
      class {
        private readonly unsubscribe: () => void;
        constructor(view: EditorView) {
          const { id } = options;
          const updateComments = () => {
            const comments = options.getComments?.() ?? [];
            // Local `id` const so the truthiness narrowing survives into the microtask closure.
            if (id && id === view.state.facet(documentId)) {
              queueMicrotask(() => view.dispatch({ effects: setComments.of({ id, comments }) }));
            }
          };

          this.unsubscribe = options.subscribe?.(updateComments) ?? (() => {});
          // Prime the initial decoration pass: `subscribe` may not emit for an already-resolved source
          // (e.g. an ECHO query whose results resolved before mount), so read once now that the source is
          // subscribed. Owning this here means adapters never have to prime manually.
          updateComments();
        }

        destroy = () => {
          this.unsubscribe();
        };
      },
    ),

    //
    // Keymap.
    //
    options.onCreate &&
      keymap.of([
        {
          key: shortcut,
          run: wrapWithCatch(createComment),
        },
      ]),

    //
    // Track deleted ranges and update ranges for decorations.
    //
    EditorView.updateListener.of(({ view, state, changes }) => {
      let mod = false;
      const { comments, ...value } = state.field(commentsState);
      changes.iterChanges((from, to, from2, to2) => {
        comments.forEach(({ comment, range }) => {
          // Only comments with a cursor reach this state (setComments filters cursor-less ones).
          if (!comment.cursor) {
            return;
          }

          // Test if range deleted.
          if (from2 === to2) {
            const newRange = Cursor.getRangeFromCursor(view.state, comment.cursor);
            if (!newRange || newRange.to - newRange.from === 0) {
              options.onDelete?.({ id: comment.id });
            }
          }

          // Update range.
          if (from <= range.to) {
            const newRange = Cursor.getRangeFromCursor(view.state, comment.cursor);
            Object.assign(range, newRange);
            mod = true;
          }
        });
      });

      if (mod) {
        view.dispatch({ effects: setCommentState.of({ comments, ...value }) });
      }
    }),

    //
    // Track selection/proximity.
    //
    EditorView.updateListener.of(({ view, state, transactions }) => {
      let min = Infinity;
      const { selection: applied, comments } = state.field(commentsState);
      const { current, closest } = applied;

      const notify = (selection: SelectionState) =>
        handleSelect({
          selection,
          id: state.facet(documentId),
          comments: comments.map(({ comment, range }) => ({
            comment,
            range,
            location: view.coordsAtPos(range.from),
          })),
        });

      // An explicit `setSelection` is authoritative: selecting a thread in the comments companion sets
      // `current` and moves the caret in one transaction, and the caret does not always land inside the
      // tracked range — proximity would then clear the very selection that transaction just made. Report
      // the applied selection so the app still learns which thread is current, but leave it as set.
      if (transactions.some((tr) => tr.effects.some((effect) => effect.is(setSelection)))) {
        notify(applied);
        return;
      }

      const { head } = state.selection.main;
      const selection: SelectionState = {};
      comments.forEach(({ comment, range }) => {
        if (head >= range.from && head <= range.to) {
          selection.current = comment.id;
          selection.closest = undefined;
        }

        if (!selection.current) {
          const d = Math.min(Math.abs(head - range.from), Math.abs(head - range.to));
          if (d < min) {
            selection.closest = comment.id;
            min = d;
          }
        }
      });

      if (selection.current !== current || selection.closest !== closest) {
        view.dispatch({ effects: setSelection.of(selection) });
        notify(selection);
      }
    }),

    options.onUpdate && trackPastedComments(options.onUpdate),
  ].filter(isNonNullable);
};

//
// State management.
//

type CommentState = {
  comment: Comment;
  range: Range;
  location?: Rect | null;
};

type SelectionState = {
  current?: string;
  closest?: string;
};

export type CommentsState = {
  id?: string;
  comments: CommentState[];
  selection: SelectionState;
};

export const setComments = StateEffect.define<{ id: string; comments: Comment[] }>();

export const setSelection = StateEffect.define<SelectionState>();

const setCommentState = StateEffect.define<CommentsState>();

/**
 * State field (reducer) that tracks comment ranges.
 * The ranges are tracked as Automerge cursors from which the absolute indexed ranges can be computed.
 * Included by {@link comments}; exported for direct state reads.
 */
export const commentsState = StateField.define<CommentsState>({
  create: (state) => ({
    id: state.facet(documentId),
    comments: [],
    selection: {},
  }),
  update: (value, tr) => {
    for (const effect of tr.effects) {
      // Update selection.
      if (effect.is(setSelection)) {
        return { ...value, selection: effect.value };
      }

      // Update range from store.
      if (effect.is(setComments)) {
        const { comments } = effect.value;
        const commentStates: CommentState[] = comments
          .map((comment) => {
            // Skip cut/deleted comments.
            if (!comment.cursor) {
              return undefined;
            }

            const range = Cursor.getRangeFromCursor(tr.state, comment.cursor);
            return range && { comment, range };
          })
          .filter(isNonNullable);

        return { ...value, comments: commentStates };
      }

      // Update entire state.
      if (effect.is(setCommentState)) {
        return effect.value;
      }
    }

    return value;
  },
});

const createCommentMark = (id: string, isCurrent: boolean) =>
  Decoration.mark({
    class: 'cm-comment',
    attributes: {
      'data-testid': 'cm-comment',
      'data-comment-id': id,
      'data-current': isCurrent ? '1' : '0',
    },
  });

/**
 * Transparent inline marks that carry `data-comment-id` for click routing (see `handleCommentClick`).
 * The colour comes from `commentsHighlightLayer`.
 */
const commentsDecorations = EditorView.decorations.compute([commentsState], (state) => {
  const {
    selection: { current },
    comments,
  } = state.field(commentsState);

  const decorations = sortBy(comments ?? [], (range) => range.range.from)
    ?.flatMap((comment) => {
      const range = comment.range;
      if (!range) {
        log.warn('Invalid range:', range);
        return undefined;
      } else if (range.from === range.to) {
        // Skip empty ranges. This can happen when a comment is cut or deleted.
        return undefined;
      }

      const mark = createCommentMark(comment.comment.id, comment.comment.id === current);
      return mark.range(range.from, range.to);
    })
    .filter(isNonNullable);

  return Decoration.set(decorations);
});

/**
 * Selection-style highlight: draws each comment range as a layer of rectangle markers below the text.
 * `RectangleMarker.forRange` extends wrapped (non-final) lines to the line edge, so a multi-line comment
 * reads as one continuous block rather than ragged per-word inline backgrounds.
 */
const commentsHighlightLayer = layer({
  above: false,
  class: 'cm-commentsHighlightLayer',
  update: (update) =>
    update.docChanged ||
    update.viewportChanged ||
    update.geometryChanged ||
    update.selectionSet ||
    update.transactions.some((tr) =>
      tr.effects.some((effect) => effect.is(setSelection) || effect.is(setComments) || effect.is(setCommentState)),
    ),
  markers: (view) => {
    const {
      selection: { current },
      comments,
    } = view.state.field(commentsState);

    return sortBy(comments ?? [], (comment) => comment.range.from).flatMap((comment) => {
      const { from, to } = comment.range;
      if (from >= to) {
        return [];
      }

      const className =
        comment.comment.id === current ? 'cm-comment-highlight cm-comment-highlight-current' : 'cm-comment-highlight';
      const range = EditorSelection.range(from, to);
      const markers = RectangleMarker.forRange(view, className, range);
      // A single rectangle means the comment fits on one visual line — round its corners.
      return markers.length === 1
        ? RectangleMarker.forRange(view, `${className} cm-comment-highlight-single`, range)
        : markers;
    });
  },
});

const handleCommentClick = EditorView.domEventHandlers({
  click: (event, view) => {
    const id = commentIdAtClick(event, view);
    if (!id) {
      return false;
    }

    // Select the clicked thread in the editor (moves the caret into the range so the proximity tracker
    // keeps it current) and notify the app for deliberate-click behaviour (e.g. reveal in a companion).
    scrollCommentIntoView(view, id);
    view.state.facet(optionsFacet).onActivate?.(id);
    return true;
  },
});

/**
 * Resolve the comment id under a click. Prefers the DOM target (a direct hit on the `.cm-comment`
 * span), then falls back to a document-position hit-test: the highlight-layer rectangles
 * (pointer-events: none) sit over the text, so a click often lands on the `.cm-line` (an ancestor of
 * the span) and an upward target walk would miss it.
 */
const commentIdAtClick = (event: MouseEvent, view: EditorView): string | undefined => {
  let target = event.target as HTMLElement | null;
  const editorRoot = view.dom;
  while (target && target !== editorRoot && !target.hasAttribute?.('data-comment-id')) {
    target = target.parentElement;
  }
  const domId = target && target !== editorRoot ? target.getAttribute('data-comment-id') : null;
  if (domId) {
    return domId;
  }

  const pos = view.posAtCoords({ x: event.clientX, y: event.clientY });
  if (pos != null) {
    const commentState = view.state.field(commentsState, false);
    return commentState?.comments.find(({ range }) => pos >= range.from && pos <= range.to)?.comment.id;
  }

  return undefined;
};

//
// Cut-and-paste.
//

type TrackedComment = { id: string; from: number; to: number };

const trackPastedComments = (onUpdate: NonNullable<CommentsOptions['onUpdate']>) => {
  // Tracks indexed selections within text.
  let tracked: { text: Text; comments: TrackedComment[] } | null = null;

  // Track cut or copy (enables cut-and-paste and copy-delete-paste to restore comment selection).
  const handleTrack = (event: Event, view: EditorView) => {
    const comments = view.state.field(commentsState);
    const { main } = view.state.selection;
    const selectedRanges = comments.comments.filter(
      ({ range }) => range.from >= main.from && range.to <= main.to && range.from < range.to,
    );

    if (!selectedRanges.length) {
      tracked = null;
    } else {
      tracked = {
        text: view.state.doc.slice(main.from, main.to),
        comments: selectedRanges.map(({ comment, range }) => ({
          id: comment.id,
          from: range.from - main.from,
          to: range.to - main.from,
        })),
      };
    }
  };

  return [
    EditorView.domEventHandlers({
      cut: handleTrack,
      copy: handleTrack,
    }),

    // Track deleted comments.
    invertedEffects.of((tr) => {
      const { comments } = tr.startState.field(commentsState);
      const effects: StateEffect<any>[] = [];
      tr.changes.iterChangedRanges((fromA, toA) => {
        for (const {
          comment: { id },
          range: { from, to },
        } of comments) {
          if (from < to && from >= fromA && to <= toA) {
            effects.push(restoreCommentEffect.of({ id, from, to }));
          }
        }
      });

      return effects;
    }),

    // Handle paste or the undo of comment deletion.
    EditorView.updateListener.of((update) => {
      const restore: TrackedComment[] = [];

      for (let i = 0; i < update.transactions.length; i++) {
        const tr = update.transactions[i];
        for (let j = 0; j < restore.length; j++) {
          restore[j] = mapTrackedComment(restore[j], tr.changes);
        }
        for (const effect of tr.effects) {
          if (effect.is(restoreCommentEffect)) {
            restore.push(effect.value);
          }
        }
      }

      if (tracked) {
        const paste = update.transactions.find((tr) => tr.isUserEvent('input.paste'));
        if (paste) {
          let found = -1;
          paste.changes.iterChanges((fromA, toA, fromB, toB, text) => {
            if (text.eq(tracked!.text)) {
              for (let i = update.transactions.indexOf(paste!) + 1; i < update.transactions.length; i++) {
                fromB = update.transactions[i].changes.mapPos(fromB);
              }

              found = fromB;
            }
          });

          if (found > -1) {
            for (const moved of tracked.comments) {
              restore.push({ id: moved.id, from: found + moved.from, to: found + moved.to });
            }
          }

          tracked = null;
        }
      }

      for (const comment of restore) {
        const { comments } = update.startState.field(commentsState);
        const exists = comments.some((c) => c.comment.id === comment.id && c.range.from < c.range.to);
        if (!exists) {
          const cursor = Cursor.getCursorFromRange(update.state, comment);
          onUpdate({ id: comment.id, cursor });
        }
      }
    }),
  ];
};

const mapTrackedComment = (comment: TrackedComment, changes: ChangeDesc) => ({
  id: comment.id,
  from: changes.mapPos(comment.from, 1),
  to: changes.mapPos(comment.to, 1),
});

/**
 * These are attached to undone/redone transactions in the editor for the purpose of restoring comments
 * that were deleted by the original changes.
 */
const restoreCommentEffect = StateEffect.define<TrackedComment>({
  map: mapTrackedComment,
});

const optionsFacet = singleValueFacet<CommentsOptions>();

/**
 * Create comment thread action.
 */
export const createComment: Command = (view) => {
  const options = view.state.facet(optionsFacet);
  // Comment creation is prohibited on this surface (e.g. a suggestion branch).
  if (options.readonly) {
    return false;
  }

  let { from, to } = view.state.selection.main;

  // Snap to the largest logical region so commenting never depends on a precise selection. Inside a
  // (unified) diff hunk the comment covers the whole changed hunk (a bare cursor in the hunk is
  // enough); otherwise, with no selection, it expands to the word under the cursor. The diff case
  // takes precedence so a comment on a change visibly covers it.
  const merge = getChunks(view.state);
  const chunk = merge?.chunks.find((candidate) => candidate.fromB <= to && candidate.endB >= from);
  if (chunk) {
    from = chunk.fromB;
    to = Math.min(chunk.endB, view.state.doc.length);
  } else if (from === to) {
    const word = view.state.wordAt(from);
    if (word) {
      from = word.from;
      to = word.to;
    }
  }

  // Exclude trailing newlines so the anchor stays within content and never reaches the document end.
  // Anchoring at the very end previously required inserting a newline to stabilise the end cursor — a
  // visible, unwanted mutation; the end cursor instead associates with the preceding character.
  while (to > from && view.state.doc.sliceString(to - 1, to) === '\n') {
    to--;
  }

  if (from === to) {
    return false;
  }

  const cursor = Cursor.getCursorFromRange(view.state, { from, to });
  if (cursor) {
    // Collapse the text selection now that the range is captured by the comment (the caret stays within
    // the new comment, keeping it the current thread).
    view.dispatch({ selection: { anchor: from } });
    // Create thread via callback; branch-tag it with the branch under review (undefined = main).
    options.onCreate?.({ cursor, from, location: view.coordsAtPos(from), branch: options.reviewBranch });
    return true;
  }

  return false;
};

//
// Utils.
//

/**
 * Whether the [from, to] document range is entirely within the editor's scroll
 * viewport. Returns false if either endpoint isn't currently rendered (off-screen).
 */
export const isRangeVisible = (view: EditorView, range: { from: number; to: number }): boolean => {
  const from = view.coordsAtPos(range.from);
  const to = view.coordsAtPos(range.to);
  if (!from || !to) {
    return false;
  }

  const { top, bottom } = view.scrollDOM.getBoundingClientRect();
  return from.top >= top && to.bottom <= bottom;
};

export type ScrollThreadOptions = {
  /** Vertical alignment when scrolling. */
  y?: 'start' | 'center' | 'end' | 'nearest';
  /** Vertical margin (px) to keep around the target when scrolling. */
  yMargin?: number;
};

/**
 * Scroll the comment thread with the given id into view (if not already entirely
 * visible) and mark it the current comment so the editor highlights it.
 */
export const scrollCommentIntoView = (
  view: EditorView,
  id: string,
  { y = 'center', yMargin }: ScrollThreadOptions = {},
) => {
  const comment = view.state.field(commentsState).comments.find((range) => range.comment.id === id);
  if (!comment?.comment.cursor) {
    return;
  }

  const range = Cursor.getRangeFromCursor(view.state, comment.comment.cursor);
  if (!range) {
    return;
  }

  const { from, to } = view.state.selection.main;
  // Never collapse a live selection: an attention change in the comments companion can land while the
  // user is selecting text, and moving the caret would silently retarget the comment they create next.
  const needsSelectionUpdate = from === to && from !== range.from;
  view.dispatch({
    selection: needsSelectionUpdate ? { anchor: range.from } : undefined,
    effects: [
      isRangeVisible(view, range) ? [] : EditorView.scrollIntoView(range.from, { y, yMargin }),
      // Always mark this thread current so the highlight follows the selected thread.
      setSelection.of({ current: id }),
    ].flat(),
  });
};

/**
 * NOTE: Matches search.
 */
// The visible highlight is drawn by `commentsHighlightLayer` (below) as a layer of rectangle markers,
// so — like a text selection — it fills wrapped lines to the line edge instead of stopping at the last
// glyph. The inline `cm-comment` mark is now transparent and exists only to carry `data-comment-id`
// (click routing) and the pointer affordance. No `border-radius`: the per-line rectangles abut, so
// rounding each would pinch the interior edges into notches on a multi-line comment.
const styles = EditorView.theme({
  // The layer only paints the background; text colour comes from the inline mark, which wraps exactly
  // the comment text. `!important` and the descendant selector win over markdown syntax-highlight tokens
  // nested inside the range.
  '.cm-comment, .cm-comment *': {
    color: 'var(--color-teal-fg) !important',
  },
  '.cm-comment[data-current="1"], .cm-comment[data-current="1"] *': {
    color: 'var(--color-orange-fg) !important',
  },
  '.cm-comment': {
    cursor: 'pointer',
  },
  // Layer wrapper: sit behind the text and never swallow clicks meant for the `cm-comment` mark.
  '.cm-commentsHighlightLayer': {
    pointerEvents: 'none',
  },
  // `box-shadow` with a 1px spread extends the block 1px on every side in the same colour — visual
  // padding without changing layout. Vertically adjacent line rectangles paint into each other's gap
  // with the same colour, so a multi-line comment stays seamless.
  // 2px horizontal breathing room around the text: `padding-inline` widens the background box and the
  // compensating negative `margin-inline` keeps it aligned to the text; `content-box` makes the padding
  // grow the box (rather than shrink the content) so the fill and border-radius extend with it.
  // `pointer-events: none` on the rectangles themselves (not only the layer wrapper) so a click over a
  // comment passes through to the `.cm-comment` mark's `data-comment-id` (and reaches `handleCommentClick`)
  // instead of landing on a highlight rectangle that carries no id.
  '.cm-comment-highlight': {
    backgroundColor: 'var(--color-teal-bg)',
    boxShadow: '0 0 0 1px var(--color-teal-bg)',
    boxSizing: 'content-box',
    marginInline: '-2px',
    paddingInline: '2px',
    pointerEvents: 'none',
  },
  '.cm-comment-highlight-current': {
    backgroundColor: 'var(--color-orange-bg)',
    boxShadow: '0 0 0 1px var(--color-orange-bg)',
  },
  // Only round a comment that occupies a single visual line: a multi-line comment is drawn as several
  // abutting rectangles, so rounding each would notch the interior edges.
  '.cm-comment-highlight-single': {
    borderRadius: '0.25rem',
  },
});