/**
 * @local/omp-ast — oh-my-pi's `ast_grep` and `ast_edit` for DSH, backed by the ast-grep CLI.
 *
 *  ast_grep  structural search: `$NAME` captures one node, `$$$NAME` zero-or-more, `$_`/`$$$`
 *            match without binding. Results are grouped per file with line numbers and captures.
 *  ast_edit  structural rewrite, preview first: returns a diff preview and a proposal id; calling
 *            again with `apply: "<id>"` writes it (refused if any file changed since the preview).
 *            `reject: "<id>"` drops it.
 */
import crypto from 'node:crypto'
import path from 'node:path'
import { findAstGrep, search, applyReplacements } from './lib/sg.js'

export const name = 'omp-ast'
export const inject = ['tools', 'fs']

const PATTERN_HELP = 'Pattern grammar: real code for the target language with metavariables — `$NAME` captures one AST node, `$$$NAME` captures zero or more nodes, `$_` / `$$$` match without binding. Names are UPPERCASE and stand for whole nodes (not parts of identifiers or strings); reusing a name requires identical code. The pattern must parse as one valid node (C/C++/Java statements need their trailing `;`).'

export function apply(ctx, config = {}) {
  const cfg = { maxMatches: 200, maxFiles: 1000, proposalTtlMinutes: 30, ...config }
  const cwdOf = (exec) => exec?.agent?.session?.header?.cwd ?? process.cwd()
  const bin = () => {
    const b = findAstGrep(cfg.binary ? [path.dirname(cfg.binary)] : [])
    if (!b) throw new Error('ast-grep is not installed. Install it with `npm i -g @ast-grep/cli` (or `npm i -g --prefix ~/.local @ast-grep/cli`).')
    return b
  }
  const splitPaths = (p) => (Array.isArray(p) ? p : String(p ?? '').split(';')).map((s) => String(s).trim()).filter(Boolean)
  const rel = (f, cwd) => { const r = path.relative(cwd, path.resolve(cwd, f)); return r && !r.startsWith('..') ? r : f }

  /** Sandbox fence for writes: same rule as DSH workspace-write (inside workspace root). */
  const policyFor = (exec) => { try { return ctx.get('sandboxPolicy')?.resolve({ ...(exec?.agent ? { session: exec.agent.session } : {}) }) } catch { return undefined } }

  // ---------------------------------------------------------------- ast_grep
  ctx.tools.register({
    name: 'ast_grep',
    description: `Structural code search with ast-grep over tree-sitter grammars (js/ts/tsx, python, go, rust, java, c/cpp, c#, kotlin, swift, ruby, php, lua, bash, html, css, json, yaml, …). Use it instead of regex grep when the shape of the code matters (calls with any arguments, functions with a decorator, a specific nesting). ${PATTERN_HELP} Results are grouped per file with \`LINE:TEXT\` and the captured metavariables.`,
    parameters: {
      type: 'object',
      additionalProperties: false,
      required: ['pat'],
      properties: {
        pat: { type: 'string', description: 'One AST pattern, e.g. `console.log($$$ARGS)` or `def $F($$$): $$$BODY`.' },
        path: { type: 'string', description: 'File, directory or glob; several separated by `;`. Defaults to the session workspace.' },
        lang: { type: 'string', description: 'Language override (e.g. ts, tsx, python, cpp). Inferred from file extensions when omitted.' },
        globs: { type: 'array', items: { type: 'string' }, description: 'Include/exclude globs (prefix ! to exclude), e.g. ["src/**", "!**/*.test.ts"].' },
        selector: { type: 'string', description: 'AST kind to select inside the pattern (advanced; e.g. `pair` for a JSON key pattern).' },
        skip: { type: 'integer', description: 'Skip this many matches (paging).' },
      },
    },
    output: { schema: { type: 'object', additionalProperties: true, properties: { text: { type: 'string' } } }, render: (_a, v) => [{ type: 'text', text: v.text }] },
    timeoutMs: 150_000,
    isConcurrencySafe: () => true,
    async execute(args, exec) {
      const pat = String(args.pat ?? '').trim()
      if (!pat) throw new Error('pat must be a non-empty AST pattern')
      const cwd = cwdOf(exec)
      const { matches, warnings } = await search({ bin: bin(), pattern: pat, lang: args.lang, paths: splitPaths(args.path), globs: args.globs, selector: args.selector, cwd, signal: exec?.signal })
      const skip = Math.max(0, Math.floor(args.skip ?? 0))
      const page = matches.slice(skip, skip + cfg.maxMatches)
      const byFile = new Map()
      for (const m of page) { if (!byFile.has(m.file)) byFile.set(m.file, []); byFile.get(m.file).push(m) }
      const out = []
      for (const [file, ms] of byFile) {
        out.push(rel(file, cwd))
        for (const m of ms) {
          const start = m.range.start.line + 1
          const lines = String(m.lines ?? m.text).split('\n')
          lines.forEach((l, i) => out.push(`${start + i}:${l.length > 300 ? `${l.slice(0, 300)}…` : l}`))
          const caps = captures(m)
          if (caps) out.push(`  ↳ ${caps}`)
        }
      }
      const files = new Set(matches.map((m) => m.file)).size
      let header = matches.length ? `${matches.length} match${matches.length === 1 ? '' : 'es'} in ${files} file${files === 1 ? '' : 's'}` : 'No matches.'
      if (matches.length > skip + cfg.maxMatches) header += ` (showing ${skip + 1}-${skip + page.length}; use skip=${skip + page.length} for more)`
      const warn = warnings.filter((w) => !/^$/.test(w))
      return { text: [header, ...out, ...(warn.length ? ['', 'ast-grep:', ...warn.map((w) => `  ${w}`)] : [])].join('\n'), count: matches.length }
    },
  })

  // ---------------------------------------------------------------- ast_edit
  /** @type {Map<string, {at:number, files: Map<string,{before:Buffer, after:Buffer}>, summary:string, owner:any}>} */
  const proposals = new Map()
  const gc = () => { const ttl = cfg.proposalTtlMinutes * 60_000; for (const [id, p] of proposals) if (Date.now() - p.at > ttl) proposals.delete(id) }

  ctx.tools.register({
    name: 'ast_edit',
    description: `Structural rewrite with ast-grep, previewed before it is applied. Pass \`ops\` (each {pat, out}: captures from \`pat\` are substituted into \`out\`; empty \`out\` deletes the node) and \`paths\`. The call returns a diff preview and a proposal id WITHOUT modifying files; call ast_edit again with \`apply: "<id>"\` to write it (refused if a file changed since the preview) or \`reject: "<id>"\` to drop it. Each rewrite is a 1:1 node substitution. ${PATTERN_HELP}`,
    parameters: {
      type: 'object',
      additionalProperties: false,
      properties: {
        ops: {
          type: 'array',
          description: 'Rewrite rules applied in order.',
          items: { type: 'object', additionalProperties: false, required: ['pat', 'out'], properties: { pat: { type: 'string' }, out: { type: 'string' } } },
        },
        paths: { type: 'array', items: { type: 'string' }, description: 'Files, directories or globs to rewrite.' },
        lang: { type: 'string', description: 'Language override.' },
        globs: { type: 'array', items: { type: 'string' }, description: 'Include/exclude globs.' },
        apply: { type: 'string', description: 'Proposal id from a previous preview: write it.' },
        reject: { type: 'string', description: 'Proposal id to discard.' },
      },
    },
    output: {
      schema: { type: 'object', additionalProperties: true, properties: { text: { type: 'string' } } },
      render: (_a, v) => [{ type: 'text', text: v.text }],
      presentationMeta: (_a, v) => ({ diffs: Array.isArray(v.diffs) ? v.diffs : [] }),
    },
    presentResult(args, result) {
      if (result.isError) return undefined
      const diffs = result.meta?.diffs
      if (!Array.isArray(diffs) || !diffs.length) return undefined
      return { card: 'diff', title: args.apply ? 'AST edit (applied)' : 'AST edit (preview)', diffs }
    },
    timeoutMs: 300_000,
    async execute(args, exec) {
      gc()
      const cwd = cwdOf(exec)
      if (args.reject) {
        const ok = proposals.delete(args.reject)
        return { text: ok ? `Proposal ${args.reject} discarded; no files were changed.` : `No pending proposal ${args.reject}.` }
      }
      if (args.apply) return applyProposal(args.apply, exec, cwd)
      const ops = Array.isArray(args.ops) ? args.ops : []
      if (!ops.length) throw new Error('ops must contain at least one {pat, out} rewrite (or pass apply/reject with a proposal id)')
      const seen = new Set()
      for (const [i, o] of ops.entries()) {
        if (!o || typeof o.pat !== 'string' || !o.pat.trim()) throw new Error(`ops[${i}].pat must be a non-empty pattern`)
        if (typeof o.out !== 'string') throw new Error(`ops[${i}].out must be a string (empty deletes)`)
        if (seen.has(o.pat)) throw new Error(`ops[${i}] duplicates pattern ${JSON.stringify(o.pat)}`)
        seen.add(o.pat)
      }
      const paths = splitPaths(args.paths)
      if (!paths.length) throw new Error('paths must list at least one file, directory or glob')
      const fsp = await import('node:fs/promises')
      // Apply ops sequentially in memory: each op's rewrite is computed on the result of the previous.
      const files = new Map() // abs -> {before, after}
      const counts = []
      const warnings = []
      for (const op of ops) {
        const { matches, warnings: w } = await search({ bin: bin(), pattern: op.pat, rewrite: op.out, lang: args.lang, paths, globs: args.globs, cwd, signal: exec?.signal })
        warnings.push(...w)
        // ast-grep reads from disk. Files an earlier op already rewrote are re-searched on a temp copy
        // of their in-memory version; the rest use this op's disk matches.
        const rep = (m) => ({ start: m.replacementOffsets?.start ?? m.range.byteOffset.start, end: m.replacementOffsets?.end ?? m.range.byteOffset.end, text: m.replacement ?? '' })
        const byFile = new Map()
        for (const m of matches) { const abs = path.resolve(cwd, m.file); if (!byFile.has(abs)) byFile.set(abs, []); byFile.get(abs).push(m) }
        let n = 0
        for (const [abs, entry] of files) {
          if (entry.after === entry.before) continue
          byFile.delete(abs)
          const tmp = path.join(path.dirname(abs), `.omp-ast-${crypto.randomBytes(4).toString('hex')}${path.extname(abs)}`)
          await fsp.writeFile(tmp, entry.after)
          try {
            const r = await search({ bin: bin(), pattern: op.pat, rewrite: op.out, lang: args.lang, paths: [tmp], cwd, signal: exec?.signal })
            if (r.matches.length) { entry.after = applyReplacements(entry.after, r.matches.map(rep)); n += r.matches.length }
          } finally { await fsp.rm(tmp, { force: true }) }
        }
        for (const [abs, ms] of byFile) {
          let entry = files.get(abs)
          if (!entry) { const b = await fsp.readFile(abs); entry = { before: b, after: b }; files.set(abs, entry) }
          entry.after = applyReplacements(entry.after, ms.map(rep))
          n += ms.length
        }
        counts.push(`${JSON.stringify(op.pat)} → ${n} replacement${n === 1 ? '' : 's'}`)
      }
      for (const [abs, e] of files) if (Buffer.compare(e.before, e.after) === 0) files.delete(abs)
      if (files.size > cfg.maxFiles) throw new Error(`rewrite touches ${files.size} files (max ${cfg.maxFiles}); narrow paths/globs`)
      if (!files.size) return { text: ['No changes: no pattern matched.', ...counts, ...(warnings.length ? ['ast-grep:', ...warnings] : [])].join('\n') }
      const id = crypto.randomBytes(3).toString('hex')
      proposals.set(id, { at: Date.now(), files, owner: exec?.agent?.session })
      const diffs = []
      const lines = []
      for (const [abs, e] of files) {
        const d = lineDiff(e.before.toString('utf8'), e.after.toString('utf8'))
        diffs.push(...d.map((h) => ({ path: rel(abs, cwd), ...h })))
        lines.push(`${rel(abs, cwd)}:`)
        for (const h of d.slice(0, 20)) {
          for (const l of (h.oldText ?? '').split('\n')) lines.push(`- ${l}`)
          for (const l of h.newText.split('\n')) lines.push(`+ ${l}`)
        }
      }
      const shown = lines.length > 400 ? [...lines.slice(0, 400), `… ${lines.length - 400} more preview lines`] : lines
      return {
        text: [
          `Staged as proposal ${id} — files NOT modified yet. ${files.size} file${files.size === 1 ? '' : 's'}:`,
          ...counts, '', ...shown, '',
          `Apply: ast_edit {"apply": "${id}"} · Discard: ast_edit {"reject": "${id}"}`,
          ...(warnings.length ? ['', 'ast-grep:', ...warnings.map((w) => `  ${w}`)] : []),
        ].join('\n'),
        proposal: id,
        diffs,
      }
    },
  })

  async function applyProposal(id, exec, cwd) {
    const p = proposals.get(id)
    if (!p) throw new Error(`No pending proposal ${id} (expired after ${cfg.proposalTtlMinutes} min, already applied, or never issued). Run the preview again.`)
    if (p.owner && exec?.agent?.session && p.owner !== exec.agent.session) throw new Error(`Proposal ${id} belongs to another session.`)
    const fsp = await import('node:fs/promises')
    // Verify nothing changed since the preview, then write all.
    for (const [abs, e] of p.files) {
      const now = await fsp.readFile(abs).catch(() => null)
      if (!now || Buffer.compare(now, e.before) !== 0) throw new Error(`${rel(abs, cwd)} changed since the preview; proposal ${id} not applied. Run ast_edit again.`)
    }
    const policy = policyFor(exec)
    const written = []
    const diffs = []
    for (const [abs, e] of p.files) {
      const target = await ctx.fs.resolve(abs, { signal: exec?.signal })
      const outcome = await ctx.fs.writeText(target, e.after.toString('utf8'), undefined, exec?.signal, policy)
      ctx.emit('fs/observed', target, { kind: 'present', version: outcome.version }, exec)
      written.push(rel(abs, cwd))
      diffs.push(...lineDiff(e.before.toString('utf8'), e.after.toString('utf8')).map((h) => ({ path: rel(abs, cwd), ...h })))
    }
    proposals.delete(id)
    return { text: `Applied proposal ${id} to ${written.length} file${written.length === 1 ? '' : 's'}:\n${written.join('\n')}`, results: written.map((f) => ({ path: f, operation: 'update' })), diffs }
  }

  ctx.effect(() => () => proposals.clear(), 'omp-ast: proposals')
}

