import { zeroes } from "@heydovetail/array";
import { Node } from "prosemirror-model";
import { EditorState, TextSelection, Transaction } from "prosemirror-state";
import { changedDescendants } from "../pquery";
import { CellAttrs, EditorSchema, NodeTypeName } from "../schema";
import { DocNode, DocPos, TableColumnIndex, TableNode, TableRowIndex, TableStart } from "../types";
import { CellSelection } from "./CellSelection";
import { TableMap, TableMapProblemType } from "./TableMap";
import { closestTableFromSelection, expandColSpan, findTable, isColumnHeader, shrinkColSpan } from "./util";

/**
 * Toggle the entire table being selected.
 *
 * @param tr Transaction to contribute to
 * @param tableStart Start position of the table in doc
 * @param tableMap A pre-computed table map
 */
export function toggleTableSelection(
  tr: Transaction,
  table: TableNode,
  tableStart: TableStart,
  tableMap: TableMap = TableMap.get(table)
) {
  const sel = tr.selection;

  // If the entire table is selected, clear the cell selection (replace with a
  // text selection), else select the whole table.
  if (sel instanceof CellSelection && sel.$anchorCell.start(-1) == tableStart && sel.isTotalSelection()) {
    const clearedSelection = TextSelection.create(tr.doc, sel.anchor);
    tr.setSelection(clearedSelection);
  } else {
    tr.setSelection(CellSelection.createTotalSelection(tr.doc, tableStart, tableMap));
  }
}

/**
 * Select all cells in the active table.
 */
export function selectAll(tr: Transaction): boolean {
  if (!(tr.selection instanceof CellSelection) || !tr.selection.isTotalSelection()) {
    const tableDescriptor = closestTableFromSelection(tr.selection);
    if (tableDescriptor != null) {
      tr.setSelection(CellSelection.createTotalSelection(tr.doc, tableDescriptor.start, TableMap.get(tableDescriptor.node)));
      return true;
    }
  }
  return false;
}

/**
 * Add a column at a specific index in a table.
 *
 * @param tr Transaction to contribute to
 * @param table ProseMirror node of the table
 * @param tableStart Start position of the table in doc
 * @param tableMap A pre-computed table map
 * @param columnEdgeIndex 0-indexed column edge of where to insert the column
 */
export function addColumn(
  tr: Transaction,
  table: TableNode,
  tableStart: TableStart,
  tableMap: TableMap,
  columnEdgeIndex: number
) {
  let refColumn: number | null = columnEdgeIndex > 0 ? -1 : 0;
  if (isColumnHeader(tableMap, table, columnEdgeIndex + refColumn)) {
    refColumn = columnEdgeIndex == 0 || columnEdgeIndex == tableMap.width ? null : 0;
  }

  for (let row = 0; row < tableMap.height; row++) {
    const index = row * tableMap.width + columnEdgeIndex;
    // If this position falls inside a col-spanning cell
    if (columnEdgeIndex > 0 && columnEdgeIndex < tableMap.width && tableMap.map[index - 1] == tableMap.map[index]) {
      const pos = tableMap.map[index];
      const cell = table.nodeAt(pos);
      if (cell != null) {
        tr.setNodeMarkup(
          tr.mapping.map(tableStart + pos),
          undefined,
          expandColSpan(cell.attrs as CellAttrs, columnEdgeIndex - tableMap.columnIndex(pos))
        );
        // Skip ahead if rowspan > 1
        row += cell.attrs.rowspan - 1;
      }
    } else {
      const type = refColumn === null ? table.type.schema.nodes.td : table.nodeAt(tableMap.map[index + refColumn])!.type;
      const pos = tableMap.positionAt(row, columnEdgeIndex, table);
      tr.insert(tr.mapping.map(tableStart + pos), type.createAndFill()!);
    }
  }
  return tr;
}

