// This file defines a ProseMirror selection subclass that models
// table cell selections. The table plugin needs to be active to wire
// in the user interaction part of table selections (so that you
// actually get such selections when you select across cells).

import { Fragment, Node, ResolvedPos, Slice } from "prosemirror-model";
import {
  EditorState,
  NodeSelection,
  Selection,
  SelectionBookmark,
  SelectionRange,
  TextSelection,
  Transaction
} from "prosemirror-state";
import { Mapping } from "prosemirror-transform";
import { CellAttrs, NodeTypeName } from "../schema";
import { TableNode, TablePosOfCell, TableStart } from "../types";
import { TableMap } from "./TableMap";
import { inSameTable, pointsAtCell, shrinkColSpan } from "./util";

/**
 * A [`Selection`](http://prosemirror.net/docs/ref/#state.Selection)
 * subclass that represents a cell selection spanning part of a table.
 * With the plugin enabled, these will be created when the user
 * selects across cells, and will be drawn by giving selected cells a
 * `selectedCell` CSS class.
 */
export class CellSelection extends Selection {
  public readonly $anchorCell: Readonly<ResolvedPos>;
  public readonly $headCell: Readonly<ResolvedPos>;

  /**
   * A table selection is identified by its anchor and head cells. The positions
   * given to this constructor should point _before_ two cells in the same
   * table. They may be the same, to select a single cell.
   *
   * @param $anchorCell A resolved position pointing _in front of_ the anchor
   * cell (the one that doesn't move when extending the selection).
   * @param $headCell A resolved position pointing in front of the head cell
   * (the one moves when extending the selection).
   */
  constructor($anchorCell: ResolvedPos, $headCell: ResolvedPos = $anchorCell) {
    const table = $anchorCell.node(-1) as TableNode;
    const map = TableMap.get(table);
    const start = $anchorCell.start(-1);
    const rect = map.rectBetween(($anchorCell.pos - start) as TablePosOfCell, ($headCell.pos - start) as TablePosOfCell);
    const doc = $anchorCell.node(0);
    const cells = map.cellsInRect(rect).filter(p => p != $headCell.pos - start);
    // Make the head cell the first range, so that it counts as the
    // primary part of the selection
    cells.unshift($headCell.pos - start);
    const ranges = cells.map(pos => {
      const cell = table.nodeAt(pos);
      if (cell == null) {
        throw new Error("Unable to find table cell");
      }
      const from = pos + start + 1;
      return new SelectionRange(doc.resolve(from), doc.resolve(from + cell.content.size));
    });
    super(ranges[0].$from, ranges[0].$to, ranges);
    this.$anchorCell = $anchorCell;
    this.$headCell = $headCell;
  }

  public map(doc: Node, mapping: Mapping): Selection {
    const $anchorCell = doc.resolve(mapping.map(this.$anchorCell.pos));
    const $headCell = doc.resolve(mapping.map(this.$headCell.pos));
    if (pointsAtCell($anchorCell) && pointsAtCell($headCell) && inSameTable($anchorCell, $headCell)) {
      const tableChanged = this.$anchorCell.node(-1) != $anchorCell.node(-1);
      if (tableChanged && this.isRowSelection()) {
        return CellSelection.rowSelection($anchorCell, $headCell);
      } else if (tableChanged && this.isColSelection()) {
        return CellSelection.colSelection($anchorCell, $headCell);
      } else {
        return new CellSelection($anchorCell, $headCell);
      }
    }
    return TextSelection.between($anchorCell, $headCell);
  }

