wellformed

import {
  AGGREGATE_NAMES,
  FUNCTIONS_BY_NAME,
  FormulaError,
  isComputed,
  parseCached,
  q,
  type ColumnMeta,
  type Node,
} from '@wellformed/shared';
import { findColumn, type MetaSnapshot } from '../db/meta.js';

/**
 * Compiles the Wellformed formula language into SQLite expressions.
 *
 * This is what makes the SQL bar honest: a formula column is not evaluated in
 * JavaScript after the fact, it becomes part of the SELECT list, so the same
 * expression can also be filtered, sorted and aggregated by the database.
 *
 * Values from the expression are always bound, never interpolated. Identifiers
 * are interpolated, which is safe because every name passed the identifier
 * rules in `@wellformed/shared` before it was ever written to the catalogue.
 */

export interface Fragment {
  sql: string;
  params: unknown[];
}

const MAX_COLUMN_NESTING = 8;

export interface JoinNode {
  alias: string;
  targetTable: string;
  fromAlias: string;
  fromColumn: string;
}

/** One FROM clause: the outer query, or a rollup's correlated subquery. */
export class Frame {
  readonly joins: JoinNode[] = [];
  private readonly byKey = new Map<string, JoinNode>();

  link(fromAlias: string, fromColumn: string, targetTable: string, aliases: Aliases): JoinNode {
    const key = `${fromAlias}\u0000${fromColumn}`;
    const existing = this.byKey.get(key);
    if (existing) return existing;
    const node: JoinNode = { alias: aliases.next('j'), targetTable, fromAlias, fromColumn };
    this.byKey.set(key, node);
    this.joins.push(node);
    return node;
  }

  /** Joins render in creation order, so `a.b.c` always joins `a.b` first. */
  joinSql(): string {
    return this.joins
      .map(
        (j) =>
          ` LEFT JOIN ${q(j.targetTable)} ${j.alias} ON ${j.fromAlias}.${q(j.fromColumn)} = ${j.alias}."id"`,
      )
      .join('');
  }
}

export class Aliases {
  private n = 0;
  next(prefix: string): string {
    return `${prefix}${this.n++}`;
  }
}

export interface Scope {
  table: string;
  alias: string;
  frame: Frame;
}

export class Compiler {
  /** `table.Column` entries currently being expanded, for cycle detection. */
  private readonly stack: string[] = [];

  constructor(
    private readonly meta: MetaSnapshot,
    readonly aliases: Aliases = new Aliases(),
  ) {}

  /** Compile a raw expression source in the given scope. */
  compileSource(source: string, scope: Scope): Fragment {
    return this.compile(parseCached(source), scope);
  }

  compile(node: Node, scope: Scope): Fragment {
    switch (node.k) {
      case 'num':
        return { sql: String(node.v), params: [] };
      case 'bool':
        return { sql: node.v ? '1' : '0', params: [] };
      case 'null':
        return { sql: 'NULL', params: [] };
      case 'str':
        return { sql: '?', params: [node.v] };
      case 'path':
        return this.resolvePath(node.v, scope);
      case 'unary': {
        const v = this.compile(node.v, scope);
        return node.op === '-'
          ? { sql: `(-COALESCE(${v.sql}, 0))`, params: v.params }
          : { sql: `(NOT (${v.sql}))`, params: v.params };
      }
      case 'binary':
        return this.compileBinary(node, scope);
      case 'call':
        return this.compileCall(node, scope);
    }
  }

  private compileBinary(
    node: Extract<Node, { k: 'binary' }>,
    scope: Scope,
  ): Fragment {
    const l = this.compile(node.l, scope);
    const r = this.compile(node.r, scope);
    const params = [...l.params, ...r.params];

    switch (node.op) {
      case '+':
      case '-':
      case '*':
        // Spreadsheet arithmetic: an empty cell counts as zero rather than
        // poisoning the whole expression with NULL.
        return {
          sql: `(COALESCE(${l.sql}, 0) ${node.op} COALESCE(${r.sql}, 0))`,
          params,
        };
      case '/':
        // SQLite yields NULL for division by zero; `* 1.0` forces real division.
        return {
          sql: `((COALESCE(${l.sql}, 0) * 1.0) / COALESCE(${r.sql}, 0))`,
          params,
        };
      case '&':
        return {
          sql: `(COALESCE(CAST(${l.sql} AS TEXT), '') || COALESCE(CAST(${r.sql} AS TEXT), ''))`,
          params,
        };
      case '=':
        // `IS` rather than `=` so blank compares equal to blank.
        return { sql: `(${l.sql} IS ${r.sql})`, params };
      case '!=':
        return { sql: `(${l.sql} IS NOT ${r.sql})`, params };
      case '>':
      case '<':
      case '>=':
      case '<=':
        return { sql: `(${l.sql} ${node.op} ${r.sql})`, params };
      case 'AND':
        return { sql: `(${l.sql} AND ${r.sql})`, params };
      case 'OR':
        return { sql: `(${l.sql} OR ${r.sql})`, params };
    }
  }