/**
 * Add a row at a specific index in the table.
 *
 * @param tr Transaction to contribute to.
 * @param table ProseMirror node of the table
 * @param tableStart Position of the start of the table in the document (first
 * position inside it)
 * @param tableMap Precomputed `TableMap` for the table
 * @param rowEdgeIndex 0-index of the row edge to insert to
 */
export function addRow(tr: Transaction, table: TableNode, tableStart: TableStart, tableMap: TableMap, rowEdgeIndex: number) {
  let rowPos: DocPos = tableStart;
  for (let i = 0; i < rowEdgeIndex; i++) {
    rowPos = (rowPos + table.child(i).nodeSize) as DocPos;
  }
  const cells: Array<Node> = [];
  let refRow: number | null = rowEdgeIndex > 0 ? -1 : 0;
  if (rowIsHeader(table, rowEdgeIndex + refRow, tableMap)) {
    refRow = rowEdgeIndex == 0 || rowEdgeIndex == tableMap.height ? null : 0;
  }
  for (let col = 0, index = tableMap.width * rowEdgeIndex; col < tableMap.width; col++, index++) {
    // Covered by a rowspan cell
    if (rowEdgeIndex > 0 && rowEdgeIndex < tableMap.height && tableMap.map[index] == tableMap.map[index - tableMap.width]) {
      const pos = tableMap.map[index];
      const attrs = table.nodeAt(pos)!.attrs as CellAttrs;
      tr.setNodeMarkup(tableStart + pos, undefined, { ...attrs, rowspan: attrs.rowspan + 1 });
      col += attrs.colspan - 1;
    } else {
      const type =
        refRow === null ? table.type.schema.nodes.td : table.nodeAt(tableMap.map[index + refRow * tableMap.width])!.type;
      cells.push(type.createAndFill()!);
    }
  }
  tr.insert(rowPos, table.type.schema.nodes.tr.create(undefined, cells));
  return tr;
}

/**
 * Delete a table.
 *
 * @param tr Transaction to contribute to
 * @param tableStart Position of the start of the table in the document (first
 * position inside it)
 */
export function deleteTable(tr: Transaction, tableStart: TableStart): boolean {
  const table = findTable(tr.doc as DocNode, tableStart);
  if (table !== null) {
    const beforeTable = tableStart - 1;
    const afterTable = beforeTable + table.nodeSize;
    tr.delete(beforeTable, afterTable);
    return true;
  }
  return false;
}

/**
 * Remove a range of columns from a table.
 *
 * Returns `true` if it was able to successfully delete the columns. If the
 * table would end up empty (i.e. no columns), the entire table is deleted.
 *
 * @param tr Transaction to contribute to
 * @param tableStart Position of the start of the table in the document (first
 * position inside it)
 * @param columnStart 0-index of the first column to remove
 * @param columnEnd 0-index of the last column to remove
 */
export function removeColumns(
  tr: Transaction,
  tableStart: TableStart,
  columnStart: TableColumnIndex,
  columnEnd = columnStart
): boolean {
  const table = findTable(tr.doc as DocNode, tableStart);
  if (table !== null && columnStart === 0 && columnEnd === TableMap.get(table).width - 1) {
    return deleteTable(tr, tableStart);
  }

  for (let columnIndex = columnEnd; columnIndex >= columnStart; columnIndex--) {
    const table = findTable(tr.doc as DocNode, tableStart);
    if (table === null) {
      return false;
    }

    const tableMap = TableMap.get(table);
    const mapStart = tr.mapping.maps.length;

    for (let row = 0; row < tableMap.height; ) {
      const cellIndex = row * tableMap.width + columnIndex;
      const cellPos = tableMap.map[cellIndex];
      const cell = table.nodeAt(cellPos);
      if (cell != null) {
        // If this is part of a col-spanning cell
        if (
          (cellIndex > 0 && tableMap.map[cellIndex - 1] == cellPos) ||
          (cellIndex < tableMap.width - 1 && tableMap.map[cellIndex + 1] == cellPos)
        ) {
          tr.setNodeMarkup(
            tr.mapping.slice(mapStart).map(tableStart + cellPos),
            undefined,
            shrinkColSpan(cell.attrs as CellAttrs, columnIndex - tableMap.columnIndex(cellPos))
          );
        } else {
          const start = tr.mapping.slice(mapStart).map(tableStart + cellPos);
          tr.delete(start, start + cell.nodeSize);
        }
        row += cell.attrs.rowspan;
      }
    }
  }
  return true;
}

