import type { ExtensionFunction, ExtensionIdentity } from '@/contracts/extension-api';
import type { Tool } from '@/contracts/tool';
import { TreeSitterCodeIndex } from './tree-sitter-provider';
import { createCodeIndexTools } from './tools';

export const CODE_INDEX_IDENTITY: ExtensionIdentity = {
  id: 'shard.builtin.code-index', name: 'Tree-sitter WASM CodeIndex', version: '0.1.0',
};

export function createCodeIndexExtension(options: { cwd: string } | string = process.cwd()): {
  identity: ExtensionIdentity; install: ExtensionFunction;
} {
  const cwd = typeof options === 'string' ? options : options.cwd;
  return {
    identity: CODE_INDEX_IDENTITY,
    async install(api) {
      const provider = new TreeSitterCodeIndex(cwd);
      await provider.initialize();
      api.registerCodeIntelligence(provider);
      for (const tool of createCodeIndexTools(provider)) api.registerTool(tool as Tool);
      return () => provider.dispose();
    },
  };
}

export default createCodeIndexExtension;