  /** Resolve `Column`, `Link.Column`, `Link.Link.Column`, or `Link.id`. */
  resolvePath(parts: string[], scope: Scope): Fragment {
    let table = scope.table;
    let alias = scope.alias;

    for (let i = 0; i < parts.length; i++) {
      const segment = parts[i]!;
      const schema = this.meta.tables[table];
      if (!schema) throw new FormulaError(`Unknown table "${table}".`);
      const last = i === parts.length - 1;

      if (last) {
        if (segment === 'id') return { sql: `${alias}."id"`, params: [] };
        const col = findColumn(schema, segment);
        if (!col) throw new FormulaError(`"${table}" has no column named "${segment}".`);
        return this.columnRef(table, alias, col, scope.frame);
      }

      const col = findColumn(schema, segment);
      if (!col) throw new FormulaError(`"${table}" has no column named "${segment}".`);
      if (col.type !== 'link' || !col.target) {
        throw new FormulaError(
          `"${segment}" is a ${col.type} column, not a link, so "${parts.join('.')}" cannot follow it.`,
        );
      }
      const join = scope.frame.link(alias, col.name, col.target, this.aliases);
      alias = join.alias;
      table = col.target;
    }

    // Unreachable: the loop always returns on its final iteration.
    throw new FormulaError('Empty reference.');
  }

  private columnRef(table: string, alias: string, col: ColumnMeta, frame: Frame): Fragment {
    if (col.type === 'link') {
      // A bare link reference reads as the linked row's label, which is what
      // the grid shows and therefore what a filter on it should mean.
      if (!col.target) throw new FormulaError(`Link column "${col.name}" has no target table.`);
      const target = this.meta.tables[col.target];
      if (!target) throw new FormulaError(`Link column "${col.name}" points at a missing table.`);
      const join = frame.link(alias, col.name, col.target, this.aliases);
      return { sql: `${join.alias}.${q(target.display)}`, params: [] };
    }

    if (isComputed(col.type)) {
      const key = `${table}.${col.name}`;
      if (this.stack.includes(key)) {
        throw new FormulaError(`Circular reference: ${[...this.stack, key].join(' → ')}.`);
      }
      if (this.stack.length >= MAX_COLUMN_NESTING) {
        throw new FormulaError(
          `Computed columns are nested more than ${MAX_COLUMN_NESTING} deep at "${key}".`,
        );
      }
      if (!col.expr) throw new FormulaError(`Column "${key}" has no expression.`);
      this.stack.push(key);
      try {
        const inner = this.compile(parseCached(col.expr), { table, alias, frame });
        return { sql: `(${inner.sql})`, params: inner.params };
      } finally {
        this.stack.pop();
      }
    }

    return { sql: `${alias}.${q(col.name)}`, params: [] };
  }

  private compileCall(node: Extract<Node, { k: 'call' }>, scope: Scope): Fragment {
    const spec = FUNCTIONS_BY_NAME.get(node.name);
    if (!spec) {
      throw new FormulaError(
        `Unknown function ${node.name}(). Available: ${[...FUNCTIONS_BY_NAME.keys()].join(', ')}.`,
      );
    }
    if (node.args.length < spec.minArgs || node.args.length > spec.maxArgs) {
      throw new FormulaError(
        `${node.name}() takes ${describeArity(spec.minArgs, spec.maxArgs)}, but got ${node.args.length}. Signature: ${spec.signature}`,
      );
    }

    if (AGGREGATE_NAMES.has(node.name)) {
      const rollup = this.tryRollup(node, scope);
      if (rollup) return rollup;
      if (node.name === 'SUM' || node.name === 'COUNT' || node.name === 'AVG') {
        throw new FormulaError(
          `${node.name}() aggregates rows from a linked table; its first argument must name one, e.g. ${spec.signature}.`,
        );
      }
      if (node.args.length < 2) {
        throw new FormulaError(
          `${node.name}() needs either a linked column (${node.name}(orders.Total)) or two values to compare (${node.name}(a, b)).`,
        );
      }
    }

    const args = node.args.map((a) => this.compile(a, scope));
    const s = args.map((a) => a.sql);
    const params = args.flatMap((a) => a.params);
    const text = (x: string) => `COALESCE(CAST(${x} AS TEXT), '')`;

    switch (node.name) {
      case 'IF':
        return { sql: `(CASE WHEN ${s[0]} THEN ${s[1]} ELSE ${s[2] ?? 'NULL'} END)`, params };
      case 'COALESCE':
        return { sql: `COALESCE(${s.join(', ')})`, params };
      case 'ISBLANK':
        return { sql: `(${text(s[0]!)} = '')`, params };
      case 'ROUND':
        return { sql: `ROUND(${s[0]}${s[1] ? `, ${s[1]}` : ''})`, params };
      case 'ABS':
        return { sql: `ABS(${s[0]})`, params };
      case 'MIN':
      case 'MAX':
        return { sql: `${node.name}(${s.join(', ')})`, params };
      case 'CONCAT':
        return { sql: `(${s.map(text).join(' || ')})`, params };
      case 'UPPER':
      case 'LOWER':
      case 'TRIM':
        return { sql: `${node.name}(${s[0]})`, params };
      case 'LENGTH':
        return { sql: `LENGTH(${text(s[0]!)})`, params };
      case 'TODAY':
        return { sql: `date('now')`, params: [] };
      case 'NOW':
        return { sql: `datetime('now')`, params: [] };
      case 'DATEDIFF':
        return {
          sql: `CAST(julianday(${s[0]}) - julianday(${s[1]}) AS INTEGER)`,
          params,
        };
      case 'YEAR':
        return { sql: `CAST(strftime('%Y', ${s[0]}) AS INTEGER)`, params };
      case 'MONTH':
        return { sql: `CAST(strftime('%m', ${s[0]}) AS INTEGER)`, params };
      case 'DAY':
        return { sql: `CAST(strftime('%d', ${s[0]}) AS INTEGER)`, params };
      default:
        throw new FormulaError(`Function ${node.name}() is declared but not implemented.`);
    }
  }

