LexAI/packages/vscode/src/codeContext.ts

import * as vscode from 'vscode';
import * as path from 'node:path';

const MAX_DEFS = 8;
const MAX_IMPORT_FILES = 6;
const MAX_SNIPPET_CHARS = 2400;
const MAX_TOTAL_CONTEXT_CHARS = 14000;
const MAX_IDENTIFIERS = 24;
const SURROUND_LINES = 40;

const STOP_WORDS = new Set([
  'if', 'else', 'for', 'while', 'do', 'switch', 'case', 'break', 'return', 'const', 'let', 'var',
  'function', 'class', 'interface', 'type', 'enum', 'import', 'export', 'from', 'default', 'async',
  'await', 'try', 'catch', 'finally', 'throw', 'new', 'this', 'super', 'typeof', 'instanceof',
  'true', 'false', 'null', 'undefined', 'void', 'in', 'of', 'as', 'is', 'public', 'private',
  'protected', 'static', 'readonly', 'extends', 'implements', 'package', 'yield', 'with',
  'string', 'number', 'boolean', 'any', 'unknown', 'never', 'object', 'Record', 'Partial',
  'Promise', 'Array', 'Map', 'Set', 'Error', 'console', 'window', 'document', 'module',
  'require', 'exports', 'process', 'Buffer', 'self', 'global', 'Math', 'JSON', 'Date',
]);

export interface CodeContextBundle {
  /** Human-readable pack for the model user message */
  contextText: string;
  /** Short summary for UI status */
  summary: string;
  files: string[];
}

/**
 * Build workspace-aware context for a selection: surrounding code, import
 * targets, and definition-provider hits for identifiers in the selection.
 */
export async function gatherCodeContext(
  document: vscode.TextDocument,
  selection: vscode.Range,
): Promise<CodeContextBundle> {
  const selected = document.getText(selection);
  const wsFolder = vscode.workspace.getWorkspaceFolder(document.uri);
  const rel = (uri: vscode.Uri) =>
    wsFolder ? path.relative(wsFolder.uri.fsPath, uri.fsPath).replace(/\\/g, '/') : uri.fsPath;

  const parts: string[] = [];
  const files = new Set<string>();
  let budget = MAX_TOTAL_CONTEXT_CHARS;

  const push = (block: string) => {
    if (budget <= 0) return;
    const slice = block.length > budget ? block.slice(0, budget) + '\n…[truncated]' : block;
    parts.push(slice);
    budget -= slice.length;
  };

  const currentPath = rel(document.uri);
  files.add(currentPath);

  push(
    [
      '### Current file',
      `Path: ${currentPath}`,
      `Language: ${document.languageId}`,
      `Selection lines: ${selection.start.line + 1}–${selection.end.line + 1}`,
      '',
      '### Selected code',
      fence(document.languageId, selected),
    ].join('\n'),
  );

  const surroundStart = Math.max(0, selection.start.line - SURROUND_LINES);
  const surroundEnd = Math.min(document.lineCount - 1, selection.end.line + SURROUND_LINES);
  const surround = document.getText(new vscode.Range(surroundStart, 0, surroundEnd, document.lineAt(surroundEnd).text.length));
  push(
    [
      '',
      `### Surrounding code in ${currentPath} (lines ${surroundStart + 1}–${surroundEnd + 1})`,
      fence(document.languageId, surround),
    ].join('\n'),
  );

  // Import / require targets in the current file
  const importUris = await resolveImportUris(document);
  let importCount = 0;
  for (const uri of importUris) {
    if (importCount >= MAX_IMPORT_FILES || budget <= 0) break;
    if (uri.toString() === document.uri.toString()) continue;
    try {
      const doc = await vscode.workspace.openTextDocument(uri);
      const excerpt = excerptForSelection(doc, selected);
      const p = rel(uri);
      files.add(p);
      push(['', `### Imported module: ${p}`, fence(doc.languageId, excerpt)].join('\n'));
      importCount += 1;
    } catch {
      // ignore unresolved / binary
    }
  }

  // Definition provider for identifiers in the selection
  const idents = extractIdentifiers(selected);
  let defCount = 0;
  const seenDefKeys = new Set<string>();

  for (const ident of idents) {
    if (defCount >= MAX_DEFS || budget <= 0) break;
    const pos = findIdentifierPosition(document, selection, ident);
    if (!pos) continue;

    let locs: vscode.Location[] = [];
    try {
      const raw = await vscode.commands.executeCommand<
        vscode.Location | vscode.Location[] | vscode.LocationLink[] | undefined
      >('vscode.executeDefinitionProvider', document.uri, pos);
      locs = normalizeLocations(raw);
    } catch {
      continue;
    }

    for (const loc of locs) {
      if (defCount >= MAX_DEFS || budget <= 0) break;
      const key = `${loc.uri.toString()}:${loc.range.start.line}:${loc.range.start.character}`;
      if (seenDefKeys.has(key)) continue;
      seenDefKeys.add(key);
      if (loc.uri.toString() === document.uri.toString() && selection.contains(loc.range.start)) {
        continue;
      }
      try {
        const doc = await vscode.workspace.openTextDocument(loc.uri);
        const snippet = expandDefinitionSnippet(doc, loc.range);
        const p = rel(loc.uri);
        files.add(p);
        push(
          [
            '',
            `### Definition of \`${ident}\` — ${p}:${loc.range.start.line + 1}`,
            fence(doc.languageId, snippet),
          ].join('\n'),
        );
        defCount += 1;
      } catch {
        // skip
      }
    }
  }

  const summary =
    defCount || importCount
      ? `Included ${defCount} definition(s) and ${importCount} import file(s) from the workspace.`
      : 'No extra workspace definitions found (language support may be unavailable for this file).';

  return {
    contextText: parts.join('\n'),
    summary,
    files: [...files],
  };
}

