ward

export type MathAst =
  | { kind: "num"; value: number }
  | { kind: "var"; name: string }
  | { kind: "unary"; op: "+" | "-"; child: MathAst }
  | { kind: "bin"; op: BinOp; left: MathAst; right: MathAst }
  | { kind: "call"; name: string; args: MathAst[] };

type BinOp = "+" | "-" | "*" | "/" | "%" | "^";

export class MathParseError extends Error {
  constructor(message: string, public readonly source: string, public readonly pos: number) {
    super(`${message} at col ${pos + 1} of \`${source}\``);
  }
}

export interface MathEnv {
  vars: Record<string, number>;
}

const RE_IDENT_START = /[A-Za-z_]/;
const RE_IDENT_CONT = /[A-Za-z0-9_.]/;

type Tok =
  | { k: "num"; v: number; pos: number }
  | { k: "id"; v: string; pos: number }
  | { k: "op"; v: BinOp | "**"; pos: number }
  | { k: "lp"; pos: number }
  | { k: "rp"; pos: number }
  | { k: "comma"; pos: number };

function tokenize(src: string): Tok[] {
  const out: Tok[] = [];
  let i = 0;
  while (i < src.length) {
    const ch = src[i];
    if (ch === " " || ch === "\t" || ch === "\n" || ch === "\r") { i++; continue; }
    if (ch === "(") { out.push({ k: "lp", pos: i }); i++; continue; }
    if (ch === ")") { out.push({ k: "rp", pos: i }); i++; continue; }
    if (ch === ",") { out.push({ k: "comma", pos: i }); i++; continue; }
    if (ch === "*" && src[i + 1] === "*") { out.push({ k: "op", v: "**", pos: i }); i += 2; continue; }
    if ("+-*/%^".includes(ch)) { out.push({ k: "op", v: ch as BinOp, pos: i }); i++; continue; }
    if (ch >= "0" && ch <= "9" || ch === ".") {
      let j = i;
      while (j < src.length && /[0-9_]/.test(src[j])) j++;
      if (src[j] === ".") {
        j++;
        while (j < src.length && /[0-9_]/.test(src[j])) j++;
      }
      if (src[j] === "e" || src[j] === "E") {
        j++;
        if (src[j] === "+" || src[j] === "-") j++;
        while (j < src.length && /[0-9]/.test(src[j])) j++;
      }
      const raw = src.slice(i, j).replace(/_/g, "");
      const n = Number(raw);
      if (!Number.isFinite(n)) throw new MathParseError(`bad number "${raw}"`, src, i);
      out.push({ k: "num", v: n, pos: i });
      i = j;
      continue;
    }
    if (RE_IDENT_START.test(ch)) {
      let j = i + 1;
      while (j < src.length && RE_IDENT_CONT.test(src[j])) j++;
      out.push({ k: "id", v: src.slice(i, j), pos: i });
      i = j;
      continue;
    }
    throw new MathParseError(`unexpected character "${ch}"`, src, i);
  }
  return out;
}

class Parser {
  i = 0;
  constructor(public toks: Tok[], public src: string) {}

  parse(): MathAst {
    const ast = this.parseAdd();
    if (this.i < this.toks.length) {
      throw new MathParseError("trailing tokens", this.src, this.toks[this.i].pos);
    }
    return ast;
  }

  private parseAdd(): MathAst {
    let left = this.parseMul();
    while (this.peekOp("+", "-")) {
      const op = (this.advance() as Tok & { k: "op" }).v as BinOp;
      const right = this.parseMul();
      left = { kind: "bin", op, left, right };
    }
    return left;
  }

  private parseMul(): MathAst {
    let left = this.parseUnary();
    while (this.peekOp("*", "/", "%")) {
      const op = (this.advance() as Tok & { k: "op" }).v as BinOp;
      const right = this.parseUnary();
      left = { kind: "bin", op, left, right };
    }
    return left;
  }

  private parseUnary(): MathAst {
    if (this.peekOp("+", "-")) {
      const op = (this.advance() as Tok & { k: "op" }).v as "+" | "-";
      return { kind: "unary", op, child: this.parseUnary() };
    }
    return this.parsePower();
  }

  // Right-associative: `2 ^ 3 ^ 2` = `2 ^ (3 ^ 2)` = 512
  private parsePower(): MathAst {
    const left = this.parseAtom();
    if (this.peekOp("^", "**")) {
      this.advance();
      const right = this.parseUnary();
      return { kind: "bin", op: "^", left, right };
    }
    return left;
  }

