shrimp/vscode-extension/server/src/server.ts

import { TextDocument } from 'vscode-languageserver-textdocument'
import { buildDiagnostics } from './diagnostics'
import { buildSemanticTokens, TOKEN_MODIFIERS, TOKEN_TYPES } from './semanticTokens'
import { provideCompletions } from './completion/completionProvider'
import { provideSignatureHelp } from './signatureHelp'
import { PRELUDE_NAMES } from './metadata/prelude-names'
import { parser } from '../../../src/parser/shrimp'
import { setGlobals } from '../../../src/parser/tokenizer'
import { Compiler } from '../../../src/compiler/compiler'
import { Tree } from '@lezer/common'
import {
  InitializeResult,
  TextDocuments,
  TextDocumentSyncKind,
  createConnection,
  ProposedFeatures,
  CompletionItemKind,
  TextDocumentChangeEvent,
} from 'vscode-languageserver/node'
import { setGlobals } from '../../../src/parser/tokenizer'
import { globals } from '../../../src/prelude'

// Initialize parser with prelude globals so it knows dict/list/str are in scope
setGlobals(PRELUDE_NAMES)

const connection = createConnection(ProposedFeatures.all)
const documents = new TextDocuments(TextDocument)
documents.listen(connection)

const documentTrees = new Map<string, Tree>()

// Server capabilities
connection.onInitialize(handleInitialize)

// Language features
connection.languages.semanticTokens.on(handleSemanticTokens)
documents.onDidOpen(handleDocumentOpen)
documents.onDidChangeContent(handleDocumentChange)
documents.onDidClose(handleDocumentClose)
connection.onCompletion(handleCompletion)
connection.onSignatureHelp(handleSignatureHelp)

// Debug commands
connection.onRequest('shrimp/parseTree', handleParseTree)
connection.onRequest('shrimp/bytecode', handleBytecode)

// Start listening
connection.listen()

// Handler implementations
function handleInitialize(): InitializeResult {
  connection.console.log('🦐 Server initialized with capabilities')
  const result: InitializeResult = {
    capabilities: {
      textDocumentSync: TextDocumentSyncKind.Full,
      completionProvider: {
        triggerCharacters: ['.'],
      },
      signatureHelpProvider: {
        triggerCharacters: [' '],
      },
      semanticTokensProvider: {
        legend: {
          tokenTypes: TOKEN_TYPES,
          tokenModifiers: TOKEN_MODIFIERS,
        },
        full: true,
      },
    },
  }

  return result
}

function handleDocumentOpen(event: TextDocumentChangeEvent<TextDocument>) {
  const document = event.document
  setGlobals(Object.keys(globals))
  const tree = parser.parse(document.getText())
  documentTrees.set(document.uri, tree)
}

function handleSemanticTokens(params: any) {
  const document = documents.get(params.textDocument.uri)
  if (!document) return { data: [] }

  const tree = documentTrees.get(params.textDocument.uri)
  if (!tree) return { data: [] }

  const data = buildSemanticTokens(document, tree)
  return { data }
}

function handleDocumentChange(change: TextDocumentChangeEvent<TextDocument>) {
  const document = change.document

  // Parse and cache
  setGlobals(Object.keys(globals))
  const tree = parser.parse(document.getText())
  documentTrees.set(document.uri, tree)

  // Build diagnostics using cached tree
  const diagnostics = buildDiagnostics(document, tree)
  connection.sendDiagnostics({ uri: document.uri, diagnostics })
}

function handleDocumentClose(event: TextDocumentChangeEvent<TextDocument>) {
  documentTrees.delete(event.document.uri)
}

function handleCompletion(params: any) {
  const document = documents.get(params.textDocument.uri)
  if (!document) {
    console.log('❌ No document found')
    return []
  }

  const position = params.position
  const text = document.getText()
  const offset = document.offsetAt(position)
  console.log(`📍 Text around cursor: "${text.slice(Math.max(0, offset - 10), offset + 10)}"`)

  // First try context-aware completions (module/dollar)
  const contextCompletions = provideCompletions(document, position)
  console.log(`🎯 Context completions count: ${contextCompletions.length}`)
  if (contextCompletions.length > 0) {
    console.log(
      `✅ Returning ${contextCompletions.length} completions:`,
      contextCompletions.map((c) => c.label).join(', ')
    )
    return contextCompletions
  }

  // Fall back to keywords + prelude globals (for Ctrl+Space in general context)
  console.log(`⌨️  Falling back to keywords + prelude globals`)
  const keywords = ['if', 'else', 'do', 'end', 'and', 'or', 'true', 'false', 'null']
  const keywordCompletions = keywords.map((keyword) => ({
    label: keyword,
    kind: CompletionItemKind.Keyword,
  }))

  const preludeCompletions = PRELUDE_NAMES.map((name) => ({
    label: name,
    kind: CompletionItemKind.Function,
  }))

  return [...keywordCompletions, ...preludeCompletions]
}

function handleSignatureHelp(params: any) {
  const document = documents.get(params.textDocument.uri)
  if (!document) return
  return provideSignatureHelp(document, params.position)
}

function handleParseTree(params: { uri: string }) {
  connection.console.log(`🦐 Parse tree requested for: ${params.uri}`)
  const document = documents.get(params.uri)
  if (!document) return 'Document not found'

  const tree = documentTrees.get(params.uri)
  if (!tree) {
    connection.console.error(`🦐 No cached tree for ${params.uri}`)
    return 'No cached parse tree available'
  }

  const text = document.getText()
  const cursor = tree.cursor()

  let formatted = ''
  let depth = 0

  const printNode = () => {
    const nodeName = cursor.name
    const nodeText = text.slice(cursor.from, cursor.to)
    const indent = '  '.repeat(depth)

    formatted += `${indent}${nodeName}`
    if (nodeText) {
      const escapedText = nodeText.replace(/\n/g, '\\n').replace(/\r/g, '\\r')
      formatted += ` "${escapedText}"`
    }
    formatted += '\n'
  }

  const traverse = (): void => {
    printNode()

    if (cursor.firstChild()) {
      depth++
      do {
        traverse()
      } while (cursor.nextSibling())
      cursor.parent()
      depth--
    }
  }

  traverse()
  return formatted
}

function handleBytecode(params: { uri: string }) {
  connection.console.log(`🦐 Bytecode requested for: ${params.uri}`)
  const document = documents.get(params.uri)
  if (!document) return 'Document not found'

  try {
    const text = document.getText()
    const compiler = new Compiler(text)

    // Format bytecode as readable string
    let output = 'Bytecode:\n\n'
    const bytecode = compiler.bytecode

    output += bytecode.instructions
      .map((op, i) => `${i.toString().padStart(4)}: ${JSON.stringify(op)}`)
      .join('\n')

    // Strip ANSI color codes
    output = output.replace(/\x1b\[[0-9;]*m/g, '')

    return output
  } catch (error) {
    const errorMsg = error instanceof Error ? error.message : String(error)
    // Strip ANSI color codes from error message too
    return `Compilation failed: ${errorMsg.replace(/\x1b\[[0-9;]*m/g, '')}`
  }
}