/**
 * pi-codebase-map v1.1 — TF-IDF codebase search for Pi
 *
 * Index your project, search by meaning (not just grep).
 * Zero external dependencies. Pure TypeScript.
 *
 * /index          — reindex current project
 * /search <query> — search the codebase
 */
import type { ExtensionAPI } from '@mariozechner/pi-coding-agent'
import { Type } from '@sinclair/typebox'
import { execSync } from 'child_process'
import { readFileSync, writeFileSync, mkdirSync, statSync } from 'fs'
import { join, extname } from 'path'

// === Types ===
interface Chunk { file: string; line: number; text: string; tokens: string[] }
interface Index { chunks: Chunk[]; df: Record<string, number>; totalDocs: number; indexed: number; cwd: string }

const INDEX_DIR = '.pi-codebase-map'
const INDEX_FILE = 'index.json'
const CODE_EXTS = new Set(['.ts','.tsx','.js','.jsx','.py','.rs','.go','.java','.c','.cpp','.h','.hpp','.cs','.rb','.php','.swift','.kt','.scala','.sh','.bash','.zsh','.yml','.yaml','.toml','.json','.md','.txt','.sql','.html','.css','.scss','.vue','.svelte'])
const MAX_FILE_SIZE = 200_000 // 200KB

// === Tokenizer ===
function tokenize(text: string): string[] {
  return text.toLowerCase()
    .replace(/[^a-z0-9_]/g, ' ')
    .split(/\s+/)
    .filter(t => t.length > 1 && t.length < 40)
}