  /**
   * Returns a rectangular slice of table rows containing the selected cells.
   */
  public content(): Slice {
    const table = this.$anchorCell.node(-1) as TableNode;
    const map = TableMap.get(table);
    const start = this.$anchorCell.start(-1);
    const rect = map.rectBetween(
      (this.$anchorCell.pos - start) as TablePosOfCell,
      (this.$headCell.pos - start) as TablePosOfCell
    );
    const seen = [];
    const rows = [];

    for (let row = rect.top; row < rect.bottom; row++) {
      const rowContent: Node[] = [];
      for (let index = row * map.width + rect.left, col = rect.left; col < rect.right; col++, index++) {
        const pos = map.map[index];
        if (seen.indexOf(pos) == -1) {
          seen.push(pos);
          const cellRect = map.findCell(pos);
          let cell = table.nodeAt(pos);
          if (cell == null) {
            throw new Error("Unable to find table cell");
          }
          const attrs = cell.attrs as CellAttrs;
          const extraLeft = rect.left - cellRect.left;
          const extraRight = cellRect.right - rect.right;

          if (extraLeft > 0 || extraRight > 0) {
            let newAttrs = attrs;
            if (extraLeft > 0) {
              newAttrs = shrinkColSpan(newAttrs, 0, extraLeft);
            }
            if (extraRight > 0) {
              newAttrs = shrinkColSpan(newAttrs, newAttrs.colspan - extraRight, extraRight);
            }
            if (cellRect.left < rect.left) {
              cell = cell.type.createAndFill(newAttrs);
              if (cell == null) {
                throw new Error("Unable to create cell");
              }
            } else {
              cell = cell.type.create(newAttrs, cell.content);
            }
          }
          if (cellRect.top < rect.top || cellRect.bottom > rect.bottom) {
            const newAttrs = {
              ...attrs,
              rowspan: Math.min(cellRect.bottom, rect.bottom) - Math.max(cellRect.top, rect.top)
            };
            if (cellRect.top < rect.top) {
              cell = cell.type.createAndFill(newAttrs);
              if (cell == null) {
                throw new Error("Unable to create cell");
              }
            } else {
              cell = cell.type.create(newAttrs, cell.content);
            }
          }
          rowContent.push(cell);
        }
      }
      rows.push(table.child(row).copy(Fragment.from(rowContent)));
    }
    return new Slice(Fragment.from(rows), 1, 1);
  }

  public replace(tr: Transaction, content = Slice.empty) {
    const mapFrom = tr.steps.length;
    const ranges = this.ranges;
    for (let i = 0; i < ranges.length; i++) {
      const { $from, $to } = ranges[i];
      const mapping = tr.mapping.slice(mapFrom);
      tr.replace(mapping.map($from.pos), mapping.map($to.pos), i > 0 ? Slice.empty : content);
    }
    const sel = Selection.findFrom(tr.doc.resolve(tr.mapping.slice(mapFrom).map(this.to)), -1);
    if (sel != null) {
      tr.setSelection(sel);
    }
  }

  public replaceWith(tr: Transaction, node: Node) {
    this.replace(tr, new Slice(Fragment.from(node), 0, 0));
  }

  public forEachCell(f: (node: Node, pos: number) => void) {
    const table = this.$anchorCell.node(-1) as TableNode;
    const map = TableMap.get(table);
    const start = this.$anchorCell.start(-1);
    const cellPositions = map.cellsInRect(
      map.rectBetween((this.$anchorCell.pos - start) as TablePosOfCell, (this.$headCell.pos - start) as TablePosOfCell)
    );
    for (let i = 0; i < cellPositions.length; i++) {
      const cell = table.nodeAt(cellPositions[i]);
      if (cell == null) {
        throw new Error("Unable to find table cell");
      }
      f(cell, start + cellPositions[i]);
    }
  }

  /**
   * True if this selection goes all the way from the left to the right of the table.
   */
  public isRowSelection(): boolean {
    const map = TableMap.get(this.$anchorCell.node(-1) as TableNode);
    const start = this.$anchorCell.start(-1);
    const anchorLeft = map.columnIndex((this.$anchorCell.pos - start) as TablePosOfCell);
    const headLeft = map.columnIndex((this.$headCell.pos - start) as TablePosOfCell);
    if (Math.min(anchorLeft, headLeft) > 0) {
      return false;
    }
    const anchorRight = anchorLeft + this.$anchorCell.nodeAfter!.attrs.colspan;
    const headRight = headLeft + this.$headCell.nodeAfter!.attrs.colspan;
    return Math.max(anchorRight, headRight) == map.width;
  }

