dbx/apps/desktop/src/lib/sql/executableStatementRangeCache.ts

import type { ChangeSet, Text } from "@codemirror/state";
import type { DatabaseType } from "@/types/database";
import { readSqlBracedParameterAt, type SqlParameterOptions } from "@/lib/sql/sqlParameters";
import { executableStatementRanges, type SqlTextRange } from "@/lib/sql/sqlStatementRanges";
import { cursorBelongsToTrailingStatementDelimiter } from "@/lib/sql/statementDelimiter";

export interface ExecutableStatementRangeCache {
  doc: Text;
  databaseType?: DatabaseType;
  parameterOptions?: SqlParameterOptions;
  parameterSyntaxKey: string;
  byStart: Map<number, SqlTextRange>;
  byExecutableLineStart: Map<number, SqlTextRange>;
  ranges: SqlTextRange[];
}

/**
 * Membership-only view of the statement-start positions a run-statement gutter
 * needs (`byStart` keys plus `byExecutableLineStart` keys). Kept separate from
 * the full cache so it can be cheaply shifted through a ChangeSet on every
 * keystroke and fully rebuilt only after typing pauses, instead of re-parsing
 * the whole document synchronously per keystroke.
 */
export interface StatementGutterStartIndex {
  starts: ReadonlySet<number>;
  executableLineStarts: ReadonlySet<number>;
}

export function statementGutterStartIndexForCache(cache: ExecutableStatementRangeCache): StatementGutterStartIndex {
  return { starts: new Set(cache.byStart.keys()), executableLineStarts: new Set(cache.byExecutableLineStart.keys()) };
}

export function mapStatementGutterStartIndex(index: StatementGutterStartIndex, changes: ChangeSet): StatementGutterStartIndex {
  return { starts: mapStartPositions(index.starts, changes), executableLineStarts: mapStartPositions(index.executableLineStarts, changes) };
}

function mapStartPositions(positions: ReadonlySet<number>, changes: ChangeSet): Set<number> {
  const mapped = new Set<number>();
  for (const position of positions) mapped.add(changes.mapPos(position, 1));
  return mapped;
}

export function statementGutterStartIndexHasStartAt(index: StatementGutterStartIndex, lineFrom: number): boolean {
  return index.starts.has(lineFrom) || index.executableLineStarts.has(lineFrom);
}

export type ExecutableStatementRangeParser = (sql: string, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions) => SqlTextRange[];

export function executableStatementRangeCacheForDoc(
  cache: ExecutableStatementRangeCache | null,
  doc: Text,
  databaseType?: DatabaseType,
  parameterOptionsOrParse?: SqlParameterOptions | ExecutableStatementRangeParser,
  customParse: ExecutableStatementRangeParser = executableStatementRanges,
): ExecutableStatementRangeCache {
  const parameterOptions = typeof parameterOptionsOrParse === "function" ? undefined : parameterOptionsOrParse;
  const parse = typeof parameterOptionsOrParse === "function" ? parameterOptionsOrParse : customParse;
  const parameterSyntaxKey = `${parameterOptions?.enabledSyntaxes ? parameterOptions.enabledSyntaxes.join(",") : "*"}|compat=${parameterOptions?.compatibilityMode?.trim().toUpperCase() ?? ""}`;
  if (cache?.doc === doc && cache.databaseType === databaseType && cache.parameterSyntaxKey === parameterSyntaxKey) return cache;

  const byStart = new Map<number, SqlTextRange>();
  const byExecutableLineStart = new Map<number, SqlTextRange>();
  const sql = doc.toString();
  const ranges = parse(sql, databaseType, parameterOptions);
  for (const range of ranges) {
    byStart.set(range.from, range);
    const line = doc.lineAt(range.from);
    if (doc.sliceString(line.from, range.from).trim() === "") {
      byExecutableLineStart.set(line.from, range);
    }
    const executableStart = executableStartAfterLeadingDirective(sql, range, databaseType, parameterOptions);
    if (executableStart !== null) {
      byExecutableLineStart.set(doc.lineAt(executableStart).from, range);
    }
  }
  return { doc, databaseType, parameterOptions, parameterSyntaxKey, byStart, byExecutableLineStart, ranges };
}

