content: {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  import qualified Cartel
  import qualified Cartel.Ast
  import qualified Cartel.Render
  import Control.Applicative ((<$>), (<*>))
  import Control.Monad (mapM_, liftM, when)
  import Control.Exception (catch, throw, SomeException)
  import Crypto.Hash (hash, Digest, MD5)
  import Data.Aeson (Value(..), ToJSON(..), FromJSON(..), encode, decode)
  import qualified Data.ByteString as BS
  import qualified Data.ByteString.Char8 as B8
  import qualified Data.ByteString.Lazy as BL
  import qualified Data.ByteString.UTF8 as BU
  import Data.Char (isDigit)
  import Data.Data (Data)
  import Data.Default (def)
  import Data.Generics.Uniplate.Data (universeBi, transformBi)
  import Data.List (intercalate, find, filter, nub, any, takeWhile, isPrefixOf, unionBy)
  import qualified Data.Map.Strict as Map
  import Data.Map.Lazy (foldrWithKey)
  import Data.Maybe (fromMaybe, catMaybes, isJust)
  import Data.Monoid (mempty)
  import qualified Data.Text as Text
  import Data.Text.Encoding (encodeUtf8, decodeUtf8)
  import Data.Typeable (Typeable)
  import qualified Data.Version as V
  import Distribution.Hackage.DB (readHackage', hackagePath)
  import Distribution.PackageDescription (condLibrary, condTreeData, exposedModules)
  import Distribution.Text (display)
  import GHC.Generics (Generic)
  import Language.Haskell.Exts (parseFileContentsWithMode, ParseMode(..), defaultParseMode, Extension(..), KnownExtension(..), ParseResult(..), fromParseResult)
  import Language.Haskell.Exts.Fixity (applyFixities, infix_, infixl_, infixr_)
  import Language.Haskell.Exts.Syntax (Module(..), ModuleName(..), ModulePragma(..), ImportDecl(..), QName(..), Name(..), Exp(..), Stmt(..), Type(..), SrcLoc(..), QOp(..))
  import Language.Haskell.Exts.Pretty (prettyPrintWithMode, defaultMode, PPHsMode(linePragmas))
  import System.Console.CmdArgs (cmdArgs, (&=), help, typ, typDir, args, summary, name)
  import System.Directory (getTemporaryDirectory, createDirectoryIfMissing, doesFileExist)
  import System.Exit (ExitCode(..), exitFailure)
  import System.FilePath ((</>), (<.>), replaceFileName, takeBaseName)
  import System.IO (hPutStrLn, hPutStr, stderr, writeFile, hGetContents, Handle)
  import System.Posix.Process (executeFile)
  import System.Process (ProcessHandle, waitForProcess, createProcess, CreateProcess(..), shell, proc, StdStream(..))
  
  import Hesh.Process (pipeOps)
  import Hesh.Shell (desugar)
  
  data Hesh = Hesh {stdin_ :: Bool
                   ,no_sugar :: Bool
                   ,no_type_hints :: Bool
                   ,compile_only :: Bool
                   ,source_dir :: String
                   ,source_only :: Bool
                   ,verbose :: Bool
                   ,package_index :: String
                   ,module_index :: String
                   ,args_ :: [String]
                   } deriving (Data, Typeable, Show, Eq)
  
  hesh = Hesh {stdin_ = False &= help "If this option is present, or if no arguments remain after option process, then the script is read from standard input."
              ,no_sugar = False &= help "Don't expand syntax shortcuts."
              ,no_type_hints = False &= name "t" &= help "Don't add automatic type hints."
              ,compile_only = False &= help "Compile the script but don't run it."
              ,source_dir = def &= name "o" &= typDir &= help "Write to given directory instead of a temporary directory."
              ,source_only = False &= name "x" &= help "Generate expanded source code and cabal file only."
              ,verbose = False &= help "Display more information, such as Cabal output."
              ,package_index = def &= name "p" &= help "The cabal package index (normally ~/.cabal/packages/hackage.haskell.org/00-index.tar)"
              ,module_index = def &= name "m" &= help "The JSON package index (generated by hesh, normally ~/.cabal/packages/hackage.haskell.org/modules.json)"
              ,args_ = def &= args &= typ "FILE|ARG.."
              } &=
         help "Build/run a hesh script." &=
         summary ("Hesh v" ++ intercalate "." (map show heshCartelVersion))
  
  heshCartelVersion = [1, 14, 0]
  
  main = do
    opts <- cmdArgs hesh
    -- In order to work in shebang mode and to provide a more familiar
    -- commandline interface, we support taking the input filename as an
    -- argument (rather than just relying on the script being provided
    -- on stdin). If no commandline argument is provided, we assume the
    -- script is on stdin.
    let scriptFile = if stdin_ opts || null (args_ opts) then "<stdin>" else head (args_ opts)
        scriptName = if stdin_ opts || null (args_ opts) then "script" else takeBaseName (head (args_ opts))
    (source, args) <- if stdin_ opts || null (args_ opts)
                        then (\x -> (x, args_ opts)) `fmap` BS.getContents
                        else (\x -> (x, tail (args_ opts))) `fmap` BS.readFile (head (args_ opts))
    -- First, get the module -> package+version lookup table.
    hackagePath <- if null (package_index opts)
                     then hackagePath
                     else return (package_index opts)
    let modulePath = if null (module_index opts) then replaceFileName hackagePath "modules.json" else module_index opts
    modules <- fromFileCache modulePath (modulePackages hackagePath)
    let ast = (if no_type_hints opts then id else defaultToUnit) (parseScript scriptFile (no_sugar opts) source)
        -- Find all qualified module names to add module names to the import list (qualified).
        names = qualifiedNamesFromModule ast
        -- Find any import references.
        imports = importsFromModule ast
        -- Remove aliased modules from the names.
        aliases = catMaybes (map (\ (_, y, _) -> y) imports)
        fqNames = filter (`notElem` aliases) (map fst names)
        -- Insert qualified module usages back into the import list.
        (Module a b pragmas d e importDecls g) = ast
        expandedImports = importDecls ++ map importDeclQualified fqNames ++ if no_sugar opts then [] else (if isJust (find (\(m, _, _) -> m == "Hesh") imports) then [] else [importDeclUnqualified "Hesh"]) ++ if no_type_hints opts then [] else [importDeclQualified "Control.Monad.IO.Class"]
        expandedPragmas = pragmas ++ if no_sugar opts then [] else sugarPragmas
        expandedAst = Module a b expandedPragmas d e expandedImports g
        -- From the imports, build a list of necessary packages.
        -- First, remove fully qualified names that were previously
        -- imported explicitly. This is so that we don't override the
        -- package that might have been selected for that module
        -- manually in the import statement.
        packages = map (packageFromModules modules) (unionBy (\ (x, _, _) (x', _, _) -> x == x') imports (map fqNameModule fqNames) ++ if no_sugar opts then [] else [("Hesh", Nothing, Just "hesh")])
        source' = BU.fromString $ prettyPrintWithMode (defaultMode { linePragmas = True }) expandedAst
        cabal' = BU.fromString $ Cartel.Render.renderNoIndent (cartel opts packages scriptName)
    dir <- if null (source_dir opts)
             then (</> ("hesh-" ++ show (hash source' :: Digest MD5))) `liftM` getTemporaryDirectory
             else return (source_dir opts)
    let sourcePath = dir </> "Main.hs"
        cabalPath = dir </> scriptName <.> "cabal"
    -- Check if the source is different than what exists in the directory already.
    oldSource <- catchAny (BS.readFile sourcePath) (\ _ -> return "")
    oldCabal <- catchAny (BS.readFile cabalPath) (\ _ -> return "")
    when (oldSource /= source' || oldCabal /= cabal') $ do
      createDirectoryIfMissing False dir
      BS.writeFile sourcePath source'
      BS.writeFile cabalPath cabal'
      -- Cabal will complain without a LICENSE file.
      BS.writeFile (dir </> "LICENSE") ""
    if (source_only opts)
      then putStr dir
      else do
        let path = dir </> "dist/build" </> scriptName </> scriptName
        binaryExists <- doesFileExist path
        when (not binaryExists || oldSource /= source' || oldCabal /= cabal') $ do
          callCommandInDir "cabal install --only-dependencies" dir (verbose opts)
          callCommandInDir "cabal build" dir (verbose opts)
        if (compile_only opts)
          then putStr path
          else executeFile path False args Nothing
   where fqNameModule name = (name, Nothing, Nothing)
         sugarPragmas = [LanguagePragma (SrcLoc "<generated>" 0 0) [Ident "TemplateHaskell", Ident "QuasiQuotes", Ident "PackageImports"]]
  
  waitForSuccess :: String -> ProcessHandle -> Maybe Handle -> IO ()
  waitForSuccess cmd p out = do
    code <- waitForProcess p
    when (code /= ExitSuccess) $ do
      case out of
        Just o -> hPutStr stderr =<< hGetContents o
        Nothing -> return ()
      hPutStrLn stderr $ cmd ++ ": exited with code " ++ show code
      exitFailure
  
  callCommandInDir :: String -> FilePath -> Bool -> IO ()
  callCommandInDir cmd dir verbose' = do
    -- GHC is noisy on stdout. It should go to stderr instead.
    -- TODO: Move this into a more appropriate place. It just
    -- happens to work here.
    (_, out, _, p) <- createProcess (shell cmd) { cwd = Just dir, std_out = (if verbose' then UseHandle stderr else CreatePipe) }
    waitForSuccess cmd p out
  
  callCommand :: FilePath -> [String] -> IO ()
  callCommand path args = do
    (_, _, _, p) <- createProcess (proc path args)
    waitForSuccess path p Nothing
  
  catchAny :: IO a -> (SomeException -> IO a) -> IO a
  catchAny = catch
  
  -- Read a value from a cache (JSON-encoded file), writing it out if
  -- the cached value doesn't exist or is invalid.
  fromFileCache :: (FromJSON a, ToJSON a) => FilePath -> IO a -> IO a
  fromFileCache path f = do
    catchAny readCached $ \ _ -> do
      value <- f
      catchAny (BL.writeFile path (encode value)) (\ _ -> return ())
      return value
   where readCached = do
           d <- BL.readFile path
           case decode d of
             Just x -> return x
             Nothing -> error "Failed to parse JSON."
  
  qualifiedNamesFromModule :: Module -> [(String, String)]
  qualifiedNamesFromModule m = [ (mName, name) | Qual (ModuleName mName) (Ident name) <- universeBi m ]
  
  -- This is a helper until Haskell has better defaulting (something like
  -- https://ghc.haskell.org/trac/haskell-prime/wiki/Defaulting).
  -- It saves adding "IO ()" type declarations in a number of common contexts.
  defaultToUnit :: Module -> Module
  defaultToUnit = transformBi defaultExpToUnit
   where defaultExpToUnit :: Exp -> Exp
         defaultExpToUnit (Do stmts) = Do (map defaultStmtToUnit stmts)
         defaultExpToUnit exp = exp
         defaultStmtToUnit (Qualifier exp) = Qualifier (defaultQualifiedExpToUnit exp)
         defaultStmtToUnit stmt = stmt
         defaultQualifiedExpToUnit exp = if canDefaultToUnit exp then defaultToUnit exp else exp
         -- do [sh| ... |] => do [sh| ... |] :: IO ()
         canDefaultToUnit (QuasiQuote f _) = f `elem` ["sh", "Hesh.sh"]
         -- do cmd "..." [...] => do cmd "..." [...] :: IO ()
         canDefaultToUnit (App (App (Var f) _) _) = f `elem` functionNames "cmd"
         -- do ... /> "..." => ... /> "..." :: IO ()
         canDefaultToUnit (InfixApp exp1 (QVarOp op) exp2) = op `elem` (concatMap operatorNames pipeOps)
         canDefaultToUnit _ = False
         defaultToUnit exp = ExpTypeSig (SrcLoc "<generated>" 0 0) exp (TyVar (Ident "(Control.Monad.IO.Class.MonadIO m) => m ()"))
         functionNames name = [ UnQual (Ident name)
                              , Qual (ModuleName "Hesh") (Ident name) ]
         operatorNames name = [ UnQual (Symbol name)
                              , UnQual (Ident ("(" ++ name ++ ")"))
                              , Qual (ModuleName "Hesh") (Ident ("(" ++ name ++ ")")) ]
  
  importsFromModule :: Module -> [(String, Maybe String, Maybe String)]
  importsFromModule (Module _ _ _ _ _ imports _) = map importName imports
   where importName (ImportDecl _ (ModuleName m) _ _ _ pkg Nothing _) = (m, Nothing, pkg)
         importName (ImportDecl _ (ModuleName m) _ _ _ pkg (Just (ModuleName n)) _) = (m, Just n, pkg)
  
  importDeclQualified :: String -> ImportDecl
  importDeclQualified m = ImportDecl (SrcLoc "<generated>" 0 0) (ModuleName m) True False False Nothing Nothing Nothing
  
  importDeclUnqualified :: String -> ImportDecl
  importDeclUnqualified m = ImportDecl (SrcLoc "<generated>" 0 0) (ModuleName m) False False False Nothing Nothing Nothing
  
  packageFromModules modules (m, _, Just pkg)
    | pkg == "hesh" = Cartel.package "hesh" (Cartel.eq heshCartelVersion)
    | otherwise =
        let parts = Text.splitOn "-" (Text.pack pkg)
        in case parts of
             [_] -> Cartel.package pkg Cartel.anyVersion
             ps -> if Text.all (\c -> isDigit c || c == '.') (last ps)
                     then let version = map (read . Text.unpack) (Text.splitOn "." (last ps))
                              package = Text.unpack (Text.intercalate "-" (init ps))
                          in Cartel.package package (Cartel.eq version)
                     else Cartel.package pkg Cartel.anyVersion
  packageFromModules modules (m, _, Nothing)
    | m == "Hesh" || isPrefixOf "Hesh." m = Cartel.package "hesh" (Cartel.eq heshCartelVersion)
    | otherwise = constrainedPackage (Map.findWithDefault (error ("Module \"" ++ m ++ "\" not found in Hackage list.")) (Text.pack m) modules)
  
  cartel opts packages name = mempty { Cartel.Ast.properties = properties
                                     , Cartel.Ast.sections = [executable] }
   where properties = mempty
           { Cartel.name         = name
           , Cartel.version      = [0,1]
           , Cartel.cabalVersion = Just (fromIntegral 1, fromIntegral 18)
           , Cartel.buildType    = Just Cartel.simple
           , Cartel.license      = Just Cartel.allRightsReserved
           , Cartel.licenseFile  = "LICENSE"
           , Cartel.category     = "shell"
           }
         executable = Cartel.executable name fields
         fields = [ Cartel.Ast.ExeMainIs "Main.hs"
                  , Cartel.Ast.ExeInfo (Cartel.Ast.DefaultLanguage Cartel.Ast.Haskell2010)
                  , Cartel.Ast.ExeInfo (Cartel.Ast.BuildDepends ([Cartel.package "base" Cartel.anyVersion] ++ packages))
                  , Cartel.Ast.ExeInfo (Cartel.Ast.GHCOptions (["-threaded"]))
                  ]
  
  -- We make the simplifying assumption that a module only appears in a
  -- contiguous version range.
  data PackageConstraint = PackageConstraint { packageName :: Text.Text
                                             , packageMinVersion :: [Int]
                                             , packageMaxVersion :: [Int]
                                             } deriving Generic
  
  instance ToJSON PackageConstraint
  instance FromJSON PackageConstraint
  
  -- PackageConstraint -> Package
  -- Always prefer base, otherwise arbitrarily take the first module.
  constrainedPackage ps = Cartel.package (Text.unpack (packageName package)) Cartel.anyVersion
   where package = case find (\p -> packageName p == (Text.pack "base")) ps of
                     Just p' -> p'
                     Nothing -> head ps
  
  modulePackages hackagePath = foldrWithKey buildConstraints Map.empty `liftM` readHackage' hackagePath
   where buildConstraints name versions constraints = foldrWithKey (buildConstraints' name) constraints versions
         buildConstraints' name version meta constraints = foldr (\m cs -> Map.alter (alterConstraint (Text.pack name) (V.versionBranch version)) m cs) constraints (map (Text.pack . display) (exposedModules' meta))
         alterConstraint packageName' version constraint =
           case constraint of
             Nothing -> Just [PackageConstraint packageName' version version]
             Just constraints ->
               -- TODO: This could probably be more efficient.
               case find (\c -> packageName c == packageName') constraints of
                 -- If the package is already listed, update the constraint.
                 Just _ -> Just (map (updateConstraint packageName' version) constraints)
                 -- If not, add a new constraint.
                 Nothing -> Just $ constraints ++ [PackageConstraint packageName' version version]
         updateConstraint name version constraint = if packageName constraint == name
                                                      then if version < packageMinVersion constraint
                                                             then constraint { packageMinVersion = version }
                                                             else if version < packageMaxVersion constraint
                                                                    then constraint { packageMaxVersion = version }
                                                                    else constraint
                                                      else constraint
         exposedModules' = fromMaybe [] . fmap (exposedModules . condTreeData) . condLibrary
  
  parseScript :: String -> Bool -> BS.ByteString -> Module
  parseScript filename noSugar source =
    case parseFileContentsWithMode
         (defaultParseMode { parseFilename = filename
                           , extensions = exts
                           , fixities = Just (infixl_ 8 ["^..", "^?", "^?!", "^@..", "^@?", "^@?!", "^.", "^@."])})
         source' of
      ParseOk m -> m
      r@(ParseFailed _ _) -> fromParseResult r
   where exts = if noSugar then [] else [EnableExtension TemplateHaskell, EnableExtension QuasiQuotes, EnableExtension PackageImports]
         source' = if noSugar then BU.toString source'' else desugar (BU.toString source'')
         -- Remove any leading shebang line.
         source'' = if B8.isPrefixOf "#!" source
                      then B8.dropWhile (/= '\n') source
                      else source
  