import { Cursor, RawCursor } from "./cursor";
import { EditorState } from "./editorState";
import { EmptyFlow } from "./flows/empty";
import { StrandPath, strandPathsEqual, TokenPath } from "./path";
import { SelectionRange } from "./selectionRange";
import { Strand } from "./strand";
import { Token } from "./token";
import { h, t, VNode } from "./vdom";

interface TokenSplice {
  strandPath: StrandPath;
  fromPos: number;
  toPos: number;
  newTokens: readonly Token[];
}

export class Transaction {
  startState: EditorState;
  changes: TokenSplice[];
  newSelection: SelectionRange | null;

  constructor(
    startState: EditorState,
    changes: TokenSplice[],
    newSelection: SelectionRange | Cursor | null = null
  ) {
    this.startState = startState;
    this.changes = changes;
    this.newSelection =
      newSelection instanceof Cursor
        ? new SelectionRange(newSelection)
        : newSelection;
  }

  get newState(): EditorState {
    const getMatchingChanges = (strandPath: StrandPath) => {
      return this.changes.filter((change) =>
        strandPathsEqual(change.strandPath, strandPath)
      );
    };

    // I want to apply changes depth-first, so I use a recursive function.
    // The goal is to map over each strand, depth-first, returning a new strand with the changes applied.
    const applyChangesToStrand = (
      strand: Strand,
      currentPath: StrandPath
    ): Strand => {
      let newTokens: Token[] = [];
      let tokenIndexOffset = 0;

      const matchingChanges = getMatchingChanges(currentPath);

      let lastChangeEnd = 0;

      for (const change of matchingChanges) {
        // Add tokens before the change
        newTokens.push(
          ...strand.tokens
            .slice(lastChangeEnd, change.fromPos)
            .map((token) => token.clone())
        );

        // Add new tokens from the change
        newTokens.push(...change.newTokens);

        lastChangeEnd = change.toPos;
        tokenIndexOffset +=
          change.newTokens.length - (change.toPos - change.fromPos);
      }

      // Add remaining tokens after the last change
      newTokens.push(
        ...strand.tokens.slice(lastChangeEnd).map((token) => token.clone())
      );

      // Now process children strands
      newTokens = newTokens.map((token, tokenIndex) => {
        return token.mapChildren((childStrand, childIndex) => {
          return applyChangesToStrand(childStrand, [
            ...currentPath,
            { tokenIndex, childIndex },
          ]);
        });
      });

      return new Strand(newTokens);
    };

    const newContent = applyChangesToStrand(this.startState.content, []);

    // Update anchor and head positions
    const updateCursorPosition = (rawCursor: RawCursor): RawCursor => {
      let newRawCursor = { ...rawCursor };

      for (const change of this.changes) {
        if (strandPathsEqual(rawCursor.strandPath, change.strandPath)) {
          if (rawCursor.pos < change.fromPos) {
            // Cursor is before the change
            continue;
          } else if (rawCursor.pos >= change.toPos) {
            // Cursor is after the change
            newRawCursor.pos +=
              change.newTokens.length - (change.toPos - change.fromPos);
          } else {
            // Cursor is within the changed range
            newRawCursor.pos = change.fromPos + change.newTokens.length;
          }
        } else if (
          strandPathsEqual(
            rawCursor.strandPath.slice(0, change.strandPath.length),
            change.strandPath
          )
        ) {
          // Cursor is a child of a token in the changed strand
          const relativeLevel = rawCursor.strandPath[change.strandPath.length];
          if (!relativeLevel) continue; // Should not happen

          if (relativeLevel.tokenIndex < change.fromPos) {
            // Cursor's token is before the change
            continue;
          } else if (relativeLevel.tokenIndex >= change.toPos) {
            // Cursor's token is after the change
            newRawCursor.strandPath = [
              ...newRawCursor.strandPath.slice(0, change.strandPath.length),
              {
                tokenIndex:
                  relativeLevel.tokenIndex +
                  (change.newTokens.length - (change.toPos - change.fromPos)),
                childIndex: relativeLevel.childIndex,
              },
              ...newRawCursor.strandPath.slice(change.strandPath.length + 1),
            ];
          } else {
            // Cursor's token is within the changed range
            // Move cursor to the end of the newly inserted tokens
            newRawCursor.strandPath = [
              ...newRawCursor.strandPath.slice(0, change.strandPath.length),
            ];
            newRawCursor.pos = change.fromPos + change.newTokens.length;
            break;
          }
        }
      }

      return newRawCursor;
    };

    let newSelection: SelectionRange;
    if (this.newSelection) {
      newSelection = this.newSelection;
    } else {
      let anchor = this.startState.selection.anchor.toRaw();
      let head = this.startState.selection.head.toRaw();
      anchor = updateCursorPosition(anchor);
      head = updateCursorPosition(head);

      newSelection = new SelectionRange(
        new Cursor(anchor.strandPath, anchor.pos),
        new Cursor(head.strandPath, head.pos)
      );
    }

    return new EditorState(newContent, newSelection);
  }