// === Chunker ===
function chunkFile(content: string, file: string): Chunk[] {
  const lines = content.split('\n')
  const chunks: Chunk[] = []
  let buf: string[] = []
  let startLine = 1

  const flush = () => {
    if (buf.length > 0) {
      const text = buf.join('\n')
      if (text.trim().length > 20) {
        chunks.push({ file, line: startLine, text: text.slice(0, 1000), tokens: tokenize(text) })
      }
      buf = []
    }
  }

  for (let i = 0; i < lines.length; i++) {
    const line = lines[i]
    const isBreak = /^(export |function |class |def |fn |pub |interface |type |const |let |var |import |from |## |### |---- )/.test(line.trimStart())
    if (isBreak && buf.length > 3) {
      flush()
      startLine = i + 1
    }
    buf.push(line)
    if (buf.length >= 30) {
      flush()
      startLine = i + 2
    }
  }
  flush()
  return chunks
}

// === TF-IDF ===
function tf(tokens: string[]): Record<string, number> {
  const counts: Record<string, number> = {}
  for (const t of tokens) counts[t] = (counts[t] || 0) + 1
  const max = Math.max(...Object.values(counts), 1)
  const result: Record<string, number> = {}
  for (const [t, c] of Object.entries(counts)) result[t] = c / max
  return result
}

function tfidfScore(queryTokens: string[], chunk: Chunk, df: Record<string, number>, totalDocs: number): number {
  const chunkTf = tf(chunk.tokens)
  let score = 0
  for (const qt of queryTokens) {
    const termTf = chunkTf[qt] || 0
    const termDf = df[qt] || 0
    if (termTf > 0 && termDf > 0) {
      const idf = Math.log(totalDocs / termDf)
      score += termTf * idf
    }
  }
  return score
}

// === Indexing ===
function getProjectFiles(cwd: string): string[] {
  try {
    const raw = execSync('git ls-files --cached --others --exclude-standard', { encoding: 'utf-8', cwd, timeout: 15000 })
    return raw.trim().split('\n').filter(f => {
      const ext = extname(f).toLowerCase()
      return CODE_EXTS.has(ext) && !f.includes('node_modules') && !f.includes('.min.') && !f.includes('dist/') && !f.includes('build/')
    })
  } catch {
    return []
  }
}

function buildIndex(cwd: string): Index {
  const files = getProjectFiles(cwd)
  const chunks: Chunk[] = []

  for (const file of files) {
    try {
      const full = join(cwd, file)
      const stat = statSync(full)
      if (stat.size > MAX_FILE_SIZE) continue
      const content = readFileSync(full, 'utf-8')
      chunks.push(...chunkFile(content, file))
    } catch { continue }
  }

  const df: Record<string, number> = {}
  for (const chunk of chunks) {
    const unique = new Set(chunk.tokens)
    for (const t of unique) df[t] = (df[t] || 0) + 1
  }

  return { chunks, df, totalDocs: chunks.length, indexed: Date.now(), cwd }
}

function saveIndex(cwd: string, index: Index) {
  const dir = join(cwd, INDEX_DIR)
  mkdirSync(dir, { recursive: true })
  writeFileSync(join(dir, INDEX_FILE), JSON.stringify(index))
}

function loadIndex(cwd: string): Index | null {
  try {
    const data = readFileSync(join(cwd, INDEX_DIR, INDEX_FILE), 'utf-8')
    return JSON.parse(data)
  } catch { return null }
}

// === Search ===
function search(query: string, index: Index, limit = 10): { file: string; line: number; score: number; preview: string }[] {
  const queryTokens = tokenize(query)
  if (!queryTokens.length) return []

  const scored = index.chunks
    .map(chunk => ({ chunk, score: tfidfScore(queryTokens, chunk, index.df, index.totalDocs) }))
    .filter(r => r.score > 0)
    .sort((a, b) => b.score - a.score)
    .slice(0, limit)

  return scored.map(r => ({
    file: r.chunk.file,
    line: r.chunk.line,
    score: Math.round(r.score * 100) / 100,
    preview: r.chunk.text.slice(0, 200).replace(/\n/g, ' '),
  }))
}

// === Extension ===
export default function init(pi: ExtensionAPI) {
  const cwd = process.cwd()

  pi.registerTool({
    name: 'codebase_index',
    label: 'Codebase Index',
    description: 'Index the current project for TF-IDF search.',
    parameters: Type.Object({
      path: Type.Optional(Type.String({ description: 'Directory to index (default: cwd)' })),
    }),
    async execute(_toolCallId, params) {
      const target = params.path || cwd
      const t0 = Date.now()
      const index = buildIndex(target)
      saveIndex(target, index)
      return { content: [{ type: 'text' as const, text: `Indexed ${index.totalDocs} chunks from ${new Set(index.chunks.map(c => c.file)).size} files in ${Date.now() - t0}ms.` }] }
    }
  })

  pi.registerTool({
    name: 'codebase_search',
    label: 'Codebase Search',
    description: 'Search the codebase using TF-IDF. More semantic than grep — matches related terms.',
    parameters: Type.Object({
      query: Type.String({ description: 'Natural language search query' }),
      limit: Type.Optional(Type.Number({ description: 'Max results (default 10)' })),
    }),
    async execute(_toolCallId, params) {
      let index = loadIndex(cwd)
      if (!index) {
        index = buildIndex(cwd)
        saveIndex(cwd, index)
      }
      const results = search(params.query, index, params.limit || 10)
      if (!results.length) return { content: [{ type: 'text' as const, text: `No results for "${params.query}".` }] }
      const lines = results.map((r, i) => `**${i + 1}.** \`${r.file}:${r.line}\` (score: ${r.score})\n   ${r.preview}`)
      return { content: [{ type: 'text' as const, text: `## Search: "${params.query}"\n\n${lines.join('\n\n')}` }] }
    }
  })

  pi.registerTool({
    name: 'codebase_status',
    label: 'Codebase Status',
    description: 'Show codebase index status.',
    parameters: Type.Object({}),
    async execute() {
      const index = loadIndex(cwd)
      if (!index) return { content: [{ type: 'text' as const, text: 'Not indexed. Run `codebase_index` first.' }] }
      const files = new Set(index.chunks.map(c => c.file)).size
      const age = Math.round((Date.now() - index.indexed) / 60000)
      const ageStr = age < 1 ? 'just now' : age < 60 ? `${age}m ago` : `${Math.floor(age / 60)}h ago`
      return { content: [{ type: 'text' as const, text: `## Codebase Index\n\n- **Files:** ${files}\n- **Chunks:** ${index.totalDocs}\n- **Vocabulary:** ${Object.keys(index.df).length} terms\n- **Indexed:** ${ageStr}\n- **Path:** ${index.cwd}` }] }
    }
  })

  // Commands
  pi.registerCommand('index', {
    description: 'Reindex the current project',
    handler: async (_args, ctx) => {
      const t0 = Date.now()
      const index = buildIndex(cwd)
      saveIndex(cwd, index)
      const files = new Set(index.chunks.map(c => c.file)).size
      ctx.ui.notify(`Indexed ${index.totalDocs} chunks from ${files} files in ${Date.now() - t0}ms.`, 'info')
    }
  })

  pi.registerCommand('codesearch', {
    description: 'Search the codebase — /codesearch <query>',
    handler: async (args, ctx) => {
      const query = args.trim()
      if (!query) { ctx.ui.notify('Usage: /codesearch <query>', 'info'); return }
      let index = loadIndex(cwd)
      if (!index) { index = buildIndex(cwd); saveIndex(cwd, index) }
      const results = search(query, index)
      ctx.ui.notify(results.length ? results.map((r, i) => `**${i + 1}.** \`${r.file}:${r.line}\` (${r.score})\n   ${r.preview}`).join('\n\n') : `No results for "${query}".`, 'info')
    }
  })
}