  /**
   * True if this selection goes all the way from the top to the bottom of the table.
   */
  public isColSelection(): boolean {
    const anchorTop = this.$anchorCell.index(-1);
    const headTop = this.$headCell.index(-1);
    if (Math.min(anchorTop, headTop) > 0) {
      return false;
    }
    const anchorBottom = anchorTop + this.$anchorCell.nodeAfter!.attrs.rowspan;
    const headBottom = headTop + this.$headCell.nodeAfter!.attrs.rowspan;
    return Math.max(anchorBottom, headBottom) == this.$headCell.node(-1).childCount;
  }

  /**
   * True if the entire table is selected.
   */
  public isTotalSelection(): boolean {
    return this.isColSelection() && this.isRowSelection();
  }

  public eq(other: Selection): boolean {
    return (
      other instanceof CellSelection &&
      other.$anchorCell.pos == this.$anchorCell.pos &&
      other.$headCell.pos == this.$headCell.pos
    );
  }

  public toJSON() {
    return { type: "cell", anchor: this.$anchorCell.pos, head: this.$headCell.pos };
  }

  public getBookmark() {
    return new CellBookmark(this.$anchorCell.pos, this.$headCell.pos);
  }

  /**
   * Returns the smallest column selection that covers the given anchor and head
   * cell.
   */
  public static colSelection($anchorCell: ResolvedPos, $headCell = $anchorCell): CellSelection {
    const map = TableMap.get($anchorCell.node(-1) as TableNode);
    const start = $anchorCell.start(-1);
    const anchorRect = map.findCell(($anchorCell.pos - start) as TablePosOfCell);
    const headRect = map.findCell(($headCell.pos - start) as TablePosOfCell);
    const doc = $anchorCell.node(0);

    if (anchorRect.top <= headRect.top) {
      if (anchorRect.top > 0) {
        $anchorCell = doc.resolve(start + map.map[anchorRect.left]);
      }
      if (headRect.bottom < map.height) {
        $headCell = doc.resolve(start + map.map[map.width * (map.height - 1) + headRect.right - 1]);
      }
    } else {
      if (headRect.top > 0) {
        $headCell = doc.resolve(start + map.map[headRect.left]);
      }
      if (anchorRect.bottom < map.height) {
        $anchorCell = doc.resolve(start + map.map[map.width * (map.height - 1) + anchorRect.right - 1]);
      }
    }
    return new CellSelection($anchorCell, $headCell);
  }

  /**
   * Returns the smallest row selection that covers the given anchor and head
   * cell.
   */
  public static rowSelection($anchorCell: ResolvedPos, $headCell: ResolvedPos = $anchorCell): CellSelection {
    const map = TableMap.get($anchorCell.node(-1) as TableNode);
    const start = $anchorCell.start(-1);
    const anchorRect = map.findCell(($anchorCell.pos - start) as TablePosOfCell);
    const headRect = map.findCell(($headCell.pos - start) as TablePosOfCell);
    const doc = $anchorCell.node(0);

    if (anchorRect.left <= headRect.left) {
      if (anchorRect.left > 0) {
        $anchorCell = doc.resolve(start + map.map[anchorRect.top * map.width]);
      }
      if (headRect.right < map.width) {
        $headCell = doc.resolve(start + map.map[map.width * (headRect.top + 1) - 1]);
      }
    } else {
      if (headRect.left > 0) {
        $headCell = doc.resolve(start + map.map[headRect.top * map.width]);
      }
      if (anchorRect.right < map.width) {
        $anchorCell = doc.resolve(start + map.map[map.width * (anchorRect.top + 1) - 1]);
      }
    }

    return new CellSelection($anchorCell, $headCell);
  }

