Complete.hs

{-|
Copyright        : (c) Galois, Inc. 2025
Maintainer       : Langston Barrett <langston@galois.com>
-}

{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}

module Lang.Crucible.Debug.Complete
  ( CompletionEnv(..)
  , CompletionT
  , runCompletionT
  , runCompletion
  , runCompletionM
  , CompletionType(..)
  , Completions(..)
  , complete
  ) where

import Control.Lens qualified as Lens
import Control.Monad qualified as Monad
import Control.Monad.IO.Class (MonadIO)
import Control.Monad.Reader (MonadReader, ReaderT)
import Control.Monad.Reader qualified as Reader
import Control.Monad.Trans (MonadTrans)
import Control.Exception qualified as X
import Data.Foldable (toList)
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Functor.Identity (Identity, runIdentity)
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Maybe qualified as Maybe
import Data.Parameterized.Some (Some(Some), viewSome)
import Data.Parameterized.SymbolRepr (symbolRepr)
import Data.Set (Set)
import Data.Text qualified as Text
import Data.Text (Text)
import Lang.Crucible.Debug.Arg qualified as Arg
import Lang.Crucible.Debug.Arg.Type (ArgTypeRepr)
import Lang.Crucible.Debug.Arg.Type qualified as AType
import Lang.Crucible.Debug.Breakpoint qualified as Break
import Lang.Crucible.Debug.Breakpoint (Breakpoint)
import Lang.Crucible.Debug.Command (CommandExt)
import Lang.Crucible.Debug.Command qualified as Cmd
import Lang.Crucible.FunctionHandle qualified as C
import Lang.Crucible.Simulator.ExecutionTree qualified as C
import Prettyprinter qualified as PP
import System.Directory qualified as Dir
import System.FilePath qualified as FilePath
import What4.FunctionName qualified as W4

data CompletionEnv cExt
  = forall p sym ext r.
    CompletionEnv
  { envBreakpoints :: Set Breakpoint
  , envCommandExt :: CommandExt cExt
  , envState :: C.ExecState p sym ext r
  }

newtype CompletionT cExt m a
  = CompletionT { runCompletionT :: ReaderT (CompletionEnv cExt) m a }
  deriving (Applicative, Functor, Monad, MonadFail, MonadIO, MonadReader (CompletionEnv cExt), MonadTrans)

-- | Run 'CompletionT' over 'Identity'
runCompletion :: CompletionEnv cExt -> CompletionT cExt Identity a -> a
runCompletion env s = runIdentity (Reader.runReaderT (runCompletionT s) env)

-- | Run 'CompletionT' over another monad
--
-- This is kind of poorly named, but I can\'t think of a better name...
runCompletionM :: CompletionEnv cExt -> CompletionT cExt m a -> m a
runCompletionM env s = Reader.runReaderT (runCompletionT s) env

data CompletionType
  = CBreakpoint
  | CCommand
  | CExactly Text
  | CFunction
  | CPath
  deriving (Eq, Show)

instance PP.Pretty CompletionType where
  pretty =
    \case
      CBreakpoint -> "breakpoint"
      CCommand -> "command"
      CExactly t -> PP.pretty t
      CFunction -> "function"
      CPath -> "path"

completionType :: ArgTypeRepr t -> Maybe CompletionType
completionType =
  \case
    AType.TBreakpointRepr -> Just CBreakpoint
    AType.TCommandRepr -> Just CCommand
    AType.TExactlyRepr s -> Just (CExactly (symbolRepr s))
    AType.TFunctionRepr -> Just CFunction
    AType.TIntRepr -> Nothing
    AType.TPathRepr -> Just CPath
    AType.TTextRepr -> Nothing

commandsWithPrefix ::
  CompletionEnv cExt ->
  -- | Prefix
  String ->
  [String]
commandsWithPrefix cEnv pfx = filter (pfx `List.isPrefixOf`) allCmdNames
  where
    allCmdNames :: [String]
    allCmdNames =
      let cExts = envCommandExt cEnv in
      map (Text.unpack . Cmd.name cExts) (Cmd.allCmds cExts)

completePath :: FilePath -> IO [String]
completePath path =
  if List.null path
  then Dir.listDirectory =<< Dir.getCurrentDirectory
  else do
    isDir <- Dir.doesDirectoryExist path
    if isDir && not (FilePath.hasTrailingPathSeparator path)
      then pure [FilePath.addTrailingPathSeparator path]
      else do
        let dir = FilePath.addTrailingPathSeparator (FilePath.takeDirectory path)
        contents <- Dir.listDirectory dir
        let file = FilePath.takeFileName path
        if List.null file
        then pure (map (dir ++) contents)
        else pure (map (dir ++) (filter (file `List.isPrefixOf`) contents))

completionsForType ::
  CompletionEnv cExt ->
  -- | Prefix
  String ->
  CompletionType ->
  IO [String]
completionsForType cEnv pfx =
  \case
    CBreakpoint ->
      pure $
        toList (envBreakpoints cEnv) &
          map Break.toText &
          filter (Text.pack pfx `Text.isPrefixOf`) &
          map Text.unpack
    CCommand -> pure (commandsWithPrefix cEnv pfx)
    CExactly s ->
      let s' = Text.unpack s in
      pure (if pfx `List.isPrefixOf` s' then [s'] else [])
    CFunction -> do
      CompletionEnv { envState = execState } <- pure cEnv
      let binds = execState Lens.^. Lens.to C.execStateContext . C.functionBindings
      let hdls = C.handleMapToHandles (C.fnBindings binds)
      pure (map (\(C.SomeHandle hdl) -> Text.unpack (W4.functionName (C.handleName hdl))) hdls)
    CPath ->
      X.try @X.IOException (completePath pfx) <&>
        \case
          Left _ -> []
          Right cs -> cs

data Completions
  = Completions
    { completionsType :: CompletionType
    , completions :: [String]
    }

-- | Used in the test suite with @. complete@
instance PP.Pretty Completions where
  pretty (Completions ty cs) = do
    let def =
          if null cs
          then PP.pretty ty
          else
            PP.pretty ty PP.<> ":" PP.<+> PP.hsep (PP.punctuate "," (map PP.pretty cs))
    case ty of
      CBreakpoint -> def
      CCommand -> PP.pretty ty
      CExactly {} ->
        if null cs
        then PP.pretty ty PP.<+> "(x)"
        else PP.pretty ty
      CFunction {} -> def
      CPath -> def

complete ::
  CompletionEnv cExt ->
  -- | Whole line
  String ->
  -- | Word being completed
  String ->
  IO [Completions]
complete cEnv input word =
  case words input of
    [] -> pure [Completions CCommand allCmdNames]
    [_] | not (null word) ->
      pure [Completions CCommand (commandsWithPrefix cEnv word)]
    (c : args) -> do
      case Cmd.parse cExts (Text.pack c) of
        Just cmd ->
          case Cmd.regex cExts cmd of
            Some rx -> do
              let (allTypes, next) = Arg.types cExts rx (map Text.pack args)
              let types =
                    case Map.lookup (Text.pack word) allTypes of
                      Nothing -> next
                      Just ts -> ts
              let cTypes = Maybe.mapMaybe (viewSome completionType) (toList types)
              Monad.mapM (\t -> Completions t <$> completionsForType cEnv word t) cTypes
        Nothing -> pure []
  where
    cExts = envCommandExt cEnv

    allCmdNames :: [String]
    allCmdNames = map (Text.unpack . Cmd.name cExts) (Cmd.allCmds cExts)