/**
 * Remove a row from a table.
 *
 * @param tr Transaction to contribute to
 * @param tableStart Position of the start of the table in the document (first position inside it)
 * @param rowStart 0-index of the first row to remove
 * @param rowEnd 0-index of the last row to remove
 */
export function removeRows(tr: Transaction, tableStart: TableStart, rowStart: TableRowIndex, rowEnd = rowStart): boolean {
  const table = findTable(tr.doc as DocNode, tableStart);
  if (table !== null && rowStart === 0 && rowEnd === TableMap.get(table).height - 1) {
    return deleteTable(tr, tableStart);
  }

  for (let rowIndex = rowEnd; rowIndex >= rowStart; rowIndex--) {
    const table = findTable(tr.doc as DocNode, tableStart);
    if (table === null) {
      return false;
    }
    const tableMap = TableMap.get(table);

    let rowPos = 0;
    for (let i = 0; i < rowIndex; i++) {
      rowPos += table.child(i).nodeSize;
    }
    const nextRow = rowPos + table.child(rowIndex).nodeSize;

    const mapFrom = tr.mapping.maps.length;
    tr.delete(tableStart + rowPos, tableStart + nextRow);

    for (let col = 0, index = rowIndex * tableMap.width; col < tableMap.width; col++, index++) {
      const cellPos = tableMap.map[index];
      if (rowIndex > 0 && cellPos == tableMap.map[index - tableMap.width]) {
        // If this cell starts in the row above, simply reduce its rowspan
        const attrs = table.nodeAt(cellPos)!.attrs as CellAttrs;
        tr.setNodeMarkup(tr.mapping.slice(mapFrom).map(tableStart + cellPos), undefined, {
          ...attrs,
          rowspan: attrs.rowspan - 1
        });
        col += attrs.colspan - 1;
      } else if (rowIndex < tableMap.width && cellPos == tableMap.map[index + tableMap.width]) {
        // Else, if it continues in the row below, it has to be moved down
        const cell = table.nodeAt(cellPos)!;
        const attrs = cell.attrs as CellAttrs;
        const copy = cell.type.create({ ...attrs, rowspan: attrs.rowspan - 1 }, cell.content);
        const newPos = tableMap.positionAt(rowIndex + 1, col, table);
        tr.insert(tr.mapping.slice(mapFrom).map(tableStart + newPos), copy);
        col += cell.attrs.colspan - 1;
      }
    }
  }
  return true;
}

/**
 * Returns true if the row is a header row (all cells are header cells).
 *
 * @param table ProseMirror node for the table
 * @param rowIndex 0-index of the row to check
 * @param tableMap TableMap for table
 */
export function rowIsHeader(table: Node<EditorSchema>, rowIndex: number, tableMap: TableMap): boolean {
  const headerCell = table.type.schema.nodes.th;
  for (let col = 0; col < tableMap.width; col++) {
    if (table.nodeAt(tableMap.map[col + rowIndex * tableMap.width])!.type != headerCell) {
      return false;
    }
  }
  return true;
}

/**
 * Inspect all tables in the given state's document and return a
 * transaction that fixes them, if necessary. If `oldState` was
 * provided, that is assumed to hold a previous, known-good state,
 * which will be used to avoid re-scanning unchanged parts of the
 * document.
 */
