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

export class InvalidLambdaTermError extends Error {}

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 =
  | 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 ("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
): DebrujinifiedLambdaTerm => {
  if ("index" in term) {
    return term;
  }

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

    return {
      abstraction: {
        body: betaReduce(body),
        param,
      },
    };
  }

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

    return args.reduce((acc: DebrujinifiedLambdaTerm, x) => {
      if ("abstraction" in acc) {
        const { body } = acc.abstraction;
        const newBody = substitute(body, 1, x);
        return newBody;
      }
      if ("application" in acc) {
        const {
          application: { left, args },
        } = acc;
        return {
          application: {
            left,
            args: [...args, x],
          },
        };
      }
      return { application: { left: acc, args: [x] } };
    }, left);
  }

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

export const emitDebrujin = (term: DebrujinifiedLambdaTerm): string => {
  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
): DebrujinifiedLambdaTerm => {
  const ast = parse(term, allowUnderscores);
  const debrujined = debrujinify(ast, symbolTable);

  let prev = debrujined;
  let next = betaReduce(prev);

  while (emitDebrujin(prev) !== emitDebrujin(next)) {
    // alpha equivalence
    prev = next;
    next = betaReduce(prev);
  }
  return next;
};