import {
  parse,
  type LambdaTerm,
  isVariable,
  isApplication,
  isAbstraction,
  SymbolTable,
} from ".";

export class InvalidLambdaTermError extends Error {}

export class MaxRecursionDepthError extends Error {}

export type Visitors = Map<
  string,
  (term: DebrujinifiedLambdaTerm) => DebrujinifiedLambdaTerm
>;

export type DebrujinAbstraction = {
  abstraction: {
    param: string;
    body: DebrujinifiedLambdaTerm;
  };
};

export type DebrujinApplication = {
  application: {
    left: DebrujinifiedLambdaTerm;
    args: Array<DebrujinifiedLambdaTerm>;
  };
};

export type DebrujinIndex = { name: string; index: number };

export type DebrujinifiedLambdaTerm =
  | ((t: DebrujinifiedLambdaTerm) => DebrujinifiedLambdaTerm)
  | DebrujinAbstraction
  | DebrujinApplication
  | DebrujinIndex;

export const debrujinify = (
  term: LambdaTerm,
  symbolTable: SymbolTable,
): DebrujinifiedLambdaTerm => {
  if (isVariable(term)) {
    if (!symbolTable.has(term)) {
      throw new InvalidLambdaTermError(`Undefined variable: ${term}`);
    }
    return { index: symbolTable.get(term), name: term };
  }

  if (isAbstraction(term)) {
    const newSymbolTable = symbolTable.createChild();
    newSymbolTable.add(term.abstraction.param);

    const { body, param } = term.abstraction;
    return {
      abstraction: {
        param,
        body: debrujinify(body, newSymbolTable),
      },
    };
  }

  if (isApplication(term)) {
    const { left, args } = term.application;
    return {
      application: {
        left: debrujinify(left, symbolTable),
        args: args.map((arg) => debrujinify(arg, symbolTable)),
      },
    };
  }

  throw new InvalidLambdaTermError(
    `Invalid lambda term: ${JSON.stringify(term)}`,
  );
};

export const substitute = (
  inTerm: DebrujinifiedLambdaTerm,
  index: number,
  withTerm: DebrujinifiedLambdaTerm,
): DebrujinifiedLambdaTerm => {
  if (typeof inTerm === "function") {
    return inTerm(withTerm);
  }

  if ("index" in inTerm) {
    if (inTerm.index > index) {
      return adjustIndices(inTerm, -1);
    }
    if (index === inTerm.index) {
      return withTerm;
    }
    return inTerm;
  }

  if ("application" in inTerm) {
    const { left, args } = inTerm.application;

    return {
      application: {
        left: substitute(left, index, withTerm),
        args: args.map((arg) => substitute(arg, index, withTerm)),
      },
    };
  }

  if ("abstraction" in inTerm) {
    const { param, body } = inTerm.abstraction;
    const newBody = substitute(body, index + 1, withTerm);

    return {
      abstraction: {
        param,
        body: newBody,
      },
    };
  }

  throw new InvalidLambdaTermError(
    `Invalid lambda term: ${JSON.stringify(inTerm)}`,
  );
};

export const adjustIndices = (
  term: DebrujinifiedLambdaTerm,
  delta: number,
): DebrujinifiedLambdaTerm => {
  if ("index" in term) {
    return {
      ...term,
      index: term.index + delta,
    };
  }

  if ("application" in term) {
    const { left, args } = term.application;

    return {
      application: {
        left: adjustIndices(left, delta),
        args: args.map((arg) => adjustIndices(arg, delta)),
      },
    };
  }

  if ("abstraction" in term) {
    const { body, param } = term.abstraction;

    return {
      abstraction: {
        body: adjustIndices(body, delta),
        param,
      },
    };
  }

  throw new InvalidLambdaTermError(
    `Invalid lambda term: ${JSON.stringify(term)}`,
  );
};

export const betaReduce = (
  term: DebrujinifiedLambdaTerm,
  visitors: Visitors,
  maxDepth: number,
): DebrujinifiedLambdaTerm => {
  if (maxDepth === 0) {
    throw new MaxRecursionDepthError("max recursion depth identified");
  }
  if (typeof term === "function") {
    return term;
  }

  if ("index" in term) {
    const replacement = visitors.get(term.name)?.apply(null, [term]);
    return replacement ?? term;
  }

  if ("abstraction" in term) {
    const { body, param } = term.abstraction;

    return {
      abstraction: {
        body: betaReduce(body, visitors, maxDepth - 1),
        param,
      },
    };
  }

  if ("application" in term) {
    const { left, args } = term.application;
    const [reducedLeft, ...reducedArgs] = [left, ...args].map((term) =>
      betaReduce(term, visitors, maxDepth - 1),
    );
    return reducedArgs.reduce((acc: DebrujinifiedLambdaTerm, x) => {
      if ("abstraction" in acc) {
        const { body } = acc.abstraction;
        const substituted = substitute(body, 1, x);
        return substituted;
      }
      return acc;
    }, reducedLeft);
  }

  throw new InvalidLambdaTermError(
    `Invalid lambda term: ${JSON.stringify(term)}`,
  );
};

export const emitDebrujin = (term: DebrujinifiedLambdaTerm): string => {
  if (typeof term === "function") {
    return term.toString();
  }

  if ("index" in term) {
    return term.index.toString();
  }

  if ("abstraction" in term) {
    return `λ.${emitDebrujin(term.abstraction.body)}`;
  }

  if ("application" in term) {
    return `(${emitDebrujin(term.application.left)} ${term.application.args
      .map(emitDebrujin)
      .join(" ")})`;
  }

  throw new InvalidLambdaTermError(
    `Invalid lambda term: ${JSON.stringify(term)}`,
  );
};

export const emitNamed = (term: DebrujinifiedLambdaTerm): string => {
  if ("name" in term) {
    return term.name;
  }

  if ("abstraction" in term) {
    return `(λ (${term.abstraction.param}) . ${emitNamed(
      term.abstraction.body,
    )})`;
  }

  if ("application" in term) {
    return `(${emitNamed(term.application.left)} ${term.application.args
      .map(emitNamed)
      .join(" ")})`;
  }

  throw new InvalidLambdaTermError(`Invalid lambda term: ${term}`);
};

export const interpret = (
  term: string,
  symbolTable = new SymbolTable(),
  allowUnderscores = false, // in our world, underscores should be internal to the game.
  visitors: Visitors = new Map(),
  maxDepth = 20,
): DebrujinifiedLambdaTerm => {
  const ast = parse(term, allowUnderscores);
  const debrujined = debrujinify(ast, symbolTable);

  let prev = debrujined;
  let next = betaReduce(prev, visitors, maxDepth);

  while (emitDebrujin(prev) !== emitDebrujin(next)) {
    // alpha equivalence
    prev = next;
    next = betaReduce(prev, visitors, maxDepth);
  }
  console.log(next);
  return next;
};
