import { tokenize } from './lexer.js';
import { Parser, inferPathKey, mapChildren } from './parser.js';
import { SqlSyntaxError, SqlEvaluationError } from './errors.js';
import { evaluate } from './evaluate.js';
import { aggregates, isAggregateName } from './aggregates.js';

/**
 * @typedef {{kind: 'Star'}} StarSelect
 * @typedef {{kind: 'Shape', expr: import('./parser.js').ObjectLiteralNode}} ShapeSelect
 * @typedef {{kind: 'Columns', items: {expr: import('./parser.js').Expr, key: string}[]}} ColumnsSelect
 * @typedef {StarSelect|ShapeSelect|ColumnsSelect} Select
 * @typedef {{kind: 'Source', expr: import('./parser.js').Expr, alias: string|null}} SourceFromItem
 * @typedef {{kind: 'Unnest', expr: import('./parser.js').Expr, alias: string|null}} UnnestFromItem
 * @typedef {SourceFromItem|UnnestFromItem} FromItem
 * @typedef {{type: 'INNER'|'LEFT'|'CROSS', item: FromItem, on: import('./parser.js').Expr|null}} JoinClause
 * @typedef {{first: FromItem, joins: JoinClause[]}} From
 * @typedef {{expr: import('./parser.js').Expr, dir: 'ASC'|'DESC'}} OrderItem
 * @typedef {{type: 'Query', select: Select, from: From, where: import('./parser.js').Expr|null, groupBy: import('./parser.js').Expr[]|null, having: import('./parser.js').Expr|null, orderBy: OrderItem[]|null, limit: import('./parser.js').Expr|null}} Query
 * @typedef {{key: string, fold: (values: any[]) => any, arg: import('./parser.js').Expr|import('./parser.js').StarNode}} PlannedAggregate
 * @typedef {{select: Select, having: import('./parser.js').Expr|null, orderBy: OrderItem[]|null, groupBy: import('./parser.js').Expr[]|null, aggregates: PlannedAggregate[], grouped: boolean}} Plan
 */

class QueryParser extends Parser {
  /** @returns {Query} */
  parseQuery() {
    this.expect('SELECT');
    const select = this.parseSelect();
    this.expect('FROM');
    const from = this.parseFrom();
    let where = null;
    if (this.check('WHERE')) {
      this.next();
      where = this.parseOr();
    }
    let groupBy = null;
    if (this.check('GROUP')) {
      this.next();
      this.expect('BY');
      groupBy = this.parseGroupByList();
    }
    let having = null;
    if (this.check('HAVING')) {
      this.next();
      having = this.parseOr();
    }
    let orderBy = null;
    if (this.check('ORDER')) {
      this.next();
      this.expect('BY');
      orderBy = this.parseOrderByList();
    }
    let limit = null;
    if (this.check('LIMIT')) {
      this.next();
      limit = this.parseOr();
    }
    this.expect('eof');
    /** @type {Query} */
    const ast = { type: 'Query', select, from, where, groupBy, having, orderBy, limit };
    // Planning is what rejects misplaced/nested aggregates and SELECT * with
    // GROUP BY; running it here surfaces those at parse time. The result is
    // recomputed (cheaply) by executeQuery, which also accepts hand-built ASTs.
    planQuery(ast);
    return ast;
  }

  /** @returns {import('./parser.js').Expr[]} */
  parseGroupByList() {
    const items = [this.parseOr()];
    while (this.check(',')) {
      this.next();
      items.push(this.parseOr());
    }
    return items;
  }

  /** @returns {OrderItem[]} */
  parseOrderByList() {
    const items = [this.parseOrderByItem()];
    while (this.check(',')) {
      this.next();
      items.push(this.parseOrderByItem());
    }
    return items;
  }

  /** @returns {OrderItem} */
  parseOrderByItem() {
    const expr = this.parseOr();
    let dir = 'ASC';
    if (this.check('ASC')) {
      this.next();
    } else if (this.check('DESC')) {
      this.next();
      dir = 'DESC';
    }
    return { expr, dir };
  }