  setNewSelection(newSelection: SelectionRange | Cursor | null): Transaction {
    if (newSelection instanceof Cursor) {
      newSelection = new SelectionRange(newSelection);
    }
    return new Transaction(this.startState, this.changes, newSelection);
  }

  static void(startState: EditorState): Transaction {
    return new Transaction(startState, []);
  }

  // TODO: Improve & test this logic
  static insertAtSelection(
    startState: EditorState,
    newTokens:
      | Token[]
      | ((params: {
          setCursor: (
            type?: "anchor" | "head" | "collapsed"
          ) => TransactionNewCursor;
        }) => Token[]),
    selectionRange: SelectionRange = startState.selection
  ): Transaction {
    const {
      tokens,
      anchor: newAnchor,
      head: newHead,
    } = extractNewCursors(
      newTokens instanceof Function
        ? newTokens({ setCursor: (type) => new TransactionNewCursor(type) })
        : newTokens
    );

    if (selectionRange.isCollapsed()) {
      return new Transaction(
        startState,
        [
          {
            strandPath: selectionRange.anchor.strandPath,
            fromPos: selectionRange.anchor.pos,
            toPos: selectionRange.anchor.pos,
            newTokens: tokens,
          },
        ],
        newAnchor || newHead
          ? new SelectionRange(
              newAnchor
                ? new Cursor(newAnchor.strandPath, newAnchor.pos)
                : undefined,
              newHead ? new Cursor(newHead.strandPath, newHead.pos) : undefined
            )
          : null
      );
    } else {
      // If the selection is not collapsed, we need to delete the selected range first
      const anchor = selectionRange.anchor;
      const head = selectionRange.head;

      // Ensure anchor is before head
      let fromCursor: Cursor;
      let toCursor: Cursor;
      if (
        strandPathsEqual(anchor.strandPath, head.strandPath) &&
        anchor.pos <= head.pos
      ) {
        fromCursor = anchor;
        toCursor = head;
      } else {
        fromCursor = head;
        toCursor = anchor;
      }

      return new Transaction(
        startState,
        [
          {
            strandPath: fromCursor.strandPath,
            fromPos: fromCursor.pos,
            toPos: toCursor.pos,
            newTokens: tokens,
          },
        ],
        newAnchor || newHead
          ? new SelectionRange(
              newAnchor
                ? new Cursor(newAnchor.strandPath, newAnchor.pos)
                : undefined,
              newHead ? new Cursor(newHead.strandPath, newHead.pos) : undefined
            )
          : null
      );
    }
  }

