complete.ts

import {
  parse,
  parseFromClause,
  SelectStatement,
  FromTableNode,
  IncompleteSubqueryNode,
  FromClauseParserResult,
  DeleteStatement,
  ParseError,
  ExpectedLiteralNode,
  AST,
  AlterTableStatement,
  DropTableStatement,
} from '@deepnote/sql-parser'
import { CompletionItem } from 'vscode-languageserver-types'
import { Schema, Table } from '../database_libs/AbstractClient'
import { stubLogger } from '../logger'
import { getRidOfAfterPosString } from './StringUtils'
import { getLastToken } from './utils/getLastToken'
import {
  isPosInLocation,
  createTablesFromFromNodes,
  findColumnAtPosition,
  getAllNestedFromNodes,
  getNearestFromTableFromPos,
} from './AstUtils'
import { createBasicKeywordCandidates } from './candidates/createBasicKeywordCandidates'
import { createCatalogDatabaseAndTableCandidates } from './candidates/createTableCandidates'
import { createJoinCondidates } from './candidates/createJoinCandidates'
import {
  createCandidatesForColumnsOfAnyTable,
  createCandidatesForScopedColumns,
  createCandidatesForUnscopedColumns,
} from './candidates/createColumnCandidates'
import { createAliasCandidates } from './candidates/createAliasCandidates'
import { createSelectAllColumnsCandidates } from './candidates/createSelectAllColumnsCandidates'
import { createFunctionCandidates } from './candidates/createFunctionCandidates'
import { createKeywordCandidatesFromExpectedLiterals } from './candidates/createKeywordCandidatesFromExpectedLiterals'
import { createJoinTablesCandidates } from './candidates/createJoinTableCndidates'
import { ICONS, toCompletionItemForKeyword } from './CompletionItemUtils'

export type Pos = { line: number; column: number }

const logger = stubLogger()

function getFromNodesFromClause(sql: string): FromClauseParserResult | null {
  try {
    return parseFromClause(sql)
  } catch (_e) {
    // no-op
    return null
  }
}

type CompletionError = {
  label: string
  detail: string
  line: number
  offset: number
}
class Completer {
  lastToken = ''
  candidates: CompletionItem[] = []
  schema: Schema
  error: CompletionError | null = null
  sql: string
  pos: Pos
  isSpaceTriggerCharacter = false
  isDotTriggerCharacter = false
  jupyterLabMode: boolean

  constructor(schema: Schema, sql: string, pos: Pos, jupyterLabMode: boolean) {
    this.schema = schema
    this.sql = sql
    this.pos = pos
    this.jupyterLabMode = jupyterLabMode
  }

  complete() {
    const target = getRidOfAfterPosString(this.sql, this.pos)
    logger.debug(`target: ${target}`)
    this.lastToken = getLastToken(target)
    const idx = this.lastToken.lastIndexOf('.')
    this.isSpaceTriggerCharacter = this.lastToken === ''
    this.isDotTriggerCharacter =
      !this.isSpaceTriggerCharacter && idx == this.lastToken.length - 1

    try {
      const ast = parse(target)
      this.addCandidatesForParsedStatement(ast)
    } catch (_e: unknown) {
      logger.debug('error')
      logger.debug(_e)
      if (!(_e instanceof Error)) {
        throw _e
      }
      if (_e.name !== 'SyntaxError') {
        throw _e
      }
      const e = _e as ParseError
      const parsedFromClause = getFromNodesFromClause(this.sql)
      if (parsedFromClause) {
        const fromNodes = getAllNestedFromNodes(
          parsedFromClause?.from?.tables || []
        )
        const fromNodeOnCursor = getNearestFromTableFromPos(fromNodes, this.pos)
        if (
          fromNodeOnCursor &&
          fromNodeOnCursor.type === 'incomplete_subquery'
        ) {
          // Incomplete sub query 'SELECT sub FROM (SELECT e. FROM employees e) sub'
          this.addCandidatesForIncompleteSubquery(fromNodeOnCursor)
        } else {
          this.addCandidatesForSelectQuery(e, fromNodes, parsedFromClause)
          const expectedLiteralNodes =
            e.expected?.filter(
              (v): v is ExpectedLiteralNode => v.type === 'literal'
            ) || []
          this.addCandidatesForJoins(expectedLiteralNodes, fromNodes)
        }
      } else if (e.message === 'EXPECTED COLUMN NAME') {
        this.addCandidatesForInsert()
      } else {
        this.addCandidatesForError(e)
      }
      this.error = {
        label: e.name,
        detail: e.message,
        line: e.location.start.line,
        offset: e.location.start.offset,
      }
    }
    return this.candidates
  }