  /** @returns {Select} */
  parseSelect() {
    if (this.check('*')) {
      this.next();
      return { kind: 'Star' };
    }
    if (this.check('{')) {
      const expr = this.parseObjectLiteral();
      return { kind: 'Shape', expr };
    }
    const items = [this.parseSelectItem()];
    while (this.check(',')) {
      this.next();
      items.push(this.parseSelectItem());
    }
    return { kind: 'Columns', items };
  }

  /** @returns {{expr: import('./parser.js').Expr, key: string}} */
  parseSelectItem() {
    const startTok = this.peek();
    const expr = this.parseOr();
    if (this.check('AS')) {
      this.next();
      const key = this.expectIdentLike().value;
      return { expr, key };
    }
    return { expr, key: inferPathKey(expr, startTok) };
  }

  /** @returns {From} */
  parseFrom() {
    const first = this.parseFromItem();
    /** @type {JoinClause[]} */
    const joins = [];
    while (this.isJoinStart()) {
      joins.push(this.parseJoinClause());
    }
    validateFromAliases(first, joins);
    return { first, joins };
  }

  /** @returns {boolean} */
  isJoinStart() {
    return this.check('JOIN') || this.check('INNER') || this.check('LEFT') || this.check('CROSS');
  }

  /** @returns {FromItem} */
  parseFromItem() {
    if (this.check('UNNEST')) {
      this.next();
      this.expect('(');
      const expr = this.parseOr();
      this.expect(')');
      let alias = null;
      if (this.check('AS')) {
        this.next();
        alias = this.expectIdentLike().value;
      }
      return { kind: 'Unnest', expr, alias };
    }
    const expr = this.parsePostfix();
    let alias = null;
    if (this.check('AS')) {
      this.next();
      alias = this.expectIdentLike().value;
    }
    return { kind: 'Source', expr, alias };
  }

  /** @returns {JoinClause} */
  parseJoinClause() {
    const startTok = this.peek();
    /** @type {'INNER'|'LEFT'|'CROSS'} */
    let type = 'INNER';
    if (this.check('INNER')) {
      this.next();
      this.expect('JOIN');
    } else if (this.check('LEFT')) {
      this.next();
      this.expect('JOIN');
      type = 'LEFT';
    } else if (this.check('CROSS')) {
      this.next();
      this.expect('JOIN');
      type = 'CROSS';
    } else {
      this.expect('JOIN');
    }

    const item = this.parseFromItem();

    let on = null;
    if (this.check('ON')) {
      this.next();
      on = this.parseOr();
    }

    if (type === 'CROSS' && on) {
      throw new SqlSyntaxError('CROSS JOIN cannot have an ON condition', startTok);
    }
    if (type !== 'CROSS' && item.kind !== 'Unnest' && !on) {
      throw new SqlSyntaxError('JOIN requires an ON condition (except CROSS JOIN or JOIN UNNEST(...))', startTok);
    }

    return { type, item, on };
  }
}

/**
 * Works out how a query has to be run: which aggregates it computes, what the
 * SELECT/HAVING/ORDER BY expressions look like once those aggregates have been
 * lifted out, and whether the query is grouped at all. Pure - it never mutates
 * `ast`, so a parsed query can be executed repeatedly.
 * @param {Query} ast
 * @returns {Plan}
 */
function planQuery(ast) {
  assertNoAggregates(ast.where, 'WHERE');
  assertNoAggregates(ast.from.first.expr, 'FROM');
  for (const expr of ast.groupBy ?? []) assertNoAggregates(expr, 'GROUP BY');
  for (const join of ast.from.joins) {
    assertNoAggregates(join.item.expr, 'JOIN');
    assertNoAggregates(join.on, 'JOIN ... ON');
  }

  /** @type {PlannedAggregate[]} */
  const found = [];
  const select = mapSelect(ast.select, (expr) => extractAggregates(expr, found));
  const having = ast.having === null ? null : extractAggregates(ast.having, found);
  const orderBy =
    ast.orderBy === null
      ? null
      : resolveOrderAliases(ast.orderBy, ast.select).map((item) => ({
          ...item,
          expr: extractAggregates(item.expr, found),
        }));

  const grouped = ast.groupBy !== null || found.length > 0;
  if (grouped && ast.select.kind === 'Star') {
    throw new SqlSyntaxError('SELECT * cannot be combined with GROUP BY or aggregate functions; list the columns instead');
  }
  if (ast.having !== null && !grouped) {
    throw new SqlSyntaxError('HAVING requires GROUP BY or an aggregate function; use WHERE to filter rows');
  }

  return { select, having, orderBy, groupBy: ast.groupBy, aggregates: found, grouped };
}

