wellformed

import { FormulaError, type BinaryOp, type Node } from './ast.js';
import { lex, type Token } from './lex.js';

const MAX_SOURCE_LENGTH = 2000;
const MAX_DEPTH = 32;

/**
 * Recursive-descent parser for the Wellformed formula language.
 *
 * Precedence, loosest first:
 *   OR · AND · comparison · & (concat) · + - · * / · unary - NOT · primary
 */
export function parse(source: string): Node {
  if (source.length > MAX_SOURCE_LENGTH) {
    throw new FormulaError(`Expression is too long (max ${MAX_SOURCE_LENGTH} characters)`);
  }
  const tokens = lex(source);
  const p = new Parser(tokens);
  const node = p.expression(0);
  p.expectEof();
  return node;
}

const parseCache = new Map<string, Node | FormulaError>();
const PARSE_CACHE_LIMIT = 500;

/** Parse with memoisation. Throws the same error every time for a bad source. */
export function parseCached(source: string): Node {
  const hit = parseCache.get(source);
  if (hit) {
    if (hit instanceof FormulaError) throw hit;
    return hit;
  }
  let result: Node | FormulaError;
  try {
    result = parse(source);
  } catch (err) {
    result = err instanceof FormulaError ? err : new FormulaError(String(err));
  }
  if (parseCache.size >= PARSE_CACHE_LIMIT) parseCache.clear();
  parseCache.set(source, result);
  if (result instanceof FormulaError) throw result;
  return result;
}

class Parser {
  private pos = 0;

  constructor(private readonly tokens: Token[]) {}

  private peek(): Token {
    return this.tokens[this.pos]!;
  }

  private next(): Token {
    const t = this.tokens[this.pos]!;
    if (t.kind !== 'eof') this.pos++;
    return t;
  }

  private atPunct(...values: string[]): boolean {
    const t = this.peek();
    return t.kind === 'punct' && values.includes(t.value);
  }

  private atKeyword(...words: string[]): boolean {
    const t = this.peek();
    return t.kind === 'ident' && words.includes(t.value.toUpperCase());
  }

  expectEof(): void {
    const t = this.peek();
    if (t.kind !== 'eof') {
      throw new FormulaError(`Unexpected ${JSON.stringify(t.value)} after the end of the expression`, t.pos);
    }
  }

  private expectPunct(value: string): void {
    const t = this.peek();
    if (t.kind !== 'punct' || t.value !== value) {
      throw new FormulaError(
        `Expected ${JSON.stringify(value)} but found ${t.kind === 'eof' ? 'end of expression' : JSON.stringify(t.value)}`,
        t.pos,
      );
    }
    this.next();
  }

  expression(depth: number): Node {
    if (depth > MAX_DEPTH) throw new FormulaError('Expression nests too deeply');
    return this.or(depth);
  }

  private or(depth: number): Node {
    let left = this.and(depth);
    while (this.atKeyword('OR')) {
      this.next();
      left = { k: 'binary', op: 'OR', l: left, r: this.and(depth + 1) };
    }
    return left;
  }

  private and(depth: number): Node {
    let left = this.comparison(depth);
    while (this.atKeyword('AND')) {
      this.next();
      left = { k: 'binary', op: 'AND', l: left, r: this.comparison(depth + 1) };
    }
    return left;
  }

  private comparison(depth: number): Node {
    let left = this.concat(depth);
    while (this.atPunct('=', '!=', '>', '<', '>=', '<=')) {
      const op = this.next().value as BinaryOp;
      left = { k: 'binary', op, l: left, r: this.concat(depth + 1) };
    }
    return left;
  }

  private concat(depth: number): Node {
    let left = this.additive(depth);
    while (this.atPunct('&')) {
      this.next();
      left = { k: 'binary', op: '&', l: left, r: this.additive(depth + 1) };
    }
    return left;
  }

  private additive(depth: number): Node {
    let left = this.multiplicative(depth);
    while (this.atPunct('+', '-')) {
      const op = this.next().value as BinaryOp;
      left = { k: 'binary', op, l: left, r: this.multiplicative(depth + 1) };
    }
    return left;
  }

  private multiplicative(depth: number): Node {
    let left = this.unary(depth);
    while (this.atPunct('*', '/')) {
      const op = this.next().value as BinaryOp;
      left = { k: 'binary', op, l: left, r: this.unary(depth + 1) };
    }
    return left;
  }

  private unary(depth: number): Node {
    if (this.atPunct('-')) {
      this.next();
      return { k: 'unary', op: '-', v: this.unary(depth + 1) };
    }
    if (this.atPunct('+')) {
      this.next();
      return this.unary(depth + 1);
    }
    if (this.atKeyword('NOT')) {
      this.next();
      return { k: 'unary', op: 'NOT', v: this.unary(depth + 1) };
    }
    return this.primary(depth);
  }

  private primary(depth: number): Node {
    const t = this.next();

    if (t.kind === 'eof') throw new FormulaError('Unexpected end of expression', t.pos);

    if (t.kind === 'num') {
      const v = Number(t.value);
      if (!Number.isFinite(v)) throw new FormulaError(`Invalid number ${JSON.stringify(t.value)}`, t.pos);
      return { k: 'num', v };
    }

    if (t.kind === 'str') return { k: 'str', v: t.value };

    if (t.kind === 'punct' && t.value === '(') {
      const inner = this.expression(depth + 1);
      this.expectPunct(')');
      return inner;
    }

    if (t.kind === 'ident') {
      const upper = t.value.toUpperCase();
      if (upper === 'TRUE') return { k: 'bool', v: true };
      if (upper === 'FALSE') return { k: 'bool', v: false };
      if (upper === 'NULL') return { k: 'null' };

      if (this.atPunct('(')) {
        if (t.value.includes('.')) {
          throw new FormulaError(`Function names cannot contain a dot: ${JSON.stringify(t.value)}`, t.pos);
        }
        this.next();
        const args: Node[] = [];
        if (!this.atPunct(')')) {
          args.push(this.expression(depth + 1));
          while (this.atPunct(',')) {
            this.next();
            args.push(this.expression(depth + 1));
          }
        }
        this.expectPunct(')');
        return { k: 'call', name: upper, args };
      }

      const segments = t.value.split('.');
      if (segments.some((s) => s === '')) {
        throw new FormulaError(`Malformed reference ${JSON.stringify(t.value)}`, t.pos);
      }
      if (segments.length > 6) {
        throw new FormulaError(`Reference ${JSON.stringify(t.value)} follows too many links (max 5)`, t.pos);
      }
      return { k: 'path', v: segments };
    }

    throw new FormulaError(`Unexpected ${JSON.stringify(t.value)}`, t.pos);
  }
}

/** Every column/table reference in an expression, as dotted paths. */
export function collectPaths(node: Node, out: string[][] = []): string[][] {
  switch (node.k) {
    case 'path':
      out.push(node.v);
      break;
    case 'unary':
      collectPaths(node.v, out);
      break;
    case 'binary':
      collectPaths(node.l, out);
      collectPaths(node.r, out);
      break;
    case 'call':
      for (const a of node.args) collectPaths(a, out);
      break;
    default:
      break;
  }
  return out;
}