wellformed

import { literal, q, SQL_KEYWORDS, type SqlToken } from '@wellformed/shared';

/**
 * Substitute bound parameters back into a statement for display.
 *
 * Execution never uses this string — it exists so the SQL bar can show the
 * query with real values instead of a row of question marks. Placeholders
 * inside string literals are skipped, and values are escaped the same way
 * SQLite would.
 */
export function inlineParams(sql: string, params: readonly unknown[]): string {
  let out = '';
  let paramIndex = 0;
  let inString = false;

  for (let i = 0; i < sql.length; i++) {
    const ch = sql[i]!;
    if (inString) {
      out += ch;
      if (ch === "'") {
        if (sql[i + 1] === "'") {
          // Doubled quote: an escaped apostrophe, still inside the literal.
          out += "'";
          i++;
        } else {
          inString = false;
        }
      }
      continue;
    }
    if (ch === "'") {
      inString = true;
      out += ch;
      continue;
    }
    if (ch === '?') {
      out += paramIndex < params.length ? literal(params[paramIndex]) : '?';
      paramIndex++;
      continue;
    }
    out += ch;
  }
  return out;
}

const TOKEN_RE =
  /('(?:[^']|'')*')|(\b\d+(?:\.\d+)?\b)|("(?:[^"]|"")*")|([A-Za-z_][A-Za-z0-9_]*)|([(),.*/+\-<>=|]+)|(\s+)/g;

const FUNCTION_NAMES = new Set([
  'COUNT',
  'SUM',
  'AVG',
  'MIN',
  'MAX',
  'COALESCE',
  'ROUND',
  'ABS',
  'UPPER',
  'LOWER',
  'TRIM',
  'LENGTH',
  'CAST',
  'DATE',
  'DATETIME',
  'JULIANDAY',
  'STRFTIME',
  'IFNULL',
  'NULLIF',
]);

/* ---- the artificial variable layer -------------------------------------- */

/** A computed column's name and its expression as display SQL, literals inlined. */
export interface SqlDef {
  name: string;
  sql: string;
}

/** Collect a definition once; later sightings of the same column add nothing. */
export function addDef(defs: SqlDef[], name: string, sql: string): void {
  if (!defs.some((d) => d.name === name)) defs.push({ name, sql });
}

/**
 * Marks a folded-in variable: `\x00name\x00`. A sentinel rather than printable
 * syntax, so no column name — however SQL-shaped — can collide with real
 * statement text; the token pass turns it into `:name` at the very end.
 */
const VAR_MARK = '\x00';

function varMark(name: string): string {
  return `${VAR_MARK}${name}${VAR_MARK}`;
}

/**
 * Split into quoted and unquoted stretches, so alias detection and clause
 * breaking never act on the inside of a string, a quoted identifier, or a
 * folded-in variable name.
 */
function segments(sql: string): { text: string; quoted: boolean }[] {
  const out: { text: string; quoted: boolean }[] = [];
  let plain = '';
  for (let i = 0; i < sql.length; i++) {
    const ch = sql[i]!;
    if (ch === VAR_MARK) {
      if (plain) {
        out.push({ text: plain, quoted: false });
        plain = '';
      }
      let s = ch;
      for (i++; i < sql.length; i++) {
        s += sql[i];
        if (sql[i] === VAR_MARK) break;
      }
      out.push({ text: s, quoted: true });
      continue;
    }
    if (ch !== "'" && ch !== '"') {
      plain += ch;
      continue;
    }
    if (plain) {
      out.push({ text: plain, quoted: false });
      plain = '';
    }
    let s = ch;
    for (i++; i < sql.length; i++) {
      s += sql[i];
      if (sql[i] === ch) {
        if (sql[i + 1] === ch) {
          s += ch;
          i++;
        } else break;
      }
    }
    out.push({ text: s, quoted: true });
  }
  if (plain) out.push({ text: plain, quoted: false });
  return out;
}

/** The compiler's aliases: t for the root, j for joins, r for subqueries. */
const ALIAS_RE = /\b[tjr]\d+\b/g;

function escapeRe(s: string): string {
  return s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}