/**
 * Replaces every aggregate call in `node` with a reference to a generated `$aggN`
 * binding, recording in `sink` how to compute it. `$` is not a legal identifier
 * character in the lexer, so these names can never collide with anything the
 * query itself could have written.
 * @param {import('./parser.js').Expr} node
 * @param {PlannedAggregate[]} sink
 * @returns {import('./parser.js').Expr}
 */
function extractAggregates(node, sink) {
  if (node.type === 'Call' && isAggregateName(node.name)) {
    const name = node.name.toUpperCase();
    if (node.args.length !== 1) {
      throw new SqlSyntaxError(`${name} takes exactly one argument`);
    }
    const [arg] = node.args;
    if (arg.type !== 'Star' && containsAggregate(arg)) {
      throw new SqlSyntaxError('Aggregate functions cannot be nested');
    }
    const key = `$agg${sink.length}`;
    sink.push({ key, fold: aggregates[name], arg });
    return { type: 'Identifier', name: key };
  }
  return mapChildren(node, (child) => extractAggregates(child, sink));
}

/** @param {import('./parser.js').Expr} node */
function containsAggregate(node) {
  if (node.type === 'Call' && isAggregateName(node.name)) return true;
  let found = false;
  mapChildren(node, (child) => {
    if (containsAggregate(child)) found = true;
    return child;
  });
  return found;
}

/** @param {import('./parser.js').Expr|null} node @param {string} clause */
function assertNoAggregates(node, clause) {
  if (node !== null && containsAggregate(node)) {
    throw new SqlSyntaxError(`Aggregate functions are not allowed in ${clause}`);
  }
}

/**
 * @param {Select} select
 * @param {(expr: import('./parser.js').Expr) => import('./parser.js').Expr} fn
 * @returns {Select}
 */
function mapSelect(select, fn) {
  if (select.kind === 'Star') return select;
  if (select.kind === 'Shape') {
    return { kind: 'Shape', expr: /** @type {import('./parser.js').ObjectLiteralNode} */ (fn(select.expr)) };
  }
  return { kind: 'Columns', items: select.items.map((item) => ({ ...item, expr: fn(item.expr) })) };
}

/**
 * Lets ORDER BY name an output column, as SQL does, so `COUNT(*) AS aantal ...
 * ORDER BY aantal DESC` sorts by the aggregate instead of silently finding
 * nothing. Only bare identifiers are resolved - `ORDER BY p.aantal` stays a path
 * into the row.
 * @param {OrderItem[]} orderBy
 * @param {Select} select
 * @returns {OrderItem[]}
 */
function resolveOrderAliases(orderBy, select) {
  const outputs = selectOutputs(select);
  if (outputs === null) return orderBy;
  return orderBy.map((item) => {
    if (item.expr.type !== 'Identifier') return item;
    const aliased = outputs.get(item.expr.name);
    return aliased === undefined ? item : { ...item, expr: aliased };
  });
}

/** @param {Select} select @returns {Map<string, import('./parser.js').Expr>|null} */
function selectOutputs(select) {
  if (select.kind === 'Star') return null;
  if (select.kind === 'Shape') return new Map(select.expr.properties.map((p) => [p.key, p.value]));
  return new Map(select.items.map((item) => [item.key, item.expr]));
}

/**
 * @param {FromItem} first
 * @param {JoinClause[]} joins
 */
function validateFromAliases(first, joins) {
  const hasJoins = joins.length > 0;
  if (first.kind === 'Unnest' && !first.alias) {
    throw new SqlSyntaxError('UNNEST(...) in FROM requires AS alias');
  }
  if (hasJoins && !first.alias) {
    throw new SqlSyntaxError('FROM sources must be aliased when using JOIN');
  }
  for (const join of joins) {
    if (!join.item.alias) {
      throw new SqlSyntaxError('Joined sources must be aliased');
    }
  }
}

