import { parametersSchema as z, defineCustomTool } from "@roo-code/types"
import { statSync, readdirSync } from "fs"
import { join, relative } from "path"

export default defineCustomTool({
  name: "context_budget",
  description: "Estimate token cost of reading files or directories. Returns file count, total lines, chars, and estimated tokens. Helps decide whether to load full files or use targeted reads.",
  parameters: z.object({
    paths: z.array(z.string()).describe("File or directory paths to analyze"),
    recursive: z.boolean().optional().describe("Recurse into directories. Default: true"),
    extensions: z.array(z.string()).optional().describe("Filter by file extensions, e.g. ['.java', '.ts']. Default: all files"),
  }),
  async execute({ paths, recursive = true, extensions }) {
    let totalFiles = 0
    let totalLines = 0
    let totalChars = 0
    const breakdown: Array<{ path: string; lines: number; chars: number; tokens: number }> = []

    function processFile(filePath: string) {
      try {
        const stat = statSync(filePath)
        if (!stat.isFile()) return
        if (extensions && !extensions.some(ext => filePath.endsWith(ext))) return
        
        const { readFileSync } = require("fs")
        const content = readFileSync(filePath, "utf-8")
        const lines = content.split("\n").length
        const chars = content.length
        const tokens = Math.ceil(chars / 4)
        
        totalFiles++
        totalLines += lines
        totalChars += chars
        breakdown.push({ path: filePath, lines, chars, tokens })
      } catch (e) {
        // skip unreadable files
      }
    }

    function processDir(dirPath: string, recurse: boolean) {
      try {
        const entries = readdirSync(dirPath, { withFileTypes: true })
        for (const entry of entries) {
          if (entry.name.startsWith(".") || entry.name === "node_modules" || entry.name === "target" || entry.name === ".git") continue
          const fullPath = join(dirPath, entry.name)
          if (entry.isFile()) {
            processFile(fullPath)
          } else if (entry.isDirectory() && recurse) {
            processDir(fullPath, recurse)
          }
        }
      } catch (e) {
        // skip unreadable dirs
      }
    }

    for (const p of paths) {
      try {
        const stat = statSync(p)
        if (stat.isFile()) {
          processFile(p)
        } else if (stat.isDirectory()) {
          processDir(p, recursive !== false)
        }
      } catch (e) {
        breakdown.push({ path: p, lines: 0, chars: 0, tokens: 0 })
      }
    }

    const totalTokens = Math.ceil(totalChars / 4)
    
    // Sort by tokens descending, take top 15
    breakdown.sort((a, b) => b.tokens - a.tokens)
    const topFiles = breakdown.slice(0, 15)

    const result = {
      summary: {
        files: totalFiles,
        totalLines,
        totalChars,
        estimatedTokens: totalTokens,
        warning: totalTokens > 50000 ? "⚠️ LARGE — will consume significant context budget" : 
                 totalTokens > 20000 ? "⚠️ MEDIUM — consider targeted reads" : 
                 "✅ Fits comfortably in context",
      },
      largestFiles: topFiles,
    }

    return JSON.stringify(result, null, 2)
  },
})