Utility.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}

module PMS.Infra.CmdRun.DS.Utility where

import System.Log.FastLogger
import qualified Control.Exception.Safe as E
import qualified Control.Concurrent as CC
import Control.Monad.IO.Class
import Control.Monad.Except
import Control.Monad.Reader
import qualified Control.Concurrent.STM as STM
import System.Exit
import Control.Lens
import System.Process
import System.IO
import Control.Concurrent.Async (async, cancel, concurrently, wait)
import qualified Data.ByteString as BS

import qualified PMS.Domain.Model.DM.Type as DM
import qualified PMS.Domain.Model.DS.Utility as DM
import PMS.Infra.CmdRun.DM.Type

-- |
--
runApp :: DM.DomainData -> AppData -> TimedFastLogger -> AppContext a -> IO (Either DM.ErrorData a)
runApp domDat appDat logger ctx =
  DM.runFastLoggerT domDat logger
    $ runExceptT
    $ flip runReaderT domDat
    $ runReaderT ctx appDat


-- |
--
liftIOE :: IO a -> AppContext a
liftIOE f = liftIO (go f) >>= liftEither
  where
    go :: IO b -> IO (Either String b)
    go x = E.catchAny (Right <$> x) errHdl

    errHdl :: E.SomeException -> IO (Either String a)
    errHdl = return . Left . show


-- |
--
toolsCallResponse :: STM.TQueue DM.McpResponse
                  -> DM.JsonRpcRequest
                  -> ExitCode
                  -> String
                  -> String
                  -> IO ()
toolsCallResponse resQ jsonRpc code outStr errStr = do
  let content = [ DM.McpToolsCallResponseResultContent "text" outStr
                , DM.McpToolsCallResponseResultContent "text" errStr
                ]
      result = DM.McpToolsCallResponseResult {
                  DM._contentMcpToolsCallResponseResult = content
                , DM._isErrorMcpToolsCallResponseResult = (ExitSuccess /= code)
                }
      resDat = DM.McpToolsCallResponseData jsonRpc result
      res = DM.McpToolsCallResponse resDat

  STM.atomically $ STM.writeTQueue resQ res


-- |
--
errorToolsCallResponse :: DM.JsonRpcRequest -> String -> AppContext ()
errorToolsCallResponse jsonRpc errStr = do
  let content = [ DM.McpToolsCallResponseResultContent "text" errStr ]
      result = DM.McpToolsCallResponseResult {
                  DM._contentMcpToolsCallResponseResult = content
                , DM._isErrorMcpToolsCallResponseResult = True
                }
      resDat = DM.McpToolsCallResponseData jsonRpc result
      res = DM.McpToolsCallResponse resDat

  resQ <- view DM.responseQueueDomainData <$> lift ask
  liftIOE $ STM.atomically $ STM.writeTQueue resQ res


-- | Run a shell command and collect stdout and stderr concurrently.
-- Reading stdout and stderr sequentially risks a deadlock when the child
-- process produces large stderr output: the OS pipe buffer fills up,
-- blocking the process before it can write to stdout, which in turn
-- prevents hGetContents hout from returning.
-- Using 'concurrently' drains both handles in separate threads so neither
-- pipe can stall the other.
runCommandBS :: String -> IO (ExitCode, BS.ByteString, BS.ByteString)
runCommandBS cmd = do
    (hout, herr, ph) <- openCommandPipes cmd
    (out, err) <- concurrently
                    (BS.hGetContents hout)  -- read stdout in one thread
                    (BS.hGetContents herr)  -- read stderr concurrently
    exitCode <- waitForProcess ph
    return (exitCode, out, err)


-- | Run a shell command with a timeout and collect stdout and stderr concurrently.
-- The timeout is handled while the ProcessHandle is still available, so the
-- child process and pipe readers can be released immediately.
runCommandBSWithTimeout :: Int -> String -> IO (ExitCode, BS.ByteString, BS.ByteString)
runCommandBSWithTimeout tout cmd = do
    (hout, herr, ph) <- openCommandPipes cmd
    aOut <- async (BS.hGetContents hout)
    aErr <- async (BS.hGetContents herr)
    ret <- waitForProcessOrTimeout tout ph
    case ret of
      Just exitCode -> do
        out <- wait aOut
        err <- wait aErr
        return (exitCode, out, err)
      Nothing -> do
        terminateProcess ph
        cancel aOut
        cancel aErr
        return (ExitFailure 1, BS.empty, "timeout occurred.")
  where
    waitForProcessOrTimeout :: Int -> ProcessHandle -> IO (Maybe ExitCode)
    waitForProcessOrTimeout timeoutMicrosec ph = go timeoutMicrosec
      where
        pollIntervalMicrosec :: Int
        pollIntervalMicrosec = 100 * 1000

        go :: Int -> IO (Maybe ExitCode)
        go remaining
          | remaining <= 0 = return Nothing
          | otherwise = do
              mExitCode <- getProcessExitCode ph
              case mExitCode of
                Just exitCode -> return $ Just exitCode
                Nothing -> do
                  let waitMicrosec = min pollIntervalMicrosec remaining
                  CC.threadDelay waitMicrosec
                  go (remaining - waitMicrosec)


openCommandPipes :: String -> IO (Handle, Handle, ProcessHandle)
openCommandPipes cmd = do
  (hin, hout, herr, ph) <- createShellProcess cmd
  hClose hin
  hSetBinaryMode hout True  -- disable newline translation (important on Windows)
  hSetBinaryMode herr True  -- disable newline translation (important on Windows)
  return (hout, herr, ph)


-- | Create a shell process for the given command string.
-- On Windows, the .bat scripts are expected to call "chcp 65001" at their
-- start to switch the code page to UTF-8, so that both built-in commands
-- (echo, move, copy, etc.) and external commands (tree.com, xcopy.exe, etc.)
-- emit UTF-8 output that can be decoded uniformly with decodeUtf8With.
createShellProcess :: String -> IO (Handle, Handle, Handle, ProcessHandle)
createShellProcess cmd = do
  let procSpec = shell cmd
  (mHin, mHout, mHerr, ph) <- createProcess procSpec {
      std_in = CreatePipe
    , std_out = CreatePipe
    , std_err = CreatePipe
    , use_process_jobs = True
    }
  case (mHin, mHout, mHerr) of
    (Just hin, Just hout, Just herr) -> return (hin, hout, herr, ph)
    _ -> E.throwString "createShellProcess: failed to create process pipes."