  addCandidatesForBasicKeyword() {
    createBasicKeywordCandidates().forEach((v) => {
      this.addCandidate(v)
    })
  }

  addCandidatesForExpectedLiterals(expected: ExpectedLiteralNode[]) {
    createKeywordCandidatesFromExpectedLiterals(expected).forEach((v) => {
      this.addCandidate(v)
    })
  }

  addCandidate(item: CompletionItem) {
    // A keyword completion can be occured anyplace and need to suppress them.
    if (
      item.kind &&
      item.kind === ICONS.KEYWORD &&
      !item.label.startsWith(this.lastToken)
    ) {
      return
    }
    // JupyterLab requires the dot or space character preceeding the <tab> key pressed
    // If the dot or space character are not added to the label then searching
    // in the list of suggestion does not work.
    // Here we fix this issue by adding the dot or space character
    // to the filterText and insertText.
    // TODO: report this issue to JupyterLab-LSP project.
    if (this.jupyterLabMode) {
      const text = item.insertText || item.label
      if (this.isSpaceTriggerCharacter) {
        item.insertText = ' ' + text
        item.filterText = ' ' + text
      } else if (this.isDotTriggerCharacter) {
        item.insertText = '.' + text
        item.filterText = '.' + text
      }
    }
    this.candidates.push(item)
  }

  addCandidatesForTables(tables: Table[], onFromClause: boolean) {
    createCatalogDatabaseAndTableCandidates(
      tables,
      this.lastToken,
      onFromClause
    ).forEach((item) => {
      this.addCandidate(item)
    })
  }

  addCandidatesForColumnsOfAnyTable(tables: Table[]) {
    createCandidatesForColumnsOfAnyTable(tables, this.lastToken).forEach(
      (item) => {
        this.addCandidate(item)
      }
    )
  }

  addCandidatesForIncompleteSubquery(
    incompleteSubquery: IncompleteSubqueryNode
  ) {
    const parsedFromClause = getFromNodesFromClause(incompleteSubquery.text)
    try {
      parse(incompleteSubquery.text)
    } catch (e: unknown) {
      if (!(e instanceof Error)) {
        throw e
      }
      if (e.name !== 'SyntaxError') {
        throw e
      }
      const fromText = incompleteSubquery.text
      const newPos = parsedFromClause
        ? {
            line: this.pos.line - (incompleteSubquery.location.start.line - 1),
            column:
              this.pos.column - incompleteSubquery.location.start.column + 1,
          }
        : { line: 0, column: 0 }
      const completer = new Completer(
        this.schema,
        fromText,
        newPos,
        this.jupyterLabMode
      )
      completer.complete().forEach((item) => this.addCandidate(item))
    }
  }

  /**
   * INSERT INTO TABLE1 (C
   */
  addCandidatesForInsert() {
    this.addCandidatesForColumnsOfAnyTable(this.schema.tables)
  }

  addCandidatesForError(e: ParseError) {
    const expectedLiteralNodes =
      e.expected?.filter(
        (v): v is ExpectedLiteralNode => v.type === 'literal'
      ) || []
    this.addCandidatesForExpectedLiterals(expectedLiteralNodes)
    this.addCandidatesForFunctions()
    this.addCandidatesForTables(this.schema.tables, false)
  }

