{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Hadolint.Shell where

import Control.Monad.Writer (Writer, execWriter, tell)
import Data.Functor.Identity (runIdentity)
import Data.Maybe (listToMaybe, mapMaybe)
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as Text
import ShellCheck.AST (Id (..), Token (..), pattern T_Pipe, pattern T_SimpleCommand)
import qualified ShellCheck.AST
import qualified ShellCheck.ASTLib
import ShellCheck.Checker (checkScript)
import ShellCheck.Interface
import qualified ShellCheck.Parser

data CmdPart = CmdPart
  { arg :: !Text,
    partId :: !Int
  }
  deriving (Show)

data Command = Command
  { name :: !Text.Text,
    arguments :: [CmdPart],
    flags :: [CmdPart]
  }
  deriving (Show)

data ParsedShell = ParsedShell
  { original :: !Text.Text,
    parsed :: !ParseResult,
    presentCommands :: ![Command]
  }

data ShellOpts = ShellOpts
  { shellName :: Text.Text,
    envVars :: Set.Set Text.Text
  }

defaultShellOpts :: ShellOpts
defaultShellOpts = ShellOpts "/bin/sh -c" defaultVars
  where
    defaultVars =
      Set.fromList
        [ "HTTP_PROXY",
          "http_proxy",
          "HTTPS_PROXY",
          "https_proxy",
          "FTP_PROXY",
          "ftp_proxy",
          "NO_PROXY",
          "no_proxy"
        ]

addVars :: [Text.Text] -> ShellOpts -> ShellOpts
addVars vars (ShellOpts n v) = ShellOpts n (v <> Set.fromList vars)

setShell :: Text.Text -> ShellOpts -> ShellOpts
setShell s (ShellOpts _ v) = ShellOpts s v

shellcheck :: ShellOpts -> ParsedShell -> [PositionedComment]
shellcheck (ShellOpts sh env) (ParsedShell txt _ _) =
  if "pwsh" `Text.isPrefixOf` sh
    then [] -- Do no run for powershell
    else runShellCheck
  where
    runShellCheck = crComments $ runIdentity $ checkScript si spec
    si = mockedSystemInterface [("", "")]
    spec =
      emptyCheckSpec
        { csFilename = "", -- filename can be ommited because we only want the parse results back
          csScript = script,
          csCheckSourced = False,
          csExcludedWarnings = exclusions,
          csShellTypeOverride = Nothing,
          csMinSeverity = StyleC
        }
    script = "#!" ++ extractShell sh ++ "\n" ++ printVars ++ Text.unpack txt
    exclusions =
      [ 2187, -- exclude the warning about the ash shell not being supported
        1090 -- requires a directive (shell comment) that can't be expressed in a Dockerfile
      ]

    extractShell s =
      maybe "" Text.unpack (listToMaybe . Text.words $ s)
    printVars = Text.unpack . Text.unlines . Set.toList $ Set.map (\v -> "export " <> v <> "=1") env

parseShell :: Text.Text -> ParsedShell
parseShell txt = ParsedShell {original = txt, parsed = parsedResult, presentCommands = commands}
  where
    parsedResult =
      runIdentity $
        ShellCheck.Parser.parseScript
          (mockedSystemInterface [("", "")])
          newParseSpec
            { psFilename = "", -- There is no filename
              psScript = "#!/bin/bash\n" ++ Text.unpack txt,
              psCheckSourced = False
            }

    commands = mapMaybe extractNames (findCommandsInResult parsedResult)
    extractNames token =
      case ShellCheck.ASTLib.getCommandName token of
        Nothing -> Nothing
        Just n -> Just $ Command (Text.pack n) allArgs (getAllFlags allArgs)
      where
        allArgs = extractAllArgs token

findCommandsInResult :: ParseResult -> [Token]
findCommandsInResult = extractTokensWith commandsExtractor
  where
    commandsExtractor = ShellCheck.ASTLib.getCommand

extractTokensWith :: forall a. (Token -> Maybe a) -> ParseResult -> [a]
extractTokensWith extractor ast =
  case prRoot ast of
    Nothing -> []
    Just script -> execWriter $ ShellCheck.AST.doAnalysis extract script
  where
    extract :: Token -> Writer [a] ()
    extract token =
      case extractor token of
        Nothing -> return ()
        Just t -> tell [t]

findPipes :: ParsedShell -> [Token]
findPipes (ParsedShell _ ast _) = extractTokensWith pipesExtractor ast
  where
    pipesExtractor pipe@T_Pipe {} = Just pipe
    pipesExtractor _ = Nothing

hasPipes :: ParsedShell -> Bool
hasPipes = not . null . findPipes

allCommands :: (Command -> Bool) -> ParsedShell -> Bool
allCommands check script = all check (presentCommands script)

noCommands :: (Command -> Bool) -> ParsedShell -> Bool
noCommands check = allCommands (not . check)

anyCommands :: (Command -> Bool) -> ParsedShell -> Bool
anyCommands check script = any check (presentCommands script)

findCommandNames :: ParsedShell -> [Text]
findCommandNames script = map name (presentCommands script)

cmdHasArgs :: Text.Text -> [Text.Text] -> Command -> Bool
cmdHasArgs expectedName expectedArgs (Command n args _)
  | expectedName /= n = False
  | otherwise = not $ null [arg | CmdPart arg _ <- args, arg `elem` expectedArgs]

cmdHasPrefixArg :: Text.Text -> Text.Text -> Command -> Bool
cmdHasPrefixArg expectedName expectedArg (Command n args _)
  | expectedName /= n = False
  | otherwise = not $ null [arg | CmdPart arg _ <- args, expectedArg `Text.isPrefixOf` arg]

extractAllArgs :: Token -> [CmdPart]
extractAllArgs (T_SimpleCommand _ _ (_ : allArgs)) = map mkPart allArgs
  where
    mkPart token =
      CmdPart
        (Text.pack . concat $ ShellCheck.ASTLib.oversimplify token)
        (mkId (ShellCheck.AST.getId token))
    mkId (Id i) = i
extractAllArgs _ = []

getArgs :: Command -> [Text.Text]
getArgs cmd = map arg (arguments cmd)

getAllFlags :: [CmdPart] -> [CmdPart]
getAllFlags = concatMap flag
  where
    flag (CmdPart arg pId)
      | arg == "--" || arg == "-" = []
      | "--" `Text.isPrefixOf` arg = [CmdPart (Text.drop 2 . Text.takeWhile (/= '=') $ arg) pId]
      | "-" `Text.isPrefixOf` arg = map (`CmdPart` pId) (Text.chunksOf 1 (Text.tail arg))
      | otherwise = []

getArgsNoFlags :: Command -> [Text.Text]
getArgsNoFlags args = map arg $ filter (notAFlagId . partId) (arguments args)
  where
    notAFlagId pId = pId `notElem` map partId (flags args)

hasFlag :: Text.Text -> Command -> Bool
hasFlag flag Command {flags} = not $ null [f | CmdPart f _ <- flags, f == flag]

hasAnyFlag :: [Text.Text] -> Command -> Bool
hasAnyFlag fs Command {flags} = not $ null [f | CmdPart f _ <- flags, f `elem` fs]

hasArg :: Text.Text -> Command -> Bool
hasArg arg Command {arguments} = not $ null [a | CmdPart a _ <- arguments, a == arg]

dropFlagArg :: [Text.Text] -> Command -> Command
dropFlagArg flagsToDrop Command {name, arguments, flags} = Command name filterdArgs flags
  where
    idsToDrop = Set.fromList [getValueId fId arguments | CmdPart f fId <- flags, f `elem` flagsToDrop]
    filterdArgs = [arg | arg@(CmdPart _ aId) <- arguments, not (aId `Set.member` idsToDrop)]

getValueId :: Int -> [CmdPart] -> Int
getValueId fId flags = foldl min (maxBound :: Int) $ filter (> fId) $ map partId flags