FindFiles.hs

module LLM.Tools.FindFiles
  ( findFilesToolTyped,
    FindFilesToolArgs (..),
  )
where

import Autodocodec qualified as AC
import Control.Exception (IOException, try)
import Data.Aeson (FromJSON)
import Data.List (sort)
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Generics (Generic)
import LLM.Core.Types (TypedTool (..))
import LLM.Tools.FsConfig (FsConfig, isFileHidden, isSymlink, sandboxPath, withWorkspacePaths)
import LLM.Tools.FsLimits (cheapWalkDepth, maxOutputBytes, maxVisitedEntries)
import LLM.Tools.Internal.WalkBudget (WalkBudget, budgetExhausted, budgetNote, capLines, newWalkBudget, takeEntries)
import System.Directory (doesDirectoryExist, listDirectory)
import System.FilePath ((</>))

-- | Default cap on returned paths.
defaultMaxResults :: Int
defaultMaxResults = 200

-- | Hard cap to protect the model's context window.
maxResultsCap :: Int
maxResultsCap = 2000

-- | Maximum directory depth from the search root. Sourced from
-- 'LLM.Tools.FsLimits.cheapWalkDepth' because @find_files@ only stats
-- entries, so it can afford a deeper walk than @grep@.
maxDepth :: Int
maxDepth = cheapWalkDepth

-- | What kind of filesystem entries to return.
data EntryFilter = AnyEntry | FileOnly | DirectoryOnly
  deriving (Eq, Show)

-- | 'Nothing' for an unknown kind.
parseEntryFilter :: Text -> Maybe EntryFilter
parseEntryFilter t = case T.toLower (T.strip t) of
  "file" -> Just FileOnly
  "directory" -> Just DirectoryOnly
  "dir" -> Just DirectoryOnly
  "any" -> Just AnyEntry
  "" -> Just AnyEntry
  _ -> Nothing

data FindFilesToolArgs = FindFilesToolArgs
  { _ffPattern :: Text,
    _ffPath :: Text,
    _ffType :: Text,
    _ffIncludeHidden :: Bool,
    _ffMaxResults :: Int
  }
  deriving (Generic)
  deriving (FromJSON) via (AC.Autodocodec FindFilesToolArgs)

instance AC.HasCodec FindFilesToolArgs where
  codec :: AC.JSONCodec FindFilesToolArgs
  codec =
    AC.object "find files by glob pattern" $
      FindFilesToolArgs
        <$> AC.requiredField
          "pattern"
          "Glob pattern matched against paths relative to 'path'. Supports '**' (any number of path segments, including zero), '*' (any chars except '/'), and '?' (any single char except '/'). Examples: '**/*.hs', 'src/**/Tools/*.hs', '*.cabal'."
          AC..= (._ffPattern)
        <*> AC.optionalFieldWithDefault
          "path"
          "."
          "Relative directory to search under (default '.')"
          AC..= (._ffPath)
        <*> AC.optionalFieldWithDefault
          "type"
          "any"
          "Entry kind to return: 'file', 'directory', or 'any' (default 'any')"
          AC..= (._ffType)
        <*> AC.optionalFieldWithDefault
          "includeHidden"
          False
          "Include dotfiles and dot-directories (default false)"
          AC..= (._ffIncludeHidden)
        <*> AC.optionalFieldWithDefault
          "maxResults"
          defaultMaxResults
          "Maximum number of paths to return (default 200, hard cap 2000)"
          AC..= (._ffMaxResults)

findFilesToolTyped :: FsConfig -> TypedTool ctx FindFilesToolArgs
findFilesToolTyped cfg =
  TypedTool
    { ttoolName = "find_files",
      ttoolDescription =
        "Find files and directories by glob pattern under a workspace-relative root. "
          <> "Supports '**' (any number of path segments, including zero), '*' (any chars except '/'), and '?'. "
          <> "Returns one path per line, relative to the search root, sorted. "
          <> "Skips hidden entries unless 'includeHidden' is true. "
          <> "Use 'type' to restrict to files or directories. "
          <> "Output is capped at 'maxResults' (default 200, hard cap 2000); "
          <> "the header reports whether results were truncated and how many entries were scanned. "
          <> "A search looks at most "
          <> T.pack (show maxVisitedEntries)
          <> " entries and says so if it stopped there.",
      ttoolReadonly = True,
      ttoolExecute = const (withWorkspacePaths cfg . execute cfg)
    }

execute :: FsConfig -> FindFilesToolArgs -> IO Text
execute cfg args = do
  let rel = T.unpack args._ffPath
      pat = args._ffPattern
      includeHidden = args._ffIncludeHidden
      cap = min maxResultsCap args._ffMaxResults
  case parseEntryFilter args._ffType of
    _ | T.null pat -> pure "Error: pattern must be non-empty"
    Nothing ->
      pure $ "Error: unknown type " <> T.pack (show args._ffType) <> "; use 'file', 'directory' or 'any'"
    _ | args._ffMaxResults <= 0 ->
      pure $ "Error: maxResults must be positive, got " <> T.pack (show args._ffMaxResults)
    Just kind -> do
      resolved <- sandboxPath cfg rel
      isDir <- doesDirectoryExist resolved
      if not isDir
        then pure $ "Error: path is not a directory: " <> args._ffPath
        else do
          let segs = splitGlob pat
          budget <- newWalkBudget maxVisitedEntries
          (matches, scanned, truncated) <- walk budget resolved includeHidden kind segs cap
          stopped <- budgetExhausted budget
          let note = if stopped then Just (budgetNote budget) else Nothing
          pure $ renderResults args._ffPath pat matches scanned truncated note

