import Parser from 'tree-sitter'
import TypeScript from 'tree-sitter-typescript'

interface CodeChunk {
  content: string
  metadata: {
    language: string
    nodeType: string
    name: string
    startLine: number
    endLine: number
  }
}

export function chunkCodeByAST(sourceCode: string, language = 'typescript'): CodeChunk[] {
  const parser = new Parser()
  parser.setLanguage(TypeScript.typescript)
  const tree = parser.parse(sourceCode)

  const chunks: CodeChunk[] = []
  const targetTypes = new Set(['function_declaration', 'class_declaration', 'method_definition', 'interface_declaration'])

  function traverse(node: Parser.SyntaxNode) {
    if (targetTypes.has(node.type)) {
      // 提取标识符名称（函数名/类名）
      const identifierNode = node.childForFieldName('name')
      const name = identifierNode ? identifierNode.text : 'anonymous'

      chunks.push({
        content: node.text,
        metadata: {
          language,
          nodeType: node.type,
          name,
          startLine: node.startPosition.row + 1,
          endLine: node.endPosition.row + 1
        }
      })
    }

    for (let i = 0; i < node.childCount; i++) {
      const child = node.child(i)
      if (child) traverse(child)
    }
  }

  traverse(tree.rootNode)
  return chunks
}

export function chunkByWindow(text: string, size = 512, overlap = 64): string[] {
  const chunks: string[] = []
  const step = size - overlap
  for (let i = 0; i < text.length; i += step) {
    chunks.push(text.slice(i, i + size))
  }
  return chunks
}

interface SearchResult {
  id: string
  content: string
  metadata: Record<string, any>
}

export function reciprocalRankFusion(
  rankedLists: SearchResult[][],
  k = 60,
  topN = 5
): SearchResult[] {
  const scoreMap = new Map<string, { score: number; item: SearchResult }>()

  for (const list of rankedLists) {
    list.forEach((item, rank) => {
      // rank 从 0 开始，排名越高（rank 越小），增益越大
      const rrfScore = 1 / (k + rank + 1)
      const existing = scoreMap.get(item.id)
      if (existing) {
        existing.score += rrfScore
      } else {
        scoreMap.set(item.id, { score: rrfScore, item })
      }
    })
  }

  return [...scoreMap.values()]
    .sort((a, b) => b.score - a.score)
    .slice(0, topN)
    .map(entry => entry.item)
}

import { pipeline } from '@xenova/transformers'

export class RerankerService {
  private ranker: any

  async init() {
    // 也可直接请求本地或远程部署的 bge-reranker-v2-m3 服务
    this.ranker = await pipeline('text-classification', 'Xenova/bge-reranker-base')
  }

  async rerank(query: string, candidates: SearchResult[], topN = 5): Promise<SearchResult[]> {
    const scored = await Promise.all(
      candidates.map(async item => {
        const out = await this.ranker({ text: query, text_pair: item.content })
        return { item, score: out[0].score }
      })
    )

    return scored
      .sort((a, b) => b.score - a.score)
      .slice(0, topN)
      .map(entry => entry.item)
  }
}

import { ToolDefinition } from './types'

export const searchCodebaseTool: ToolDefinition = {
  name: 'search_codebase',
  description: '在代码库或技术文档中执行语义检索与符号匹配。当用户询问宏观架构、涉及跨文件关联、或需要定位未知函数时调用此工具。',
  parameters: {
    type: 'object',
    properties: {
      query: { type: 'string', description: '搜索关键词或语义提问' },
      k: { type: 'integer', description: '期望返回的候选片段数量', default: 5 }
    },
    required: ['query']
  },
  execute: async ({ query, k = 5 }) => {
    // 1. 稠密向量检索
    const denseResults = await vectorStore.similaritySearch(query, k * 2)
    // 2. BM25 稀疏检索
    const sparseResults = await bm25Index.search(query, k * 2)
    // 3. RRF 融合
    const fused = reciprocalRankFusion([denseResults, sparseResults], 60, k * 2)
    // 4. 重排序
    const finalChunks = await reranker.rerank(query, fused, k)

    return finalChunks
      .map((c, i) => `### Result ${i + 1} (${c.metadata.filePath}:${c.metadata.startLine})\n\`\`\`${c.metadata.language}\n${c.content}\n\`\`\``)
      .join('\n\n')
  }
}