/**
 * Parses a `SELECT ... FROM ... [WHERE ...] [GROUP BY ...] [HAVING ...] [ORDER BY ...] [LIMIT ...]`
 * query string into an AST.
 * @param {string} source
 * @returns {Query}
 */
export function parseQuery(source) {
  const tokens = tokenize(source);
  const parser = new QueryParser(tokens);
  return parser.parseQuery();
}

/**
 * Executes a parsed query against a JSON context object.
 * @param {Query} ast
 * @param {Record<string, any>} context
 * @param {{functions?: Record<string, (...args: any[]) => any>}} [options]
 * @returns {any[]}
 */
export function executeQuery(ast, context, options = {}) {
  const hasJoins = ast.from.joins.length > 0;
  const firstArray = resolveFromItemArray(ast.from.first, context, options);

  /** @type {{ctx: any, pristineRow: any}[]} */
  let rows = firstArray.map((row) => ({
    ctx: ast.from.first.alias ? bind({}, row, ast.from.first.alias) : row,
    pristineRow: row,
  }));

  for (const join of ast.from.joins) {
    /** @type {{ctx: any, pristineRow: any}[]} */
    const nextRows = [];
    for (const left of rows) {
      const scopeCtx = { ...context, ...left.ctx };
      const rightArray = resolveFromItemArray(join.item, scopeCtx, options);
      let matchedAny = false;
      for (const rightRow of rightArray) {
        const ctx = bind(left.ctx, rightRow, /** @type {string} */ (join.item.alias));
        if (join.on && evaluate(join.on, ctx, options) !== true) continue;
        nextRows.push({ ctx, pristineRow: undefined });
        matchedAny = true;
      }
      if (!matchedAny && join.type === 'LEFT') {
        nextRows.push({ ctx: bind(left.ctx, null, /** @type {string} */ (join.item.alias)), pristineRow: undefined });
      }
    }
    rows = nextRows;
  }

  const plan = planQuery(ast);

  /** @type {{ctx: any, pristineRow: any}[]} */
  let matched = [];
  for (const r of rows) {
    if (ast.where && evaluate(ast.where, r.ctx, options) !== true) continue;
    matched.push(r);
  }

  if (plan.grouped) {
    matched = groupRows(matched, plan, context, options);
    if (plan.having) {
      matched = matched.filter((r) => evaluate(plan.having, r.ctx, options) === true);
    }
  }

  if (plan.orderBy) {
    const comparator = makeRowComparator(plan.orderBy, options);
    matched = matched.slice().sort(comparator);
  }

  if (ast.limit) {
    const n = evaluate(ast.limit, context, options);
    if (n !== null) {
      if (typeof n !== 'number' || !Number.isInteger(n) || n < 0) {
        throw new SqlEvaluationError('LIMIT must be a non-negative integer or NULL');
      }
      matched = matched.slice(0, n);
    }
  }

  return matched.map((r) =>
    plan.select.kind === 'Star' && !hasJoins ? r.pristineRow : project(plan.select, r.ctx, options),
  );
}

/**
 * Folds the matched rows into one row per distinct GROUP BY key - or a single row
 * when the query has no GROUP BY but does use aggregates. A group's context is its
 * first row's context plus the computed `$aggN` bindings, so expressions that
 * aren't aggregated still resolve: against an arbitrary row of the group, the way
 * SQLite treats bare columns, rather than being rejected outright.
 * @param {{ctx: any, pristineRow: any}[]} rows
 * @param {Plan} plan
 * @param {Record<string, any>} context
 * @param {{functions?: Record<string, (...args: any[]) => any>}} options
 * @returns {{ctx: any, pristineRow: any}[]}
 */
function groupRows(rows, plan, context, options) {
  /** @type {Map<string, {ctx: any, pristineRow: any}[]>} */
  const groups = new Map();
  for (const r of rows) {
    // Serializing the key values is what makes grouping by an object or array
    // (not just a scalar) work at all; it costs key-order sensitivity, which
    // only shows up for objects built with their keys in differing orders.
    const key = plan.groupBy === null ? '' : JSON.stringify(plan.groupBy.map((e) => evaluate(e, r.ctx, options)));
    const group = groups.get(key);
    if (group) group.push(r);
    else groups.set(key, [r]);
  }

  // `SELECT COUNT(*) AS n FROM empty` is still one row (n = 0), but grouping an
  // empty set of rows yields no groups at all.
  if (plan.groupBy === null && groups.size === 0) groups.set('', []);

  const out = [];
  for (const group of groups.values()) {
    const ctx = { ...(group.length > 0 ? group[0].ctx : context) };
    for (const agg of plan.aggregates) {
      // COUNT(*) has no expression to evaluate; a non-null placeholder per row
      // is what turns "count non-nulls" into "count rows".
      const values = group.map((r) => (agg.arg.type === 'Star' ? true : evaluate(agg.arg, r.ctx, options)));
      ctx[agg.key] = agg.fold(values);
    }
    out.push({ ctx, pristineRow: undefined });
  }
  return out;
}