/**
 * A regex matching this definition's SQL wherever it reappears in a statement,
 * however the aliases were renumbered there — the compiler hands every
 * occurrence of the same column fresh ones. Each alias becomes a capture group
 * on first sight and a backreference after, so a match must rename them
 * consistently or not at all.
 */
function defPattern(defSql: string): RegExp {
  let src = '';
  const seen = new Map<string, number>();
  for (const seg of segments(defSql)) {
    if (seg.quoted) {
      src += escapeRe(seg.text);
      continue;
    }
    let last = 0;
    for (const m of seg.text.matchAll(ALIAS_RE)) {
      src += escapeRe(seg.text.slice(last, m.index));
      const group = seen.get(m[0]);
      if (group !== undefined) src += `\\${group}`;
      else {
        seen.set(m[0], seen.size + 1);
        src += String.raw`(\b[tjr]\d+\b)`;
      }
      last = m.index + m[0].length;
    }
    src += escapeRe(seg.text.slice(last));
  }
  return new RegExp(src, 'g');
}

function foldOnce(text: string, def: SqlDef): string {
  const folded = text.replace(defPattern(def.sql), varMark(def.name));
  // In the SELECT list the fold leaves `:X AS "X"`, which says its one thing twice.
  return folded.split(`${varMark(def.name)} AS ${q(def.name)}`).join(varMark(def.name));
}

/**
 * Replace the compiler's aliases with the names of the tables they stand for.
 * Where the alias would only repeat its table (`FROM "orders" orders`) it is
 * dropped — the table name is its own alias — and a second use of the same
 * name gets a numbered one, so a self-join or a rollup over its own table
 * cannot capture the outer row's name.
 *
 * `inherited` carries the enclosing statement's names into a definition's
 * body, but only the ones the body actually mentions claim their name there.
 */
function nameAliases(
  text: string,
  inherited: ReadonlyMap<string, string>,
): { text: string; names: Map<string, string> } {
  const segs = segments(text);
  const names = new Map<string, string>();
  const taken = new Set<string>();

  for (const seg of segs) {
    if (seg.quoted) continue;
    for (const m of seg.text.matchAll(ALIAS_RE)) {
      const outer = inherited.get(m[0]);
      if (outer !== undefined && !names.has(m[0])) {
        names.set(m[0], outer);
        taken.add(outer);
      }
    }
  }

  // Introductions in order: FROM or JOIN, the quoted table, then the alias.
  const intros = new Map<number, string>();
  for (let i = 0; i + 2 < segs.length; i++) {
    const before = segs[i]!;
    const table = segs[i + 1]!;
    const after = segs[i + 2]!;
    if (before.quoted || !/\b(?:FROM|JOIN)\s+$/.test(before.text)) continue;
    if (!table.quoted || !table.text.startsWith('"')) continue;
    const alias = /^ ([tjr]\d+)\b/.exec(after.quoted ? '' : after.text)?.[1];
    if (!alias || names.has(alias)) continue;

    const base = table.text.slice(1, -1).replaceAll('""', '"');
    let name = base;
    for (let n = 2; taken.has(name); n++) name = `${base}${n}`;
    names.set(alias, name);
    taken.add(name);
    if (name === base) intros.set(i + 2, alias);
  }

  let out = '';
  for (let i = 0; i < segs.length; i++) {
    const seg = segs[i]!;
    if (seg.quoted) {
      out += seg.text;
      continue;
    }
    let t = seg.text;
    const dropped = intros.get(i);
    if (dropped) t = t.replace(new RegExp(`^ ${dropped}\\b`), '');
    out += t.replace(ALIAS_RE, (alias) => names.get(alias) ?? alias);
  }
  return { text: out, names };
}

/** Clause keywords that start a new line when they occur at nesting depth zero. */
const CLAUSE_STARTS = ['FROM', 'LEFT JOIN', 'WHERE', 'AND', 'GROUP BY', 'ORDER BY', 'LIMIT'];

/**
 * For the main statement clauses break at depth zero, flush left; a
 * definition's body is one parenthesised subquery, so its clauses sit at
 * depth one and their lines indent under the `:name =` that owns them.
 */
