-- This file is part of Diohsc
-- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of version 3 of the GNU General Public License as
-- published by the Free Software Foundation, or any later version.
--
-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see http://www.gnu.org/licenses/.

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module RunExternal where

import           Control.Concurrent   (forkIO)
import           Control.Monad        (void)
import           Control.Monad.Catch  (bracket_, finally)
import           Control.Monad.State  (State, put, runState)
import           Data.Char            (isAlphaNum)
import           System.Environment   (setEnv, unsetEnv)
import           System.Exit          (ExitCode)
import           System.IO
import           System.IO.Temp       (withTempFile)
import           System.Process

import qualified Data.ByteString.Lazy as BL

import           ANSIColour
import           Mundanities
import           Prompt
import           Util

-- |Wrapper to ensure we don't accidentally allow use of shell commands in
-- restricted mode!
newtype RestrictedIO a = RestrictedIO (IO a)
    deriving (Functor,Applicative,Monad)

runRestrictedIO :: RestrictedIO a -> IO a
runRestrictedIO (RestrictedIO m) = m


subPercentOrAppend :: String -> String -> String
subPercentOrAppend sub str =
    let (s',subbed) = subPercent str `runState` False
    in if subbed then s' else s' ++ (' ':shellQuote sub)
    where
    -- |based on specification for $BROWSER in 'man 1 man'
    subPercent :: String -> State Bool String
    subPercent ""          = return []
    subPercent ('%':'%':s) = ('%':) <$> subPercent s
    subPercent ('%':'c':s) = (':':) <$> subPercent s
    subPercent ('%':'s':s) = put True >> (sub ++) <$> subPercent s
    subPercent (c:s)       = (c:) <$> subPercent s
    shellQuote s
        | all shellSafe s && not (null s) = s
        | otherwise = '\'' : substAll '\'' "'\\''" s <> "'"
    shellSafe c = isAlphaNum c || c `elem` ".,_-+="
    substAll c r s
        | (s',_:t) <- break (== c) s = s' <> r <> substAll c r t
        | otherwise = s

confirmShell :: String -> String -> IO Bool
confirmShell desc cmd = promptYN True False $
    desc ++ " following shell command?: " ++ withBoldStr cmd

pipeToCmdLazily :: String -> [String] -> [(String,String)] -> BL.ByteString -> RestrictedIO ()
pipeToCmdLazily cmd cArgs = pipeLazily $ proc cmd cArgs

pipeToShellLazily :: String -> [(String,String)] -> BL.ByteString -> RestrictedIO ()
pipeToShellLazily = pipeLazily . shell

filterShell :: String -> [(String,String)] -> BL.ByteString -> RestrictedIO BL.ByteString
filterShell = filterProcess . shell

withExtraEnv :: [(String,String)] -> IO a -> IO a
withExtraEnv envir = bracket_
    (mapM_ (uncurry setEnv) envir)
    (mapM_ (unsetEnv . fst) envir)

pipeLazily :: CreateProcess -> [(String,String)] -> BL.ByteString -> RestrictedIO ()
pipeLazily cp envir b = RestrictedIO . withExtraEnv envir $ do
    (Just inp, _, _, pid) <- createProcess $
        cp { std_in = CreatePipe , std_out = Inherit }
    hSetBuffering inp NoBuffering
    ignoreIOErr . finally (BL.hPut inp b) . void $ do
        hClose inp
        waitForProcess pid

filterProcess :: CreateProcess -> [(String,String)] -> BL.ByteString -> RestrictedIO BL.ByteString
filterProcess cp envir b = RestrictedIO . withExtraEnv envir $ do
    (Just inp, Just outp, _, pid) <- createProcess $
        cp { std_in = CreatePipe , std_out = CreatePipe }
    hSetBuffering inp NoBuffering
    hSetBuffering outp NoBuffering
    _ <- forkIO $ ignoreIOErr . finally (BL.hPut inp b) . void $ do
        hClose inp
        waitForProcess pid
    BL.hGetContents outp

runMailcap :: Bool -> String -> String -> String -> BL.ByteString -> RestrictedIO ()
runMailcap noConfirm action dir mimetype b =
    RestrictedIO . withTempFile dir "runtmp" $ \path h -> do
        BL.hPut h b
        hClose h
        let cArgs = ["--action=" ++ action, mimetype ++ ":" ++ path]
        rawSystem "run-mailcap" ("--norun":cArgs) >>
            (if noConfirm then return True else promptYN True False "Run this command?") >>?
            void $ rawSystem "run-mailcap" cArgs

runShellCmd :: String -> [(String,String)] -> RestrictedIO ExitCode
runShellCmd cmd envir = RestrictedIO . withExtraEnv envir $ spawnCommand cmd >>= waitForProcess

shellOnData :: Bool -> String -> FilePath -> [(String,String)] -> BL.ByteString -> RestrictedIO ()
shellOnData noConfirm cmd dir envir b = RestrictedIO . withTempFile dir "runtmp" $ \path h ->
    let cmd' = subPercentOrAppend path cmd
    in (if noConfirm then return True else confirmShell "Run" cmd') >>? void $ do
        BL.hPut h b >> hClose h
        runRestrictedIO $ runShellCmd cmd' envir