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

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}

module Lang.Crucible.Debug.Commands
  ( execStateSimState
  , backtrace
  , block
  , call
  , cfg
  , frame
  , help
  , load
  , prove
  , reg
  , secret
  , source
  , usage
  ) where

import Control.Exception qualified as X
import Control.Lens qualified as Lens
import Control.Monad.Except (runExceptT)
import Control.Monad.IO.Class (liftIO)
import Data.List qualified as List
import Data.Maybe qualified as Maybe
import Data.Parameterized.Classes (ixF')
import Data.Parameterized.Context qualified as Ctx
import Data.Parameterized.Nonce qualified as Nonce
import Data.Parameterized.Some (Some(Some), viewSome)
import Data.Parameterized.TraversableFC (toListFC)
import Data.Sequence qualified as Seq
import Data.Text.IO qualified as IO
import Data.Text qualified as Text
import Lang.Crucible.Analysis.Postdom qualified as C
import Lang.Crucible.Backend.Prove qualified as Prove
import Lang.Crucible.Backend qualified as C
import Lang.Crucible.CFG.Core qualified as C
import Lang.Crucible.CFG.Extension qualified as C
import Lang.Crucible.CFG.Reg qualified as C.Reg
import Lang.Crucible.CFG.SSAConversion qualified as C
import Lang.Crucible.Debug.Arg qualified as Arg
import Lang.Crucible.Debug.Command.Base qualified as BCmd
import Lang.Crucible.Debug.Command qualified as Cmd
import Lang.Crucible.Debug.Complete (CompletionT)
import Lang.Crucible.Debug.Complete qualified as Complete
import Lang.Crucible.Debug.Context (Context)
import Lang.Crucible.Debug.Context qualified as Ctxt
import Lang.Crucible.Debug.Inputs (Inputs)
import Lang.Crucible.Debug.Inputs qualified as Inps
import Lang.Crucible.Debug.Outputs qualified as Outps
import Lang.Crucible.Debug.Regex qualified as Rgx
import Lang.Crucible.Debug.Response qualified as Resp
import Lang.Crucible.Debug.Response (Response)
import Lang.Crucible.Debug.Statement qualified as Stmt
import Lang.Crucible.Debug.Statement (Statement)
import Lang.Crucible.Debug.Style qualified as Style
import Lang.Crucible.Debug.Style (StyleT)
import Lang.Crucible.FunctionHandle qualified as C
import Lang.Crucible.Pretty (ppRegVal)
import Lang.Crucible.Simulator.CallFrame qualified as C
import Lang.Crucible.Simulator.EvalStmt qualified as C
import Lang.Crucible.Simulator.ExecutionTree qualified as C
import Lang.Crucible.Simulator qualified as C
import Lang.Crucible.Syntax.Atoms qualified as C
import Lang.Crucible.Syntax.Concrete qualified as C
import Lang.Crucible.Syntax.SExpr qualified as C
import Lang.Crucible.Utils.Seconds qualified as Sec
import Lang.Crucible.Utils.Timeout qualified as CTO
import Prettyprinter qualified as PP
import System.Exit qualified as Exit
import Text.Megaparsec qualified as MP
import What4.Config qualified as W4
import What4.Expr.Builder qualified as W4
import What4.FunctionName qualified as W4
import What4.Interface qualified as W4
import What4.ProgramLoc qualified as W4
import What4.Solver qualified as W4

---------------------------------------------------------------------
-- Helpers, exported

execStateSimState ::
  C.ExecState p sym ext r ->
  Either Resp.NotApplicable (C.SomeSimState p sym ext r)
execStateSimState =
  \case
    C.ResultState _                  -> Left Resp.DoneSimulating
    C.AbortState _ st                -> Right (C.SomeSimState st)
    C.UnwindCallState _ _ st         -> Right (C.SomeSimState st)
    C.CallState _ _ st               -> Right (C.SomeSimState st)
    C.TailCallState _ _ st           -> Right (C.SomeSimState st)
    C.ReturnState _ _ _ st           -> Right (C.SomeSimState st)
    C.ControlTransferState _ st      -> Right (C.SomeSimState st)
    C.RunningState _ st              -> Right (C.SomeSimState st)
    C.SymbolicBranchState _ _ _ _ st -> Right (C.SomeSimState st)
    C.OverrideState _ st             -> Right (C.SomeSimState st)
    C.BranchMergeState _ st          -> Right (C.SomeSimState st)
    C.InitialState _ _ _ _ _         -> Left Resp.NotYetSimulating

---------------------------------------------------------------------
-- Helpers, not exported

insertCfg ::
  C.IsSyntaxExtension ext =>
  C.Reg.CFG ext blocks init ret ->
  C.FnHandleMap (C.FnState p sym ext) ->
  C.FnHandleMap (C.FnState p sym ext)
insertCfg regCFG hdlMap =
  case C.toSSA regCFG of
    C.SomeCFG ssaCFG ->
      C.insertHandleMap
        (C.cfgHandle ssaCFG)
        (C.UseCFG ssaCFG (C.postdomInfo ssaCFG))
        hdlMap

data SomeFn f = forall args ret. SomeFn (f args ret)

lookupHandleMapByName ::
  W4.FunctionName ->
  C.FnHandleMap f ->
  Maybe (SomeFn f)
lookupHandleMapByName fNm hdls = do
  let sHdls = C.handleMapToHandles hdls
  C.SomeHandle h <- List.find (\(C.SomeHandle hdl) -> C.handleName hdl == fNm) sHdls
  SomeFn <$> C.lookupHandleMap h hdls

-- TODO: Upstream to Crucible as a Lens
modifySimContext ::
  (C.SimContext p sym ext -> C.SimContext p sym ext) ->
  C.ExecState p sym ext r ->
  C.ExecState p sym ext r
modifySimContext f =
  \case
    C.ResultState r ->
      case r of
        C.FinishedResult ctx x -> C.ResultState (C.FinishedResult (f ctx) x)
        C.AbortedResult ctx x -> C.ResultState (C.AbortedResult (f ctx) x)
        C.TimeoutResult es -> modifySimContext f es
    C.AbortState x s -> C.AbortState x (s Lens.& C.stateContext Lens.%~ f)
    C.UnwindCallState x y s -> C.UnwindCallState x y (s Lens.& C.stateContext Lens.%~ f)
    C.CallState x y s -> C.CallState x y (s Lens.& C.stateContext Lens.%~ f)
    C.TailCallState x y s -> C.TailCallState x y (s Lens.& C.stateContext Lens.%~ f)
    C.ReturnState x y z s -> C.ReturnState x y z (s Lens.& C.stateContext Lens.%~ f)
    C.ControlTransferState x s -> C.ControlTransferState x (s Lens.& C.stateContext Lens.%~ f)
    C.RunningState x s -> C.RunningState x (s Lens.& C.stateContext Lens.%~ f)
    C.SymbolicBranchState x y z w s -> C.SymbolicBranchState x y z w (s Lens.& C.stateContext Lens.%~ f)
    C.OverrideState x s -> C.OverrideState x (s Lens.& C.stateContext Lens.%~ f)
    C.BranchMergeState x s -> C.BranchMergeState x (s Lens.& C.stateContext Lens.%~ f)
    C.InitialState ctx x y z w -> C.InitialState (f ctx) x y z w

---------------------------------------------------------------------
-- Commands, exported

backtrace ::
  C.ExecState p sym ext r ->
  IO (Response cExt)
backtrace execState =
  case execStateSimState execState of
    Left notApplicable -> pure (Resp.UserError (Resp.NotApplicable notApplicable))
    Right (C.SomeSimState simState) -> do
      let frames = C.activeFrames (C._stateTree simState)
      pure (Resp.Backtrace (C.ppExceptionContext frames))

block ::
  C.IsSyntaxExtension ext =>
  W4.IsExpr (W4.SymExpr sym) =>
  Context cExt sym ext t ->
  C.ExecState p sym ext r ->
  IO (Response cExt)
block _ctx execState =
  case execStateSimState execState of
    Left notApplicable -> pure (Resp.UserError (Resp.NotApplicable notApplicable))
    Right (C.SomeSimState simState) ->
      let fr = simState Lens.^. C.stateTree . C.actFrame . C.gpValue in
      case fr of
        C.MF callFrame ->
          let regs = C.regMap (callFrame Lens.^. C.frameRegs) in
          pure (Resp.Block (ppStmtSeq False (Ctx.size regs) (C._frameStmts callFrame)))
        C.OF {} -> pure (Resp.UserError (Resp.NotApplicable Resp.NoCfg))
        C.RF {} -> pure (Resp.UserError (Resp.NotApplicable Resp.NoCfg))
  where
    -- TODO: The below is taken from Crucible. Export?

    prefixLineNum :: Bool -> W4.ProgramLoc -> PP.Doc ann -> PP.Doc ann
    prefixLineNum True pl d = PP.vcat ["%" PP.<+> W4.ppNoFileName (W4.plSourceLoc pl), d]
    prefixLineNum False _ d = d

    ppStmtSeq :: C.PrettyExt ext => Bool -> Ctx.Size ctx -> C.StmtSeq ext blocks ret ctx -> PP.Doc ann
    ppStmtSeq ppLineNum h (C.ConsStmt pl s r) =
      PP.vcat
      [ prefixLineNum ppLineNum pl (C.ppStmt h s)
      , ppStmtSeq ppLineNum (C.nextStmtHeight h s) r
      ]
    ppStmtSeq ppLineNum _ (C.TermStmt pl s) =
      prefixLineNum ppLineNum pl (PP.pretty s)

call ::
  C.IsSymInterface sym =>
  C.IsSyntaxExtension ext =>
  C.ExecState p sym ext (C.RegEntry sym t) ->
  C.TypeRepr t ->
  W4.FunctionName ->
  IO (C.ExecutionFeatureResult p sym ext (C.RegEntry sym t), Response cExt)
call execState rTy fNm = do
  let fTy = C.FunctionHandleRepr Ctx.empty C.UnitRepr
  let binds = execState Lens.^. Lens.to C.execStateContext . C.functionBindings
  case C.searchHandleMap fNm fTy (C.fnBindings binds) of
    Nothing -> do
      let hdls = C.handleMapToHandles (C.fnBindings binds)
      let names = map (\(C.SomeHandle hdl) -> C.handleName hdl) hdls
      pure (C.ExecutionFeatureNoChange, Resp.UserError (Resp.FnNotFound fNm names))
    Just (hdl, bind) -> do
      case execStateSimState execState of
        Left notApplicable -> pure (C.ExecutionFeatureNoChange, Resp.UserError (Resp.NotApplicable notApplicable))
        Right (C.SomeSimState simState) -> do
          let initCtx = simState Lens.^. C.stateContext
          let globs = simState Lens.^. C.stateGlobals
          let abortHdl = simState Lens.^. C.abortHandler
          let initState =
                C.InitialState initCtx globs abortHdl rTy $
                  C.runOverrideSim rTy $
                    let args = C.RegMap Ctx.empty in
                    case bind of
                      C.UseOverride o -> do
                        _ <- C.callOverride hdl o args
                        C.exitExecution Exit.ExitSuccess
                      C.UseCFG c _ -> do
                        _ <- C.callCFG c args
                        C.exitExecution Exit.ExitSuccess
          let featRes = C.ExecutionFeatureNewState initState
          pure (featRes, Resp.Ok)

cfg ::
  C.IsSyntaxExtension ext =>
  C.ExecState p sym ext r ->
  Maybe W4.FunctionName ->
  IO (Response cExt)
cfg execState =
  \case
    Just fNm -> do
      let binds = execState Lens.^. Lens.to C.execStateContext . C.functionBindings
      let hdls = C.handleMapToHandles (C.fnBindings binds)
      let names = map (\(C.SomeHandle hdl) -> C.handleName hdl) hdls
      case lookupHandleMapByName fNm (C.fnBindings binds) of
        Nothing -> pure (Resp.UserError (Resp.FnNotFound fNm names))
        Just (SomeFn (C.UseOverride {})) -> pure (Resp.UserError (Resp.NotACfg fNm))
        Just (SomeFn (C.UseCFG c _)) -> pure (Resp.Cfg (C.ppCFG True c))
    Nothing ->
      case execStateSimState execState of
        Left notApplicable -> pure (Resp.UserError (Resp.NotApplicable notApplicable))
        Right (C.SomeSimState simState) -> do
          case simState Lens.^. C.stateTree . C.actFrame . C.gpValue of
            C.MF (C.CallFrame { C._frameCFG = c }) ->
              pure (Resp.Cfg (C.ppCFG True c))
            C.OF {} -> pure (Resp.UserError (Resp.NotApplicable Resp.NoCfg))
            C.RF {} -> pure (Resp.UserError (Resp.NotApplicable Resp.NoCfg))

frame ::
  C.IsSyntaxExtension ext =>
  W4.IsExpr (W4.SymExpr sym) =>
  Context cExt sym ext t ->
  C.ExecState p sym ext r ->
  IO (Response cExt)
frame _ctx execState =
  case execStateSimState execState of
    Left notApplicable -> pure (Resp.UserError (Resp.NotApplicable notApplicable))
    Right (C.SomeSimState simState) ->
      pure $ Resp.Frame $
        let fr = simState Lens.^. C.stateTree . C.actFrame . C.gpValue in
        let name = fr Lens.^. C.frameFunctionName in
        case fr of
          C.OF ovFrame ->
            let regs = C.regMap (ovFrame Lens.^. C.overrideRegMap) in
            PP.vcat
            [ "Override:" PP.<+> PP.pretty name
            , ""
            , "Argument types:"
            , PP.vcat $ toListFC (bullet . PP.pretty . C.regType) regs
            ]
          C.MF callFrame ->
            let regs = C.regMap (callFrame Lens.^. C.frameRegs) in
            PP.vcat
              [ "CFG:" PP.<+> PP.pretty name
              , "Block:" PP.<+> viewSome PP.pretty (C._frameBlockID callFrame)
              , ""
              , "Argument types:"
              , PP.vcat $ toListFC (bullet . PP.pretty . C.regType) regs
              ]
          C.RF {} -> "Return:" PP.<+> PP.pretty name
  where bullet = ("-" PP.<+>) . PP.align

help ::
  C.IsSyntaxExtension ext =>
  (?parserHooks :: C.ParserHooks ext) =>
  PP.Pretty cExt =>
  Cmd.CommandExt cExt ->
  Maybe (Cmd.Command cExt) ->
  Response cExt
help cExts =
  \case
    Nothing ->
      let mkHelp c = Cmd.name cExts c <> " (" <> Cmd.abbrev cExts c <> ")" <> ": " <> Cmd.help cExts c in
      Resp.Help (PP.pretty (Text.unlines (List.sort (map mkHelp (Cmd.allCmds cExts)))))
    Just cmd -> 
      Resp.Help $
        -- See 'Lang.Crucible.Debug.Command.ext{Detail,Help}' for expectations
        -- on these strings, which are why it makes sense to format them this
        -- way (e.g., adding a period, printing "details" after "help"), etc.
        let basic = [usage cExts cmd, "", PP.pretty (Cmd.help cExts cmd) <> "."]
            detail =
              case Cmd.detail cExts cmd of
                Just more -> ["", PP.pretty more]
                Nothing -> []
        in PP.vcat (basic ++ detail)

load ::
  C.IsSyntaxExtension ext =>
  (?parserHooks :: C.ParserHooks ext) =>
  C.ExecState p sym ext r ->
  FilePath ->
  IO (C.ExecutionFeatureResult p sym ext r, Response cExt)
load execState0 path = do
  eTxt <- X.try (IO.readFile path)
  case eTxt of
    Left ioErr -> pure (C.ExecutionFeatureNoChange, Resp.IOError ioErr)
    Right txt -> do
      case MP.parse (C.skipWhitespace *> MP.many (C.sexp C.atom) <* MP.eof) path txt of
        Left err -> pure (C.ExecutionFeatureNoChange, Resp.UserError (Resp.LoadParseError err))
        Right v ->
          Nonce.withIONonceGenerator $ \ng -> do
            let execCtx = C.execStateContext execState0
            let ha = C.simHandleAllocator execCtx
            parseResult <- C.top ng ha [] $ C.prog v
            case parseResult of
              Left err ->
                pure (C.ExecutionFeatureNoChange, Resp.UserError (Resp.LoadSyntaxError (PP.pretty err)))
              Right (C.ParsedProgram{ C.parsedProgCFGs = cfgs }) -> do
                let binds = execState0 Lens.^. Lens.to C.execStateContext . C.functionBindings . Lens.to C.fnBindings
                let binds' = foldl (\bs (C.Reg.AnyCFG c) -> insertCfg c bs) binds cfgs
                let execState = modifySimContext (\c -> c Lens.& C.functionBindings Lens..~ C.FnBindings binds') execState0
                let featRes = C.ExecutionFeatureModifiedState execState
                pure (featRes, Resp.Load path (length cfgs))

prove ::
  (sym ~ W4.ExprBuilder t st fs) =>
  C.IsSymBackend sym bak =>
  bak ->
  IO (Response cExt)
prove bak = do
  let sym = C.backendGetSym bak
  -- TODO(#270?): Configurable timeout
  let timeout = CTO.Timeout (Sec.secondsFromInt 5)
  -- TODO(#266): Configurable prover
  W4.extendConfig W4.z3Options (W4.getConfiguration sym)
  let prover = Prove.offlineProver timeout sym W4.defaultLogData W4.z3Adapter
  let strat = Prove.ProofStrategy prover Prove.keepGoing
  let convertResult =
        \case
          Prove.Proved {} -> Resp.Proved
          Prove.Disproved {} -> Resp.Disproved
          Prove.Unknown {} -> Resp.Unknown
  let printer = Prove.ProofConsumer $ \goal res -> do
        doc <- C.ppProofObligation sym goal
        pure (Seq.fromList [(doc, convertResult res)])
  runExceptT (Prove.proveCurrentObligations bak strat printer) >>=
    \case
      Left CTO.TimedOut -> pure Resp.ProveTimeout
      Right results -> pure (Resp.Prove results)

-- TODO: Support taking a function or block name
reg ::
  forall p sym ext r t cExt.
  W4.IsExpr (W4.SymExpr sym) =>
  Context cExt sym ext t ->
  C.ExecState p sym ext r ->
  [Int] ->
  IO (Response cExt)
reg ctx execState idxs =
  case execStateSimState execState of
    Left notApplicable -> pure (Resp.UserError (Resp.NotApplicable notApplicable))
    Right (C.SomeSimState simState) -> do
      let fr = simState Lens.^. C.stateTree . C.actFrame . C.gpValue
      let regs :: Maybe (C.Some (C.RegMap sym))
          regs =
            case fr of
              C.OF ovFrame ->
                Just (C.Some (ovFrame Lens.^. C.overrideRegMap))
              C.MF callFrame ->
                Just (C.Some (callFrame Lens.^. C.frameRegs))
              C.RF {} ->
                Nothing
      case regs of
        Nothing -> pure (Resp.UserError (Resp.NotApplicable Resp.NoFrame))
        Just (C.Some (C.RegMap rs)) -> do
          let sz = Ctx.size rs
          let mIdxs =
                if List.null idxs
                then Right (toListFC Some (Ctx.generate sz id))  -- all indices
                else traverse (\i -> maybeToEither i (Ctx.intIndex i sz)) idxs
          let ppArg (C.Some idx) =
                let C.RegEntry ty val = rs Lens.^. ixF' idx in
                ppRegVal (Ctxt.dbgPpIntrinsic ctx) ty (C.RV val)
          case mIdxs of
            Left i -> pure (Resp.UserError (Resp.NoSuchArgument i))
            Right idxs' -> pure (Resp.Arg (PP.vcat (List.map ppArg idxs')))
  where
    maybeToEither :: e -> Maybe a -> Either e a
    maybeToEither e =
      \case
        Nothing -> Left e
        Just a -> Right a

secret ::
  PP.Pretty cExt =>
  Context cExt sym ext t ->
  C.ExecState p sym ext r ->
  Rgx.Match (Arg.Arg cExt) BCmd.SecretRegex ->
  IO (Resp.Response cExt)
secret ctx execState m =
  let cExt = Ctxt.dbgCommandExt ctx in
  case m of
    -- Show what sort of tab-completion should be provided by the UI
    Rgx.MLeft (Rgx.MThen _complete (Rgx.MThen (Rgx.MLit (Arg.ACommand cmd)) (Rgx.MStar args0))) -> do
      let args = map (\(Rgx.MLit (Arg.AText t)) -> Text.unpack t) args0
      let cEnv = Ctxt.toCompletionEnv ctx execState
      let line = unwords (show (PP.pretty cmd) : map (filter (/= '_')) args)
      -- Act like we are completing as if the user pressed TAB at the "_"
      let cursor = Maybe.fromMaybe "_" (List.find ("_" `List.isSuffixOf`) args)
      completions <- Complete.complete cEnv line (List.init cursor)
      pure $
        Resp.Complete $
          case completions of
            [] -> "No completions"
            _ -> PP.vcat (List.map PP.pretty completions)
    -- Show what sort of syntax-highlighting should be provided by the UI
    Rgx.MRight (Rgx.MThen _secret (Rgx.MThen (Rgx.MLit (Arg.ACommand cmd)) (Rgx.MStar args0))) ->
      let args = map (\(Rgx.MLit (Arg.AText t)) -> t) args0 in
      case Cmd.regex cExt cmd of
        Some rx ->
          let sEnv = Ctxt.toStyleEnv ctx execState in
          pure (Resp.Style (Style.runStyle sEnv (Style.style rx args)))

source ::
  Context cExt sym ext t ->
  FilePath ->
  IO (Either X.IOException (Inputs (CompletionT cExt (StyleT cExt IO)) (Statement cExt)))
source ctx path = do
  let cExt = Ctxt.dbgCommandExt ctx
  eTxt <- liftIO (X.try (IO.readFile path))
  case eTxt of
    Left ioErr -> pure (Left ioErr)
    Right txt -> do
      let lns = Text.lines txt
      let cmds = traverse (Stmt.parse cExt) lns
      case cmds of
        Left err -> do
          Outps.send (Ctxt.dbgOutputs ctx) (Resp.UserError (Resp.SourceParseError err))
          pure (Right (Ctxt.dbgInputs ctx))
        Right stmts -> Right <$> Inps.prepend stmts (Ctxt.dbgInputs ctx)

usage ::
  PP.Pretty cExt =>
  Cmd.CommandExt cExt ->
  Cmd.Command cExt ->
  PP.Doc ann
usage cExt cmd =
  case Cmd.regex cExt cmd of
    Some rx -> PP.vcat (map (PP.pretty cmd PP.<+>) (Arg.usage rx))