  public static fromJSON(doc: Node, json: { anchor: number; head: number }) {
    return new CellSelection(doc.resolve(json.anchor), doc.resolve(json.head));
  }

  public static create(doc: Node, anchorCell: number, headCell: number = anchorCell): CellSelection {
    return new CellSelection(doc.resolve(anchorCell), doc.resolve(headCell));
  }

  public static createTotalSelection(doc: Node, tableStart: TableStart, tableMap: TableMap): CellSelection {
    return CellSelection.create(doc, tableStart + tableMap.map[0], tableStart + tableMap.map[tableMap.map.length - 1]);
  }

  public static visible = false;
}

Selection.jsonID("cell", CellSelection);

export class CellBookmark implements SelectionBookmark {
  constructor(private readonly anchor: number, private readonly head: number) {}

  public map(mapping: Mapping) {
    return new CellBookmark(mapping.map(this.anchor), mapping.map(this.head));
  }

  public resolve(doc: Node): Selection {
    const $anchorCell = doc.resolve(this.anchor);
    const $headCell = doc.resolve(this.head);

    if (
      $anchorCell.parent.type.name === ("tr" as NodeTypeName) &&
      $headCell.parent.type.name === ("tr" as NodeTypeName) &&
      $anchorCell.index() < $anchorCell.parent.childCount &&
      $headCell.index() < $headCell.parent.childCount &&
      inSameTable($anchorCell, $headCell)
    ) {
      return new CellSelection($anchorCell, $headCell);
    } else {
      return Selection.near($headCell, 1);
    }
  }
}

function isCellBoundarySelection({ $from, $to }: Selection): boolean {
  if ($from.pos == $to.pos || $from.pos < $from.pos - 6) {
    return false; // Cheap elimination
  }
  let afterFrom = $from.pos;
  let beforeTo = $to.pos;
  let depth = $from.depth;
  while (depth >= 0) {
    if ($from.after(depth + 1) < $from.end(depth)) {
      break;
    }
    depth--;
    afterFrom++;
  }
  for (let d = $to.depth; d >= 0; d--, beforeTo--) {
    if ($to.before(d + 1) > $to.start(d)) {
      break;
    }
  }

  if (afterFrom !== beforeTo) {
    return false;
  }

  const typeName = $from.node(depth).type.name;
  return (
    typeName === ("tbl" as NodeTypeName) ||
    typeName === ("td" as NodeTypeName) ||
    typeName === ("th" as NodeTypeName) ||
    typeName === ("tr" as NodeTypeName)
  );
}

export function normalizeSelection(state: EditorState, tr?: Transaction): Transaction | undefined {
  const sel = (tr !== undefined ? tr : state).selection;
  const doc = (tr !== undefined ? tr : state).doc;
  let normalize;

  if (sel instanceof NodeSelection) {
    const typeName = sel.node.type.name;
    switch (typeName) {
      case "td" as NodeTypeName:
      case "th" as NodeTypeName: {
        normalize = CellSelection.create(doc, sel.from);
        break;
      }
      case "tr" as NodeTypeName: {
        const $cell = doc.resolve(sel.from + 1);
        normalize = CellSelection.rowSelection($cell, $cell);
        break;
      }
      case "tbl" as NodeTypeName: {
        const map = TableMap.get(sel.node as TableNode);
        const start = sel.from + 1;
        const lastCell = start + map.map[map.width * map.height - 1];
        normalize = CellSelection.create(doc, start + 1, lastCell);
        break;
      }
    }
  } else if (sel instanceof TextSelection && isCellBoundarySelection(sel)) {
    normalize = TextSelection.create(doc, sel.from);
  }

  if (normalize !== undefined) {
    if (tr === undefined) {
      tr = state.tr;
    }
    tr.setSelection(normalize);
  }
  return tr;
}