/**
 * @param {FromItem} item
 * @param {Record<string, any>} scopeCtx
 * @param {{functions?: Record<string, (...args: any[]) => any>}} options
 */
function resolveFromItemArray(item, scopeCtx, options) {
  const value = evaluate(item.expr, scopeCtx, options);
  if (item.kind === 'Unnest') {
    if (value === null || value === undefined) return [];
    if (!Array.isArray(value)) throw new SqlEvaluationError('UNNEST requires an array value');
    return value;
  }
  if (!Array.isArray(value)) throw new SqlEvaluationError('FROM/JOIN source must reference an array');
  return value;
}

/**
 * Binds a row into a context: its own fields (if it's a plain object) plus its alias.
 * @param {Record<string, any>} baseCtx
 * @param {any} row
 * @param {string} alias
 */
function bind(baseCtx, row, alias) {
  const flat = isPlainObject(row) ? row : {};
  return { ...baseCtx, ...flat, [alias]: row };
}

/** @param {any} v */
function isPlainObject(v) {
  return typeof v === 'object' && v !== null && !Array.isArray(v);
}

/**
 * @param {OrderItem[]} orderBy
 * @param {{functions?: Record<string, (...args: any[]) => any>}} options
 */
function makeRowComparator(orderBy, options) {
  return (a, b) => {
    for (const item of orderBy) {
      const va = evaluate(item.expr, a.ctx, options);
      const vb = evaluate(item.expr, b.ctx, options);
      // Nulls always sort last, regardless of ASC/DESC - checked before the
      // direction flip below, so DESC never pulls nulls back to the front.
      if (va === null && vb === null) continue;
      if (va === null) return 1;
      if (vb === null) return -1;
      const cmp = compareForOrder(va, vb);
      if (cmp !== 0) return item.dir === 'DESC' ? -cmp : cmp;
    }
    return 0;
  };
}

/** @param {any} a @param {any} b */
function compareForOrder(a, b) {
  const ta = typeof a;
  const tb = typeof b;
  if (ta !== tb || (ta !== 'number' && ta !== 'string' && ta !== 'boolean')) {
    throw new SqlEvaluationError(`Cannot compare ${orderTypeName(a)} and ${orderTypeName(b)} in ORDER BY`);
  }
  if (a < b) return -1;
  if (a > b) return 1;
  return 0;
}

/** @param {any} v */
function orderTypeName(v) {
  if (v === null) return 'null';
  if (Array.isArray(v)) return 'array';
  return typeof v;
}

/**
 * @param {Select} select
 * @param {Record<string, any>} ctx
 * @param {{functions?: Record<string, (...args: any[]) => any>}} options
 */
function project(select, ctx, options) {
  if (select.kind === 'Star') return ctx;
  if (select.kind === 'Shape') return evaluate(select.expr, ctx, options);
  /** @type {Record<string, any>} */
  const out = {};
  for (const item of select.items) {
    out[item.key] = evaluate(item.expr, ctx, options);
  }
  return out;
}

/**
 * Parses (if needed) and executes a
 * `SELECT ... FROM ... [WHERE ...] [GROUP BY ...] [HAVING ...] [ORDER BY ...] [LIMIT ...]` query.
 * @param {string|Query} source - query text, or an already-parsed AST
 * @param {Record<string, any>} context
 * @param {{functions?: Record<string, (...args: any[]) => any>}} [options]
 * @returns {any[]}
 */
export function query(source, context, options) {
  const ast = typeof source === 'string' ? parseQuery(source) : source;
  return executeQuery(ast, context, options);
}