  addCandidatesForSelectQuery(
    e: ParseError,
    fromNodes: FromTableNode[],
    parsedFromClause: FromClauseParserResult
  ) {
    const subqueryTables = createTablesFromFromNodes(fromNodes)
    const schemaAndSubqueries = this.schema.tables.concat(subqueryTables)
    this.addCandidatesForSelectStar(fromNodes, schemaAndSubqueries)
    const expectedLiteralNodes =
      e.expected?.filter(
        (v): v is ExpectedLiteralNode =>
          v.type === 'literal' && hasAtLeastTwoLetters(v.text)
      ) || []
    this.addCandidatesForExpectedLiterals(expectedLiteralNodes)
    this.addCandidatesForFunctions()

    // Detect FROM clause context BEFORE adding column suggestions
    const fromNodesContainingCursor = fromNodes.filter((tableNode) =>
      isPosInLocation(tableNode.location, this.pos)
    )
    const isCursorInsideFromClause = fromNodesContainingCursor.length > 0

    // Check if cursor is right after FROM keyword or typing a table name
    const afterFromClause = parsedFromClause.after?.trim().toUpperCase() || ''
    const isCursorAfterFromKeyword =
      afterFromClause === 'FROM' ||
      afterFromClause.startsWith('FROM ') ||
      /^(INNER |LEFT |RIGHT |FULL |FULL OUTER |CROSS |NATURAL |OUTER )?JOIN( |$)/.test(
        afterFromClause
      )

    const isTypingTableName =
      isCursorInsideFromClause || isCursorAfterFromKeyword

    if (!isTypingTableName) {
      const { addedSome: addedSomeScopedColumnCandidates } =
        this.addCandidatesForScopedColumns(fromNodes, schemaAndSubqueries)

      if (!addedSomeScopedColumnCandidates) {
        this.addCandidatesForUnscopedColumns(fromNodes, schemaAndSubqueries)
      }
    }

    this.addCandidatesForAliases(fromNodes)

    if (isTypingTableName) {
      // add table candidates if the cursor is inside a FROM clause, JOIN clause,
      // or right after FROM/JOIN keyword waiting for a table name
      this.addCandidatesForTables(schemaAndSubqueries, true)
    }

    if (logger.isDebugEnabled())
      logger.debug(
        `candidates for error returns: ${JSON.stringify(this.candidates)}`
      )
  }

  addCandidatesForJoins(
    expected: ExpectedLiteralNode[],
    fromNodes: FromTableNode[]
  ) {
    createJoinTablesCandidates(
      this.schema.tables,
      expected,
      fromNodes,
      this.lastToken
    ).forEach((v) => {
      this.addCandidate(v)
    })
  }

  addCandidatesForParsedDeleteStatement(ast: DeleteStatement) {
    if (isPosInLocation(ast.table.location, this.pos)) {
      this.addCandidatesForTables(this.schema.tables, false)
    } else if (
      ast.where &&
      isPosInLocation(ast.where.expression.location, this.pos)
    ) {
      const expr = ast.where.expression
      if (expr.type === 'column_ref') {
        this.addCandidatesForColumnsOfAnyTable(this.schema.tables)
      }
    }
  }

  addCandidatesForParsedDropStatement(ast: DropTableStatement) {
    if (isPosInLocation(ast.table.location, this.pos)) {
      this.addCandidatesForTables(this.schema.tables, false)
    }
  }

  addCandidatesForParsedAlterTableStatement(ast: AlterTableStatement) {
    if (ast.command.type === 'alter_table_drop_column') {
      if (isPosInLocation(ast.command.column.location, this.pos)) {
        const table = this.schema.tables.find((v) => v.tableName === ast.table)
        this.addCandidatesForColumnsOfAnyTable(
          table ? [table] : this.schema.tables
        )
      }
    }
  }

  addCandidatesForParsedSelectQuery(ast: SelectStatement) {
    this.addCandidatesForBasicKeyword()
    if (Array.isArray(ast.columns)) {
      this.addCandidate(toCompletionItemForKeyword('FROM'))
      this.addCandidate(toCompletionItemForKeyword('AS'))
    }
    if (!ast.distinct) {
      this.addCandidate(toCompletionItemForKeyword('DISTINCT'))
    }

    // Check if cursor is inside a FROM clause table reference
    // This handles the case where "SELECT * FROM a" parses successfully
    // but we still want to suggest tables starting with "a"
    const parsedFromClause = getFromNodesFromClause(this.sql)
    const fromNodes = getAllNestedFromNodes(
      parsedFromClause?.from?.tables || []
    )
    const subqueryTables = createTablesFromFromNodes(fromNodes)
    const schemaAndSubqueries = this.schema.tables.concat(subqueryTables)

    for (const tableNode of fromNodes) {
      if (tableNode.type === 'table') {
        // Check if the lastToken matches the table name (user is typing the table name)
        // This means the cursor is ON the table name, not after it (like typing an alias)
        const tableNameMatches =
          this.lastToken.length > 0 &&
          tableNode.table.toLowerCase().startsWith(this.lastToken.toLowerCase())

        if (tableNameMatches && isPosInLocation(tableNode.location, this.pos)) {
          // Cursor is typing a table name - suggest tables
          this.addCandidatesForTables(schemaAndSubqueries, true)
          if (logger.isDebugEnabled())
            logger.debug(
              `parse query returns: ${JSON.stringify(this.candidates)}`
            )
          return
        }
      }
    }

    const columnRef = findColumnAtPosition(ast, this.pos)
    if (!columnRef) {
      this.addJoinCondidates(ast)
    } else {
      if (columnRef.table) {
        // We know what table/alias this column belongs to
        // Find the corresponding table and suggest it's columns
        this.addCandidatesForScopedColumns(fromNodes, schemaAndSubqueries)
      } else {
        // Column is not scoped to a table/alias yet
        // Could be an alias or an unscoped column
        this.addCandidatesForUnscopedColumns(fromNodes, schemaAndSubqueries)
        this.addCandidatesForAliases(fromNodes)
        this.addCandidatesForFunctions()
      }
    }
    if (logger.isDebugEnabled())
      logger.debug(`parse query returns: ${JSON.stringify(this.candidates)}`)
  }