-- ---------------------------------------------------------------------------
-- Filesystem walk
-- ---------------------------------------------------------------------------

-- | Result of the walk: matched relative paths (sorted within each directory),
-- count of entries inspected, and whether the cap truncated the output.
walk ::
  WalkBudget ->
  FilePath ->
  Bool ->
  EntryFilter ->
  [Text] ->
  Int ->
  IO ([FilePath], Int, Bool)
walk budget root includeHidden kind pat cap = do
  WalkState acc _ scanned truncated <- go root [] 0 (WalkState [] 0 0 False)
  pure (reverse acc, scanned, truncated)
  where
    -- relSegs is the path from root to current dir, in reverse order
    -- depth is current directory depth (0 = root)
    go dir relSegs depth st@(WalkState acc n scanned _)
      | n >= cap = pure (WalkState acc n scanned True)
      | depth > maxDepth = pure st
      | otherwise = do
          entries <- sort <$> safeListDirectory dir
          let visible =
                if includeHidden
                  then entries
                  else filter (not . isFileHidden) entries
          allowed <- takeEntries budget visible
          step allowed dir relSegs depth st

    step [] _ _ _ st = pure st
    step (name : rest) dir relSegs depth (WalkState acc n scanned truncated)
      | n >= cap = pure (WalkState acc n (scanned + 1) True)
      | otherwise = do
          let full = dir </> name
              entryRel = reverse (T.pack name : relSegs)
              relPath = T.unpack (T.intercalate "/" entryRel)
          -- Skip symlinks: they could point outside the sandbox.
          isLink <- isSymlink full
          if isLink
            then step rest dir relSegs depth (WalkState acc n (scanned + 1) truncated)
            else do
              isDir <- doesDirectoryExist full
              let entryKind = if isDir then DirectoryOnly else FileOnly
                  kindOk = case kind of
                    AnyEntry -> True
                    FileOnly -> entryKind == FileOnly
                    DirectoryOnly -> entryKind == DirectoryOnly
                  matched = kindOk && matchGlob pat entryRel
                  st'
                    | matched = WalkState (relPath : acc) (n + 1) (scanned + 1) truncated
                    | otherwise = WalkState acc n (scanned + 1) truncated
              st'' <-
                if isDir
                  then go full (T.pack name : relSegs) (depth + 1) st'
                  else pure st'
              step rest dir relSegs depth st''

-- | Walk state: matches so far (newest first), their count, entries
-- inspected, and whether the cap truncated the output.
data WalkState = WalkState [FilePath] Int Int Bool

safeListDirectory :: FilePath -> IO [FilePath]
safeListDirectory path = do
  r <- try (listDirectory path) :: IO (Either IOException [FilePath])
  case r of
    Right xs -> pure xs
    Left _ -> pure []

-- ---------------------------------------------------------------------------
-- Glob matcher
-- ---------------------------------------------------------------------------

-- | Split a glob into path segments. Leading "./" is dropped so callers
-- can write either "**/*.hs" or "./**/*.hs".
splitGlob :: Text -> [Text]
splitGlob pat =
  let stripped = fromMaybe pat (T.stripPrefix "./" pat)
   in filter (not . T.null) (T.splitOn "/" stripped)

-- | Match a list of glob segments against a list of path segments.
-- '**' may match zero or more path segments; other segments must match
-- exactly one segment via 'matchSegment'.
matchGlob :: [Text] -> [Text] -> Bool
matchGlob [] [] = True
matchGlob [] _ = False
matchGlob ("**" : ps) ss =
  -- '**' matches zero segments, or one+ segments then continues
  matchGlob ps ss
    || case ss of
      [] -> False
      (_ : rest) -> matchGlob ("**" : ps) rest
matchGlob _ [] = False
matchGlob (p : ps) (s : ss) = matchSegment (T.unpack p) (T.unpack s) && matchGlob ps ss

-- | Match a single glob segment (no '/' allowed in either side).
-- Supports '*' (any run of chars, possibly empty) and '?' (one char).
-- On a mismatch only the most recent '*' takes one more char, which is
-- enough for this glob syntax and keeps matching O(pattern * name) instead
-- of exponential in the number of stars.
matchSegment :: String -> String -> Bool
matchSegment pat0 str0 = go pat0 str0 Nothing
  where
    -- lastStar: the pattern after the most recent '*' and the input it
    -- was last retried from.
    go ('*' : ps) ss _ = go ps ss (Just (ps, ss))
    go (p : ps) (s : ss) lastStar
      | p == '?' || p == s = go ps ss lastStar
    go [] [] _ = True
    go _ _ (Just (ps, _ : ss)) = go ps ss (Just (ps, ss))
    go _ _ _ = False

-- ---------------------------------------------------------------------------
-- Rendering
-- ---------------------------------------------------------------------------

renderResults :: Text -> Text -> [FilePath] -> Int -> Bool -> Maybe Text -> Text
renderResults rootDisplay pat allMatches scanned truncated0 note =
  let (body, cut) = capLines maxOutputBytes (map T.pack allMatches)
      matches = body
      truncated = truncated0 || cut
      header =
        "=== find_files "
          <> T.pack (show pat)
          <> " in "
          <> rootDisplay
          <> " | "
          <> T.pack (show (length matches))
          <> " match(es) across "
          <> T.pack (show scanned)
          <> " entr"
          <> (if scanned == 1 then "y" else "ies")
          <> ( if cut
                 then " | output truncated at " <> T.pack (show maxOutputBytes) <> " bytes"
                 else if truncated then " | truncated (raise maxResults to see more)" else ""
             )
          <> maybe "" (" | " <>) note
          <> " ==="
   in T.unlines (header : body)