function breakClauses(sql: string, atDepth = 0, indent = ''): string {
  let out = '';
  let depth = 0;
  // The main SELECT list gets a line per column, tab-indented; it runs from
  // the opening SELECT until the first clause break ends it.
  let inSelect = false;
  let eatSpace = false;
  for (const seg of segments(sql)) {
    if (seg.quoted) {
      out += seg.text;
      eatSpace = false;
      continue;
    }
    const t = seg.text;
    for (let i = 0; i < t.length; i++) {
      const ch = t[i]!;
      if (eatSpace) {
        eatSpace = false;
        if (ch === ' ') continue;
      }
      if (ch === '(') depth++;
      else if (ch === ')') depth = Math.max(0, depth - 1);
      else if (depth === atDepth) {
        if (ch === ' ') {
          if (atDepth === 0 && out.endsWith('SELECT')) {
            out += '\n\t';
            inSelect = true;
            continue;
          }
          const rest = t.slice(i + 1);
          if (CLAUSE_STARTS.some((k) => rest.startsWith(k + ' '))) {
            out += '\n' + indent;
            inSelect = false;
            continue;
          }
        } else if (ch === ',' && inSelect) {
          out += ',\n\t';
          eatSpace = true;
          continue;
        }
      }
      out += ch;
    }
  }
  return out;
}

/**
 * The tokens the SQL bar actually shows: computed columns folded back into
 * `:name` variables — which SQL itself does not have — each defined once above
 * a statement broken along its clauses, with the compiler's aliases renamed to
 * the tables they stand for. A faithful renaming of the statement that ran,
 * not a different query; COPY still copies the statement verbatim.
 */
export function displayTokens(displaySql: string, defs: SqlDef[]): SqlToken[] {
  // Longest first: an expression containing another column's expansion must
  // fold as itself before the smaller one can eat its inside.
  const ordered = [...defs].sort((a, b) => b.sql.length - a.sql.length);

  let text = displaySql;
  for (const def of ordered) text = foldOnce(text, def);
  const main = nameAliases(text, new Map());

  // Definitions fold into each other too, so a rollup built on another
  // computed column shows that relationship instead of repeating its body.
  const lines = defs.map((def) => {
    let body = def.sql;
    for (const other of ordered) {
      if (other.name !== def.name && other.sql.length < def.sql.length) {
        body = foldOnce(body, other);
      }
    }
    return `${varMark(def.name)} = ${breakClauses(nameAliases(body, main.names).text, 1, '\t')}`;
  });

  const full = (lines.length ? lines.join('\n') + '\n\n' : '') + breakClauses(main.text);

  // Chunks alternate statement text and variable names; names become `:name`.
  const tokens: SqlToken[] = [];
  const parts = full.split(VAR_MARK);
  for (let i = 0; i < parts.length; i++) {
    if (i % 2 === 1) tokens.push({ c: 'var', s: `:${parts[i]!}` });
    else if (parts[i]) tokens.push(...highlight(parts[i]!));
  }
  return tokens;
}

/** Split a statement into coloured tokens for the client's SQL bar. */
export function highlight(sql: string): SqlToken[] {
  const tokens: SqlToken[] = [];
  const push = (c: SqlToken['c'], s: string) => {
    const last = tokens[tokens.length - 1];
    if (last && last.c === c) last.s += s;
    else tokens.push({ c, s });
  };

  let lastIndex = 0;
  for (const m of sql.matchAll(TOKEN_RE)) {
    if (m.index > lastIndex) push('plain', sql.slice(lastIndex, m.index));
    lastIndex = m.index + m[0].length;

    const [, str, num, quoted, word, punct, space] = m;
    if (str !== undefined) push('str', str);
    else if (num !== undefined) push('num', num);
    else if (quoted !== undefined) push('id', quoted);
    else if (word !== undefined) {
      const upper = word.toUpperCase();
      if (SQL_KEYWORDS.has(upper)) push('kw', word);
      else if (FUNCTION_NAMES.has(upper)) push('fn', word);
      else push('id', word);
    } else if (punct !== undefined) push('op', punct);
    else if (space !== undefined) push('plain', space);
  }
  if (lastIndex < sql.length) push('plain', sql.slice(lastIndex));
  return tokens;
}