function executableStartAfterLeadingDirective(sql: string, range: SqlTextRange, databaseType?: DatabaseType, parameterOptions?: SqlParameterOptions): number | null {
  const text = range.sql;
  const hasExecutableDirective = text.startsWith("/*+") || text.startsWith("/*@") || text.startsWith("/*&") || (databaseType === "mysql" && text.startsWith("/*proxy*/"));
  if (!hasExecutableDirective) return null;

  let offset = 0;
  while (offset < text.length) {
    while (/\s/.test(text[offset] ?? "")) offset += 1;

    if (text.startsWith("/*", offset)) {
      const close = text.indexOf("*/", offset + 2);
      if (close < 0) return null;
      offset = close + 2;
      continue;
    }

    if (text.startsWith("--", offset)) {
      const newline = text.indexOf("\n", offset + 2);
      if (newline < 0) return null;
      offset = newline + 1;
      continue;
    }

    if (databaseType !== "sqlserver" && text[offset] === "#" && readSqlBracedParameterAt(sql, range.from + offset, parameterOptions)?.syntax !== "mybatis") {
      const newline = text.indexOf("\n", offset + 1);
      if (newline < 0) return null;
      offset = newline + 1;
      continue;
    }

    return offset < text.length ? range.from + offset : null;
  }

  return null;
}

export function executableStatementRangeStartingAt(cache: ExecutableStatementRangeCache, lineFrom: number): SqlTextRange | null {
  return cache.byStart.get(lineFrom) ?? cache.byExecutableLineStart.get(lineFrom) ?? null;
}

export function executableStatementRangeAtCursor(cache: ExecutableStatementRangeCache, cursorPos: number): SqlTextRange | null {
  const pos = Math.max(0, Math.min(cursorPos, cache.doc.length));
  const line = cache.doc.lineAt(pos);
  const lineText = line.text.trim();
  const lineContentStart = line.from + line.text.search(/\S|$/);
  const startsHashComment = lineText.startsWith("#") && readSqlBracedParameterAt(cache.doc.toString(), lineContentStart, cache.parameterOptions)?.syntax !== "mybatis";
  if (!lineText || lineText.startsWith("--") || startsHashComment || isCursorOnLeadingBlockComment(line.text, pos - line.from)) return null;

  for (let index = 0; index < cache.ranges.length; index += 1) {
    const range = cache.ranges[index];
    if (pos >= range.from && pos <= range.to) return range;

    if (pos < range.from && range.from <= line.to && cache.doc.sliceString(pos, range.from).trim() === "") {
      return range;
    }

    const next = cache.ranges[index + 1];
    if (pos > range.to && (!next || pos < next.from)) {
      if (cursorBelongsToTrailingStatementDelimiter(cache.doc, range.to, pos)) return range;
      if (isCursorOnRangeEndLine(cache.doc, pos, range.to)) return range;
    }
  }

  return null;
}

function isCursorOnLeadingBlockComment(lineText: string, lineOffset: number): boolean {
  const commentStart = lineText.search(/\S/);
  if (commentStart < 0 || !lineText.startsWith("/*", commentStart)) return false;

  const commentEnd = lineText.indexOf("*/", commentStart + 2);
  if (commentEnd < 0) return true;
  const afterComment = lineText.slice(commentEnd + 2);
  if (!afterComment.trim()) return true;

  return lineOffset <= commentEnd + 2;
}

function isCursorOnRangeEndLine(doc: Text, pos: number, rangeTo: number): boolean {
  const line = doc.lineAt(pos);
  return rangeTo >= line.from && rangeTo <= line.to;
}