function fence(lang: string, body: string): string {
  const safe = body.replace(/\r\n/g, '\n');
  const clipped =
    safe.length > MAX_SNIPPET_CHARS ? safe.slice(0, MAX_SNIPPET_CHARS) + '\n…[truncated]' : safe;
  return '```' + (lang || '') + '\n' + clipped + '\n```';
}

function extractIdentifiers(text: string): string[] {
  const matches = text.match(/\b[_A-Za-z][_A-Za-z0-9]*\b/g) ?? [];
  const out: string[] = [];
  const seen = new Set<string>();
  for (const m of matches) {
    if (STOP_WORDS.has(m) || m.length < 2) continue;
    if (seen.has(m)) continue;
    seen.add(m);
    out.push(m);
    if (out.length >= MAX_IDENTIFIERS) break;
  }
  return out;
}

function findIdentifierPosition(
  document: vscode.TextDocument,
  selection: vscode.Range,
  ident: string,
): vscode.Position | undefined {
  const text = document.getText(selection);
  const idx = text.indexOf(ident);
  if (idx < 0) return undefined;
  const startOffset = document.offsetAt(selection.start) + idx;
  return document.positionAt(startOffset);
}

function normalizeLocations(
  raw: vscode.Location | vscode.Location[] | vscode.LocationLink[] | undefined,
): vscode.Location[] {
  if (!raw) return [];
  const arr = Array.isArray(raw) ? raw : [raw];
  return arr.map((item) => {
    if (item instanceof vscode.Location) return item;
    const link = item as vscode.LocationLink;
    return new vscode.Location(link.targetUri, link.targetSelectionRange ?? link.targetRange);
  });
}

async function resolveImportUris(document: vscode.TextDocument): Promise<vscode.Uri[]> {
  const text = document.getText();
  const specs = new Set<string>();

  // ES / TS imports
  const reFrom = /from\s+['"]([^'"]+)['"]/g;
  const reImport = /import\s+['"]([^'"]+)['"]/g;
  const reRequire = /require\s*\(\s*['"]([^'"]+)['"]\s*\)/g;
  for (const re of [reFrom, reImport, reRequire]) {
    let m: RegExpExecArray | null;
    while ((m = re.exec(text))) {
      const spec = m[1];
      if (spec.startsWith('.') || spec.startsWith('/')) specs.add(spec);
    }
  }

  const uris: vscode.Uri[] = [];
  for (const spec of specs) {
    const resolved = await resolveRelativeModule(document.uri, spec);
    if (resolved) uris.push(resolved);
  }
  return uris;
}

async function resolveRelativeModule(
  from: vscode.Uri,
  spec: string,
): Promise<vscode.Uri | undefined> {
  const baseDir = path.posix.dirname(from.path);
  const joined = path.posix.normalize(path.posix.join(baseDir, spec));
  const candidates = [
    joined,
    `${joined}.ts`,
    `${joined}.tsx`,
    `${joined}.js`,
    `${joined}.jsx`,
    `${joined}.mjs`,
    `${joined}.cjs`,
    `${joined}.json`,
    `${joined}/index.ts`,
    `${joined}/index.tsx`,
    `${joined}/index.js`,
  ];

  for (const p of candidates) {
    const uri = from.with({ path: p });
    try {
      await vscode.workspace.fs.stat(uri);
      return uri;
    } catch {
      // try next
    }
  }
  return undefined;
}

/** Prefer exporting / matching snippets from an imported file. */
function excerptForSelection(doc: vscode.TextDocument, selected: string): string {
  const idents = extractIdentifiers(selected);
  const full = doc.getText();
  if (!idents.length) {
    return full.slice(0, MAX_SNIPPET_CHARS);
  }

  const chunks: string[] = [];
  for (const id of idents.slice(0, 10)) {
    const patterns = [
      new RegExp(
        `(?:export\\s+)?(?:async\\s+)?function\\s+${id}\\b[\\s\\S]{0,800}?\\n\\}`,
        'm',
      ),
      new RegExp(
        `(?:export\\s+)?(?:const|let|var|class|interface|type|enum)\\s+${id}\\b[\\s\\S]{0,600}`,
        'm',
      ),
    ];
    for (const re of patterns) {
      const m = full.match(re);
      if (m) {
        chunks.push(m[0]);
        break;
      }
    }
  }

  if (!chunks.length) {
    // Fall back to first N lines (often exports barrel / header)
    return full.split('\n').slice(0, 80).join('\n');
  }
  return chunks.join('\n\n');
}

function expandDefinitionSnippet(doc: vscode.TextDocument, range: vscode.Range): string {
  const start = Math.max(0, range.start.line - 2);
  let end = Math.min(doc.lineCount - 1, range.end.line);
  const maxEnd = Math.min(doc.lineCount - 1, range.start.line + 80);
  while (end < maxEnd) {
    const line = doc.lineAt(end).text;
    if (line.includes('}') || line.includes(';')) break;
    end += 1;
  }
  return doc.getText(new vscode.Range(start, 0, end, doc.lineAt(end).text.length));
}