  // TODO: Improve & test this logic
  static deleteAtSelection(
    startState: EditorState,
    bias: "forward" | "backward" = "backward",
    selectionRange: SelectionRange = startState.selection
  ): Transaction {
    const changes: TokenSplice[] = [];

    if (selectionRange.isCollapsed()) {
      const cursor = selectionRange.anchor;
      if (bias === "backward") {
        if (cursor.pos === 0) {
          if (cursor.strandPath.length === 0) {
            // At the very start of the document; nothing to delete
            return Transaction.void(startState);
          } else {
            // At the start of a sub-strand; delete the parent token
            // and flatten its children
            return Transaction.deleteToken(
              startState,
              {
                strandPath: cursor.strandPath.slice(0, -1),
                tokenIndex: cursor.strandPath.at(-1)!.tokenIndex,
              },
              "flatten",
              cursor
            );
          }
        } else {
          changes.push({
            strandPath: cursor.strandPath,
            fromPos: cursor.pos - 1,
            toPos: cursor.pos,
            newTokens: [],
          });
        }
      } else if (bias === "forward") {
        if (
          cursor.pos ===
          startState.content.getStrand(cursor.strandPath)!.tokens.length
        ) {
          if (cursor.strandPath.length === 0) {
            // At the very end of the document; nothing to delete
            return Transaction.void(startState);
          } else {
            // At the end of a sub-strand; delete the parent token
            // and flatten its children
            return Transaction.deleteToken(
              startState,
              {
                strandPath: cursor.strandPath.slice(0, -1),
                tokenIndex: cursor.strandPath.at(-1)!.tokenIndex,
              },
              "flatten",
              cursor
            );
          }
        } else {
          changes.push({
            strandPath: cursor.strandPath,
            fromPos: cursor.pos,
            toPos: cursor.pos + 1,
            newTokens: [],
          });
        }
      }
    } else {
      changes.push({
        strandPath: selectionRange.commonStrandPath,
        fromPos: selectionRange.commonStartPos,
        toPos: selectionRange.commonEndPos,
        newTokens: [],
      });
    }

    return new Transaction(startState, changes);
  }

  static deleteToken(
    startState: EditorState,
    tokenPath: TokenPath,
    children: "delete" | "flatten" = "delete",
    currentCursor: Cursor | null = null
  ): Transaction {
    const changes: TokenSplice[] = [];

    const strand = startState.content.getStrand(tokenPath.strandPath);
    if (!strand) {
      throw new Error("Invalid strand path");
    }
    const token = strand.getToken(tokenPath);
    if (!token) {
      throw new Error("Invalid token path");
    }

    if (children === "delete") {
      changes.push({
        strandPath: tokenPath.strandPath,
        fromPos: tokenPath.tokenIndex,
        toPos: tokenPath.tokenIndex + 1,
        newTokens: [],
      });
      return new Transaction(startState, changes);
    } else if (children === "flatten") {
      changes.push({
        strandPath: tokenPath.strandPath,
        fromPos: tokenPath.tokenIndex,
        toPos: tokenPath.tokenIndex + 1,
        newTokens: token.children
          .flatMap((strand) => strand.tokens)
          .map((t) => t.clone()),
      });

      let newCursor: Cursor | null = null;
      if (currentCursor) {
        const childIndex =
          currentCursor.strandPath[tokenPath.strandPath.length].childIndex;
        const posInChildStrand = currentCursor.pos;
        const tokensBefore = token.children
          .slice(0, childIndex)
          .flatMap((strand) => strand.tokens).length;
        newCursor = new Cursor(
          tokenPath.strandPath,
          tokenPath.tokenIndex + tokensBefore + posInChildStrand
        );
      }

      return new Transaction(startState, changes, newCursor);
    }

    children satisfies never;
    throw new Error("Invalid children option");
  }

  static setSelection(
    startState: EditorState,
    newSelection: SelectionRange | Cursor
  ): Transaction {
    if (newSelection instanceof Cursor) {
      newSelection = new SelectionRange(newSelection);
    }
    return new Transaction(startState, [], newSelection);
  }

  static selectAll(startState: EditorState): Transaction {
    return Transaction.setSelection(
      startState,
      new SelectionRange(
        new Cursor([], 0),
        new Cursor([], startState.content.tokens.length)
      )
    );
  }

