import { Slex, Token } from "@scinorandex/slex";
import { generateStates } from "./generator";
import { Stack } from "./utils/Stack";
import { Result } from "./utils/Result";
import { Production } from "./meta/common";

export type TableAction = { type: "shift" | "reduce" | "goto"; value: number };

export class TableState {
  public readonly actions: Map<string, TableAction>;

  constructor(actions = new Map<string, TableAction>()) {
    this.actions = actions;
  }

  public getTerminalAction(terminal: string): TableAction | undefined {
    return this.actions.get(`[${terminal}]`);
  }

  public getVariableAction(variable: string): TableAction | undefined {
    return this.actions.get(variable);
  }

  toJSObject() {
    return Object.fromEntries(this.actions);
  }

  static fromJSObject(json: ReturnType<TableState["toJSObject"]>) {
    return new TableState(new Map(Object.entries(json)));
  }
}

type ReducerType<TokenType, Metadata, Node> = (
  newInput: { bag: any; name: string | null },
  oldInput: { input: LR1StackSymbol<TokenType, Metadata, Node>[]; index: number },
) => Node;

export type LR1StackSymbol<TokenType, Metadata, Node> =
  | { type: "token"; token: Token<TokenType, Metadata> }
  | { type: "node"; node: Node };

export class Sparse<TokenType, Metadata, Node> {
  public constructor(
    public readonly options: {
      productions: Production[];
      states: TableState[];
      toStringifiedTokenType: (tokenType: TokenType) => string;
    },
  ) {}

  public static tryFromProductions<TokenType, Metadata, Node>(options: {
    productions: Production[];
    toStringifiedTokenType: (tokenType: TokenType) => string;
  }): Result<Sparse<TokenType, Metadata, Node>> {
    const statesResult = generateStates(options.productions);
    if (statesResult.success === false) return statesResult;
    const states = statesResult.value;

    const sparse = new Sparse<TokenType, Metadata, Node>({
      productions: options.productions,
      states: states.toStates(),
      toStringifiedTokenType: options.toStringifiedTokenType,
    });

    return { success: true, value: sparse };
  }

  public static fromProductions<TokenType, Metadata, Node>(options: {
    productions: Production[];
    toStringifiedTokenType: (tokenType: TokenType) => string;
  }) {
    const result = Sparse.tryFromProductions<TokenType, Metadata, Node>(options);
    if (result.success === false)
      throw new Error(`Encountered error "${result.reason}" at ${result.token.line}:${result.token.column}`);
    return result.value;
  }

  public generate(
    lexer: ReturnType<Slex<TokenType, Metadata>["generate"]>,
    options: {
      reducer: ReducerType<TokenType, Metadata, Node>;
      recover?: ParserRecoveryFunction<TokenType, Metadata, Node>;
    },
  ) {
    return new LR1Parser<TokenType, Metadata, Node>(this.options, options.reducer, options.recover ?? null, lexer);
  }
}

export type ParserRecoveryFunction<TokenType, Metadata, Node> = (options: {
  lexer: ReturnType<Slex<TokenType, Metadata>["generate"]>;
  statesStack: Stack<number>;
  symbolsStack: Stack<LR1StackSymbol<TokenType, Metadata, Node>>;
  isSafe: () => boolean;
  addError: (reason: string) => void;
  crash: (reason: string) => { success: false; reason: string };
  finish: (options?: { newToken: Token<TokenType, Metadata> }) => {
    success: true;
    token?: Token<TokenType, Metadata>;
  };
  states: TableState[];
}) => { success: true; token?: Token<TokenType, Metadata> } | { success: false; reason: string };

export type ParserError<TokenType, Metadata> = { token: Token<TokenType, Metadata>; message: string };
export type ParserResult<TokenType, Metadata, Node> = {
  result: Node | null;
  errors: ParserError<TokenType, Metadata>[];
};

class LR1Parser<TokenType, Metadata, Node> {
  statesStack: Stack<number> = new Stack([0]);
  symbolsStack: Stack<LR1StackSymbol<TokenType, Metadata, Node>> = new Stack();
  exceptions: ParserError<TokenType, Metadata>[] = [];

  public constructor(
    public readonly options: {
      productions: Production[];
      states: TableState[];
      toStringifiedTokenType: (tokenType: TokenType) => string;
    },
    public readonly reducer: ReducerType<TokenType, Metadata, Node>,
    public readonly recover: ParserRecoveryFunction<TokenType, Metadata, Node> | null,
    public readonly lexer: ReturnType<Slex<TokenType, Metadata>["generate"]>,
  ) {}