  addCandidatesForParsedStatement(ast: AST) {
    if (logger.isDebugEnabled())
      logger.debug(
        `getting candidates for parse query ast: ${JSON.stringify(ast)}`
      )
    if (!ast.type) {
      this.addCandidatesForBasicKeyword()
    } else if (ast.type === 'delete') {
      this.addCandidatesForParsedDeleteStatement(ast)
    } else if (ast.type === 'select') {
      this.addCandidatesForParsedSelectQuery(ast)
    } else if (ast.type === 'alter_table') {
      this.addCandidatesForParsedAlterTableStatement(ast)
    } else if (ast.type === 'drop_table') {
      this.addCandidatesForParsedDropStatement(ast)
    } else {
      console.log(`AST type not supported yet: ${ast.type}`)
    }
  }

  addJoinCondidates(ast: SelectStatement) {
    createJoinCondidates(
      ast,
      this.schema.tables,
      this.pos,
      this.lastToken
    ).forEach((v) => {
      this.addCandidate(v)
    })
  }

  addCandidatesForFunctions() {
    console.time('addCandidatesForFunctions')
    createFunctionCandidates(this.schema.functions, this.lastToken).forEach(
      (v) => {
        this.addCandidate(v)
      }
    )
    console.timeEnd('addCandidatesForFunctions')
  }

  addCandidatesForSelectStar(fromNodes: FromTableNode[], tables: Table[]) {
    console.time('addCandidatesForSelectStar')
    createSelectAllColumnsCandidates(fromNodes, tables, this.lastToken).forEach(
      (v) => {
        this.addCandidate(v)
      }
    )
    console.timeEnd('addCandidatesForSelectStar')
  }

  addCandidatesForScopedColumns(
    fromNodes: FromTableNode[],
    tables: Table[]
  ): { addedSome: boolean } {
    console.time('addCandidatesForScopedColumns')
    let addedSome = false
    createCandidatesForScopedColumns(fromNodes, tables, this.lastToken).forEach(
      (v) => {
        addedSome = true
        this.addCandidate(v)
      }
    )
    console.timeEnd('addCandidatesForScopedColumns')
    return { addedSome }
  }

  addCandidatesForUnscopedColumns(fromNodes: FromTableNode[], tables: Table[]) {
    createCandidatesForUnscopedColumns(
      fromNodes,
      tables,
      this.lastToken
    ).forEach((v) => {
      this.addCandidate(v)
    })
  }

  addCandidatesForAliases(fromNodes: FromTableNode[]) {
    createAliasCandidates(fromNodes, this.lastToken).forEach((v) => {
      this.addCandidate(v)
    })
  }
}

export function complete(
  sql: string,
  pos: Pos,
  schema: Schema = { tables: [], functions: [] },
  jupyterLabMode = false
) {
  console.time('complete')
  if (logger.isDebugEnabled())
    logger.debug(`complete: ${sql}, ${JSON.stringify(pos)}`)
  const completer = new Completer(schema, sql, pos, jupyterLabMode)
  const candidates = completer.complete()
  console.timeEnd('complete')
  return { candidates: candidates, error: completer.error }
}

function hasAtLeastTwoLetters(value: string): boolean {
  return /[a-zA-Z].*[a-zA-Z]/.test(value)
}