function captures(m) {
  const mv = m.metaVariables
  if (!mv) return ''
  const parts = []
  for (const [k, v] of Object.entries(mv.single ?? {})) parts.push(`$${k}=${short(v.text)}`)
  for (const [k, v] of Object.entries(mv.multi ?? {})) parts.push(`$$$${k}=[${v.filter((x) => !/^[,;]$/.test(x.text)).map((x) => short(x.text)).join(', ')}]`)
  return parts.join('  ')
}
const short = (s) => { const t = String(s).replace(/\s+/g, ' '); return t.length > 60 ? `${t.slice(0, 60)}…` : t }

/** Minimal line diff → hunks with 2 lines of context (LCS on lines). */
export function lineDiff(a, b, context = 2) {
  const A = a.split('\n'); const B = b.split('\n')
  let s = 0
  while (s < A.length && s < B.length && A[s] === B[s]) s++
  let ea = A.length; let eb = B.length
  while (ea > s && eb > s && A[ea - 1] === B[eb - 1]) { ea--; eb-- }
  if (s === ea && s === eb) return []
  // Split the middle into hunks by common lines (simple LCS for moderate sizes).
  const midA = A.slice(s, ea); const midB = B.slice(s, eb)
  const n = midA.length; const m = midB.length
  const hunks = []
  if (n * m <= 1_000_000) {
    const dp = Array.from({ length: n + 1 }, () => new Uint32Array(m + 1))
    for (let i = n - 1; i >= 0; i--) for (let j = m - 1; j >= 0; j--) dp[i][j] = midA[i] === midB[j] ? dp[i + 1][j + 1] + 1 : Math.max(dp[i + 1][j], dp[i][j + 1])
    let i = 0; let j = 0; let cur = null
    const flush = () => { if (cur) { hunks.push(cur); cur = null } }
    while (i < n || j < m) {
      if (i < n && j < m && midA[i] === midB[j]) { flush(); i++; j++; continue }
      cur ??= { a0: s + i, b0: s + j, old: [], neu: [] }
      if (j < m && (i >= n || dp[i][j + 1] >= dp[i + 1][j])) cur.neu.push(midB[j++])
      else cur.old.push(midA[i++])
    }
    flush()
  } else hunks.push({ a0: s, b0: s, old: midA, neu: midB })
  return hunks.map((h) => {
    const pre = A.slice(Math.max(0, h.a0 - context), h.a0)
    const postStart = h.a0 + h.old.length
    const post = A.slice(postStart, postStart + context)
    return { oldText: [...pre, ...h.old, ...post].join('\n') || null, newText: [...pre, ...h.neu, ...post].join('\n') }
  })
}