  static moveCursor(
    startState: EditorState,
    direction: "left" | "right",
    preserveAnchor: boolean = false
  ): Transaction {
    const currentCursor = startState.selection.head;

    // Find current strand based on cursor
    let currentStrand = startState.content;
    for (const level of currentCursor.strandPath) {
      const token = currentStrand.tokens[level.tokenIndex];
      currentStrand = token.children[level.childIndex];
    }

    if (direction === "right") {
      // If the selection is not collapsed and we shouldn't preserve the anchor,
      // move to the end of the selection
      if (!startState.selection.isCollapsed() && !preserveAnchor) {
        const newCursor = new Cursor(
          startState.selection.commonStrandPath,
          startState.selection.commonEndPos
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      // If the next token has children, move into its first child strand
      const nextToken = currentStrand.tokens[currentCursor.pos];
      if (nextToken && nextToken.children.length > 0) {
        const newCursor = new Cursor(
          [
            ...currentCursor.strandPath,
            { tokenIndex: currentCursor.pos, childIndex: 0 },
          ],
          0
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      // If we are within the strand, move right past the next token
      if (currentCursor.pos < currentStrand.tokens.length) {
        const newCursor = new Cursor(
          currentCursor.strandPath,
          currentCursor.pos + 1
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      // We are at the end of the current strand. Try to move into a sibling strand or up to the parent.
      let parentStrand = startState.content;
      for (const level of currentCursor.strandPath.slice(0, -1)) {
        const token = parentStrand.tokens[level.tokenIndex];
        parentStrand = token.children[level.childIndex];
      }

      const lastStrandPathPart = currentCursor.strandPath.at(-1);

      if (!lastStrandPathPart) {
        // Already at the top-level strand; cannot move right
        const newCursor = currentCursor.clone();
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      const { tokenIndex, childIndex } = lastStrandPathPart;
      const parentToken = parentStrand.tokens[tokenIndex];

      if (childIndex + 1 < parentToken.children.length) {
        // Move to the next sibling child strand
        const newCursor = new Cursor(
          [
            ...currentCursor.strandPath.slice(0, -1),
            { tokenIndex, childIndex: childIndex + 1 },
          ],
          0
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      } else {
        // Move up to the parent strand
        const newCursor = new Cursor(
          currentCursor.strandPath.slice(0, -1),
          tokenIndex + 1
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }
    } else {
      // If the selection is not collapsed and we shouldn't preserve the anchor,
      // move to the start of the selection
      if (!startState.selection.isCollapsed() && !preserveAnchor) {
        const newCursor = new Cursor(
          startState.selection.commonStrandPath,
          startState.selection.commonStartPos
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      // If the previous token has children, move into its last child strand
      if (currentCursor.pos > 0) {
        const prevToken = currentStrand.tokens[currentCursor.pos - 1];
        if (prevToken && prevToken.children.length > 0) {
          const newCursor = new Cursor(
            [
              ...currentCursor.strandPath,
              {
                tokenIndex: currentCursor.pos - 1,
                childIndex: prevToken.children.length - 1,
              },
            ],
            prevToken.children[prevToken.children.length - 1].tokens.length
          );
          return Transaction.setSelection(
            startState,
            preserveAnchor
              ? new SelectionRange(startState.selection.anchor, newCursor)
              : newCursor
          );
        }
      }

      // If we are within the strand, move left before the previous token
      if (currentCursor.pos > 0) {
        const newCursor = new Cursor(
          currentCursor.strandPath,
          currentCursor.pos - 1
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      // We are at the start of the current strand. Try to move into a sibling strand or up to the parent.
      let parentStrand = startState.content;
      for (const level of currentCursor.strandPath.slice(0, -1)) {
        const token = parentStrand.tokens[level.tokenIndex];
        parentStrand = token.children[level.childIndex];
      }

      const lastStrandPathPart = currentCursor.strandPath.at(-1);
      if (!lastStrandPathPart) {
        // Already at the top-level strand; cannot move left
        const newCursor = currentCursor.clone();
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }

      const { tokenIndex, childIndex } = lastStrandPathPart;
      const parentToken = parentStrand.tokens[tokenIndex];

      if (childIndex > 0) {
        // Move to the previous sibling child strand
        const siblingStrand = parentToken.children[childIndex - 1];
        const newCursor = new Cursor(
          [
            ...currentCursor.strandPath.slice(0, -1),
            { tokenIndex, childIndex: childIndex - 1 },
          ],
          siblingStrand.tokens.length
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      } else {
        // Move up to the parent strand
        const newCursor = new Cursor(
          currentCursor.strandPath.slice(0, -1),
          tokenIndex
        );
        return Transaction.setSelection(
          startState,
          preserveAnchor
            ? new SelectionRange(startState.selection.anchor, newCursor)
            : newCursor
        );
      }
    }
  }
}

class TransactionNewCursor extends Token {
  flow = new EmptyFlow();

  constructor(public type: "anchor" | "head" | "collapsed" = "collapsed") {
    super();
  }

  get children(): readonly Strand[] {
    return [];
  }

  mapChildren(
    _fn: (child: Strand, index: number) => Strand
  ): TransactionNewCursor {
    return new TransactionNewCursor();
  }

  renderToDebugText(): string {
    return `TransactionNewCursor()`;
  }

  renderToDebugHTML(): VNode {
    return h("span", {}, t("▮"));
  }

  renderToDebugMathML(): VNode {
    return h("mtext", {}, t("▮"));
  }
}

function extractNewCursors(tokens: readonly Token[]): {
  tokens: readonly Token[];
  anchor: RawCursor | null;
  head: RawCursor | null;
} {
  let newCursorTokenPaths: TokenPath[] = [];
  let anchor: RawCursor | null = null;
  let head: RawCursor | null = null;

  const strand = new Strand(tokens);
  for (const [str, strandPath] of strand.traverseStrands()) {
    const tokens = [...str.tokens];
    for (let i = 0; i < tokens.length; i++) {
      const token = tokens[i];
      if (token instanceof TransactionNewCursor) {
        newCursorTokenPaths.push({ strandPath, tokenIndex: i });
        if (token.type === "anchor" || token.type === "collapsed") {
          if (anchor) {
            throw new Error(
              "Multiple new `anchor` or `collapsed` cursors found in transaction."
            );
          }
          anchor = { strandPath, pos: i };
        }
        if (token.type === "head" || token.type === "collapsed") {
          if (head) {
            throw new Error(
              "Multiple new `head` or `collapsed` cursors found in transaction."
            );
          }
          head = { strandPath, pos: i };
        }
      }
    }
  }

  function filterOutNewCursors(strand: Strand): Strand {
    return new Strand(
      strand.tokens
        .filter((token) => !(token instanceof TransactionNewCursor))
        .map((token) =>
          token.mapChildren((childStrand) => filterOutNewCursors(childStrand))
        )
    );
  }

  const newTokens = filterOutNewCursors(strand).tokens;

  // Adjust cursor positions after removing new cursor tokens
  function adjustCursorPosition(rawCursor: RawCursor): RawCursor {
    let adjustedCursor = { ...rawCursor };

    for (const tokenPath of newCursorTokenPaths) {
      tokenPath.strandPath;
      rawCursor.strandPath;

      // Example: 1X+a|/b
      // tokenPath: { strandPath: [], tokenIndex: 1 }
      // adjustedCursor: { strandPath: [{ tokenIndex: 3, childIndex: 0 }], pos: 1 }
      // Result (new adjustedCursor): { strandPath: [{ tokenIndex: 2, childIndex: 0 }], pos: 1 }
      const isSameStrandPath = strandPathsEqual(
        tokenPath.strandPath,
        adjustedCursor.strandPath
      );

      if (isSameStrandPath) {
        if (tokenPath.tokenIndex < adjustedCursor.pos) {
          adjustedCursor.pos -= 1;
        }
      } else if (
        strandPathsEqual(
          tokenPath.strandPath.slice(0, adjustedCursor.strandPath.length),
          adjustedCursor.strandPath
        )
      ) {
        // Cursor is a child of a token in the changed strand
        const relativeLevel =
          adjustedCursor.strandPath[tokenPath.strandPath.length];
        if (!relativeLevel) continue; // Should not happen

        if (tokenPath.tokenIndex < relativeLevel.tokenIndex) {
          // Cursor's token is after the removed cursor token
          adjustedCursor.strandPath = [
            ...adjustedCursor.strandPath.slice(0, tokenPath.strandPath.length),
            {
              tokenIndex: relativeLevel.tokenIndex - 1,
              childIndex: relativeLevel.childIndex,
            },
            ...adjustedCursor.strandPath.slice(tokenPath.strandPath.length + 1),
          ];
        }
      }
    }

    return adjustedCursor;
  }

  return {
    tokens: newTokens,
    anchor: anchor ? adjustCursorPosition(anchor) : null,
    head: head ? adjustCursorPosition(head) : null,
  };
}