  /**
   * `SUM(orders.Total)` becomes a correlated subquery over the rows of
   * `orders` whose link column points back at the current row. Returns null
   * when the call is not shaped like a rollup, so MIN/MAX can fall through to
   * their scalar meaning.
   */
  private tryRollup(node: Extract<Node, { k: 'call' }>, scope: Scope): Fragment | null {
    const first = node.args[0];
    if (!first || first.k !== 'path') return null;

    const [head, ...rest] = first.v;
    if (!head) return null;
    const target = this.meta.tables[head];
    if (!target) return null;

    let backLink: ColumnMeta | undefined;
    const candidates = target.columns.filter((c) => c.type === 'link' && c.target === scope.table);
    const second = node.args[1];

    if (second) {
      if (second.k !== 'str') {
        throw new FormulaError(
          `The second argument to ${node.name}() must be a quoted link column name, e.g. ${node.name}(${head}.Total, "Customer").`,
        );
      }
      backLink = candidates.find((c) => c.name === second.v);
      if (!backLink) {
        throw new FormulaError(
          `"${head}.${second.v}" is not a link column pointing at "${scope.table}".`,
        );
      }
    } else if (candidates.length === 0) {
      throw new FormulaError(
        `"${head}" has no link column pointing at "${scope.table}", so ${node.name}() has nothing to group by.`,
      );
    } else if (candidates.length > 1) {
      throw new FormulaError(
        `"${head}" links to "${scope.table}" more than once (${candidates
          .map((c) => c.name)
          .join(', ')}). Say which to use: ${node.name}(${head}.…, "${candidates[0]!.name}").`,
      );
    } else {
      backLink = candidates[0]!;
    }

    const subAlias = this.aliases.next('r');
    const subFrame = new Frame();
    const subScope: Scope = { table: head, alias: subAlias, frame: subFrame };

    let inner: Fragment | null = null;
    if (rest.length) {
      inner = this.resolvePath(rest, subScope);
    } else if (node.name !== 'COUNT') {
      throw new FormulaError(`${node.name}(${head}) needs a column, e.g. ${node.name}(${head}.Total).`);
    }

    const agg =
      node.name === 'COUNT'
        ? inner
          ? `COUNT(${inner.sql})`
          : 'COUNT(*)'
        : `${node.name}(${inner!.sql})`;

    // An empty set sums to zero and counts to zero; min/max/avg stay blank.
    const wrapped = node.name === 'SUM' || node.name === 'COUNT' ? `COALESCE(${agg}, 0)` : agg;

    const sql =
      `(SELECT ${wrapped} FROM ${q(head)} ${subAlias}${subFrame.joinSql()}` +
      ` WHERE ${subAlias}.${q(backLink.name)} = ${scope.alias}."id")`;

    return { sql, params: inner?.params ?? [] };
  }
}

function describeArity(min: number, max: number): string {
  if (min === max) return `${min} argument${min === 1 ? '' : 's'}`;
  return `${min}–${max} arguments`;
}

/**
 * Type-check an expression against the catalogue without running it.
 * Returns the error message, or null when it compiles.
 */
export function validateExpression(meta: MetaSnapshot, table: string, source: string): string | null {
  try {
    const compiler = new Compiler(meta);
    compiler.compileSource(source, { table, alias: 't0', frame: new Frame() });
    return null;
  } catch (err) {
    if (err instanceof FormulaError) return err.message;
    throw err;
  }
}