  private parseAtom(): MathAst {
    const t = this.peek();
    if (!t) throw new MathParseError("unexpected end of expression", this.src, this.src.length);
    if (t.k === "lp") {
      this.advance();
      const inner = this.parseAdd();
      const close = this.peek();
      if (close?.k !== "rp") {
        throw new MathParseError("expected `)`", this.src, close?.pos ?? this.src.length);
      }
      this.advance();
      return inner;
    }
    if (t.k === "num") { this.advance(); return { kind: "num", value: t.v }; }
    if (t.k === "id") {
      this.advance();
      if (this.peek()?.k === "lp") {
        this.advance();
        const args: MathAst[] = [];
        if (this.peek()?.k !== "rp") {
          args.push(this.parseAdd());
          while (this.peek()?.k === "comma") {
            this.advance();
            args.push(this.parseAdd());
          }
        }
        const close = this.peek();
        if (close?.k !== "rp") {
          throw new MathParseError("expected `)` in call", this.src, close?.pos ?? this.src.length);
        }
        this.advance();
        return { kind: "call", name: t.v, args };
      }
      return { kind: "var", name: t.v };
    }
    throw new MathParseError(`unexpected token`, this.src, t.pos);
  }

  private peek(): Tok | undefined { return this.toks[this.i]; }
  private advance(): Tok { return this.toks[this.i++]; }
  private peekOp(...ops: string[]): boolean {
    const t = this.peek();
    return !!t && t.k === "op" && ops.includes(t.v);
  }
}

export function compileMath(source: string): MathAst {
  return new Parser(tokenize(source), source).parse();
}

const CONSTANTS: Record<string, number> = {
  pi: Math.PI,
  e: Math.E,
  tau: Math.PI * 2,
  inf: Infinity,
};

export function evalMath(ast: MathAst, env: MathEnv): number {
  switch (ast.kind) {
    case "num":
      return ast.value;
    case "var":
      if (Object.prototype.hasOwnProperty.call(env.vars, ast.name)) {
        return env.vars[ast.name];
      }
      if (ast.name in CONSTANTS) return CONSTANTS[ast.name];
      // JS `Math` namespace: `Math.PI`, `Math.LN2`, `Math.SQRT2`, ...
      if (ast.name.startsWith("Math.")) {
        const v = (Math as unknown as Record<string, unknown>)[
          ast.name.slice(5)
        ];
        if (typeof v === "number") return v;
      }
      throw new Error(`unknown variable: ${ast.name}`);
    case "unary": {
      const v = evalMath(ast.child, env);
      return ast.op === "-" ? -v : v;
    }
    case "bin": {
      const l = evalMath(ast.left, env);
      const r = evalMath(ast.right, env);
      switch (ast.op) {
        case "+": return l + r;
        case "-": return l - r;
        case "*": return l * r;
        case "/": return l / r;
        case "%": return l % r;
        case "^": return Math.pow(l, r);
      }
      throw new Error(`unknown op ${ast.op as string}`);
    }
    case "call": {
      const args = ast.args.map((a) => evalMath(a, env));
      // JS `Math` namespace passthrough: `Math.hypot(3,4)`, `Math.log1p(x)`,
      // `Math.random()`, anything else on the global Math object.
      if (ast.name.startsWith("Math.")) {
        const fn = (Math as unknown as Record<string, unknown>)[
          ast.name.slice(5)
        ];
        if (typeof fn === "function") {
          return (fn as (...a: number[]) => number)(...args);
        }
      }
      return callFn(ast.name, args);
    }
  }
}

function callFn(name: string, args: number[]): number {
  switch (name) {
    case "floor": return Math.floor(args[0]);
    case "ceil":  return Math.ceil(args[0]);
    case "round":
      if (args.length <= 1) return Math.round(args[0]);
      const f = Math.pow(10, args[1]);
      return Math.round(args[0] * f) / f;
    case "trunc": return Math.trunc(args[0]);
    case "abs":   return Math.abs(args[0]);
    case "sign":  return Math.sign(args[0]);
    case "sqrt":  return Math.sqrt(args[0]);
    case "cbrt":  return Math.cbrt(args[0]);
    case "exp":   return Math.exp(args[0]);
    case "log":
      return args.length === 1
        ? Math.log(args[0])
        : Math.log(args[0]) / Math.log(args[1]);
    case "log2":  return Math.log2(args[0]);
    case "log10": return Math.log10(args[0]);
    case "sin":   return Math.sin(args[0]);
    case "cos":   return Math.cos(args[0]);
    case "tan":   return Math.tan(args[0]);
    case "asin":  return Math.asin(args[0]);
    case "acos":  return Math.acos(args[0]);
    case "atan":  return Math.atan(args[0]);
    case "atan2": return Math.atan2(args[0], args[1]);
    case "pow":   return Math.pow(args[0], args[1]);
    case "min":   return Math.min(...args);
    case "max":   return Math.max(...args);
    case "clamp": return Math.max(args[1], Math.min(args[2], args[0]));
    case "mod":   return ((args[0] % args[1]) + args[1]) % args[1];
    default:
      throw new Error(`unknown function: ${name}`);
  }
}

/** Convenience: parse + evaluate in one shot. Throws on either failure. */
export function evaluate(source: string, vars: Record<string, number>): number {
  return evalMath(compileMath(source), { vars });
}