  public parse(): ParserResult<TokenType, Metadata, Node> {
    const { productions, states, toStringifiedTokenType } = this.options;

    while (true) {
      let currentState = this.statesStack.peek();
      let token = this.lexer.peekNextToken();
      let action = states[currentState].getTerminalAction(toStringifiedTokenType(token.type));

      if (action == null) {
        if (this.recover === null)
          throw new LR1ParserGraveError(
            `Invalid syntax. Expected ${[...states[currentState].actions.keys()].join(
              ",",
            )} but got ${toStringifiedTokenType(token.type)}`,
            token,
          );

        const recoveryResult = this.recover({
          lexer: this.lexer,
          states,
          statesStack: this.statesStack,
          symbolsStack: this.symbolsStack,
          addError: (reason: string) => this.exceptions.push({ message: reason, token: this.lexer.peekNextToken() }),
          crash: (reason: string) => ({ success: false, reason }),
          finish: (options?: { newToken: Token<TokenType, Metadata> }) => ({ success: true, token: options?.newToken }),
          isSafe: () =>
            states[this.statesStack.peek()].actions.has(`[${toStringifiedTokenType(this.lexer.peekNextToken().type)}]`),
        });

        if (recoveryResult.success === false) throw new LR1ParserGraveError(recoveryResult.reason, token);

        token = recoveryResult.token ?? this.lexer.peekNextToken();
        currentState = this.statesStack.peek();
        action = states[currentState].getTerminalAction(toStringifiedTokenType(token.type));
      }

      const fixedAction = action as TableAction;
      if (fixedAction.type === "shift") {
        // add current token to the stack and push the next state
        token = this.lexer.getNextToken();
        this.statesStack.push(fixedAction.value);
        this.symbolsStack.push({ type: "token", token: token });
      }

      // handle reductions
      else if (fixedAction.type === "reduce") {
        if (fixedAction.value === 0) break;
        const production = productions[fixedAction.value];

        // pop the stack and reduce by this production
        const popped: LR1StackSymbol<TokenType, Metadata, Node>[] = [];
        for (let i = 0; i < production.rhs.length; i++) {
          popped.unshift(this.symbolsStack.pop());
          this.statesStack.pop();
        }

        // try to perform the reduction
        const productionIndex = production.originalProductionIndex;

        try {
          // prepare the object that you're going to send to the reducer function
          const bag = {} as { [key: string]: Node | Token<TokenType, Metadata> };

          for (let i = 0; i < production.rhs.length; i++) {
            const rhsItem = production.rhs[i];
            const poppedItem = popped[i];
            if (rhsItem.name != null)
              bag[rhsItem.name] = poppedItem.type === "token" ? poppedItem.token : poppedItem.node;
          }

          const newInput = { bag, name: production.name };
          const oldInput = { input: popped, index: productionIndex };

          const node = this.reducer(newInput, oldInput);
          this.symbolsStack.push({ type: "node", node });
        } catch (err) {
          const input = JSON.stringify(popped, null, 2);
          console.log(err);
          throw new LR1ParserGraveError(
            `Error while performing reduction for production: ${productionIndex}. Reduction input: ${input}`,
            token,
          );
        }

        // get the top node and figure out what state to add to state stack
        const topNode = this.statesStack.peek();
        const gotoAction = states[topNode].getVariableAction(production.identifier);

        if (gotoAction == undefined)
          throw new LR1ParserGraveError(
            `No GOTO action for for state: ${topNode} and production: ${production.identifier}. Actions: Reduce by production ${productionIndex}`,
            token,
          );
        else if (gotoAction.type !== "goto")
          throw new LR1ParserGraveError(
            `Expected GOTO action for state: ${topNode} and production: ${production.identifier}. Actions: Reduce by production ${productionIndex}. Received: ${gotoAction.type}`,
            token,
          );

        this.statesStack.push(gotoAction.value);
      }
    }

    const top = this.symbolsStack.peek();
    return { result: top.type === "node" ? top.node : null, errors: this.exceptions };
  }
}

export class LR1ParserGraveError<TokenType, Metadata> extends Error {
  constructor(
    public readonly reason: string,
    public readonly currentToken: Token<TokenType, Metadata>,
    err?: Error,
  ) {
    super(reason);
  }
}