export function fixTables(state: EditorState, oldState?: EditorState): Transaction | undefined {
  let tr: Transaction | undefined;

  function check(node: Node<EditorSchema>, pos: number) {
    if (node.type.name === ("tbl" as NodeTypeName)) {
      tr = fixTable(state, node as TableNode, (pos + 1) as TableStart, tr);
    }
  }

  if (oldState === undefined) {
    state.doc.descendants(check);
  } else if (oldState.doc != state.doc) {
    changedDescendants(oldState.doc, state.doc, 0, check);
  }

  return tr;
}

/**
 * Fix the given table, if necessary. Will append to the transaction it was
 * given, if non-null, or create a new one if necessary.
 */
export function fixTable(
  state: EditorState<EditorSchema>,
  table: TableNode,
  tableStart: TableStart,
  tr?: Transaction
): Transaction | undefined {
  const map = TableMap.get(table);
  if (map.problems === null) {
    return tr;
  }

  // Track which rows we must add cells to, so that we can adjust that
  // when fixing collisions.
  const missingCellCountForRow = zeroes(map.height);

  for (const prob of map.problems) {
    if (prob.type == TableMapProblemType.COLLISION) {
      const cell = table.nodeAt(prob.pos);
      if (cell != null) {
        const attrs = cell.attrs as CellAttrs;
        for (let j = 0; j < attrs.rowspan; j++) {
          missingCellCountForRow[prob.row + j] += prob.n;
        }
        if (tr === undefined) {
          tr = state.tr;
        }
        tr.setNodeMarkup(
          tr.mapping.map(tableStart + prob.pos),
          undefined,
          shrinkColSpan(attrs, attrs.colspan - prob.n, prob.n)
        );
      }
    } else if (prob.type == TableMapProblemType.MISSING) {
      missingCellCountForRow[prob.row] += prob.n;
    } else if (prob.type == TableMapProblemType.OVERLONG_ROWSPAN) {
      const cell = table.nodeAt(prob.pos);
      if (cell != null) {
        const attrs = cell.attrs as CellAttrs;
        if (tr === undefined) {
          tr = state.tr;
        }
        tr.setNodeMarkup(tr.mapping.map(tableStart + prob.pos), undefined, { ...attrs, rowspan: attrs.rowspan - prob.n });
      }
    } else if (prob.type == TableMapProblemType.COLWIDTH_MISMATCH) {
      const cell = table.nodeAt(prob.pos);
      if (cell != null) {
        const attrs = cell.attrs as CellAttrs;
        if (tr === undefined) {
          tr = state.tr;
        }
        tr.setNodeMarkup(tr.mapping.map(tableStart + prob.pos), undefined, { ...attrs, colwidth: prob.colwidth });
      }
    }
  }

  let firstRowMissingCellsIndex, lastRowMissingCellsIndex;
  for (let i = 0; i < missingCellCountForRow.length; i++) {
    if (missingCellCountForRow[i] > 0) {
      if (firstRowMissingCellsIndex === undefined) {
        firstRowMissingCellsIndex = i;
      }
      lastRowMissingCellsIndex = i;
    }
  }

  // Add the necessary cells, using a heuristic for whether to add the
  // cells at the start or end of the rows (if it looks like a 'bite'
  // was taken out of the table, add cells at the start of the row
  // after the bite. Otherwise add them at the end).
  for (let rowIndex = 0, pos = tableStart as DocPos; rowIndex < map.height; rowIndex++) {
    const afterRowPos = (pos + table.child(rowIndex).nodeSize) as DocPos;
    const addCellCount = missingCellCountForRow[rowIndex];
    if (addCellCount > 0) {
      const nodes = [];
      for (let j = 0; j < addCellCount; j++) {
        const cell = state.schema.nodes.td.createAndFill()!;
        nodes.push(cell);
      }
      const side =
        (rowIndex == 0 || firstRowMissingCellsIndex == rowIndex - 1) && lastRowMissingCellsIndex == rowIndex
          ? pos + 1
          : afterRowPos - 1;
      if (tr === undefined) {
        tr = state.tr;
      }
      tr.insert(tr.mapping.map(side), nodes);
    }
    pos = afterRowPos;
  }
  return tr;
}