Stack.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NoImplicitPrelude #-}

module HWM.Integrations.Toolchain.Stack
  ( Stack (..),
    syncStackYaml,
    setupStackMatrixEnvironment,
    scanStackFiles,
    buildMatrix,
  )
where

import Control.Monad.Except
import Data.Aeson
  ( FromJSON (..),
    ToJSON (..),
    genericParseJSON,
    genericToJSON,
  )
import qualified Data.Map as Map
import qualified Data.Text as T
import HWM.Core.Common (Name)
import HWM.Core.Formatting (Format (..), Status (..), slugify)
import HWM.Core.Options (Options (..), askOptions)
import HWM.Core.Parsing (Parse (..))
import HWM.Core.Pkg (Pkg (..), PkgName)
import HWM.Core.Result (Issue (..), fromEither)
import HWM.Core.Version (Version, latestGHCVersion, parseGHCVersion)
import HWM.Domain.Build (Builder (StackBuilder))
import HWM.Domain.ConfigT (ConfigT)
import HWM.Domain.Environments (BuildEnvironment (..), EnviromentProfile (..), Environments (..), Feature (..), StackEnvironment (..), getBuildEnvironment, hkgRefs, mkEnvironment, mkEnvironments)
import HWM.Domain.Workspace (toWorkspaceRef)
import HWM.Runtime.Cache (getSnapshotGHC)
import HWM.Runtime.Files (aesonYAMLOptions, readYaml, rewrite_)
import HWM.Runtime.Process (EnvVars)
import Relude hiding (head, tail)
import System.Directory (createDirectoryIfMissing, doesFileExist, makeAbsolute)
import System.FilePath (dropExtension, (</>))
import System.FilePath.Glob (compile, globDir1)
import System.FilePath.Posix (takeFileName)

data Stack = Stack
  { packages :: [FilePath],
    resolver :: Name,
    allowNewer :: Maybe Bool,
    saveHackageCreds :: Maybe Bool,
    extraDeps :: Maybe [Name],
    compiler :: Maybe Text
  }
  deriving
    ( Show,
      Generic
    )

type VersionRegistry = Map PkgName Version

instance FromJSON Stack where
  parseJSON = genericParseJSON aesonYAMLOptions

instance ToJSON Stack where
  toJSON = genericToJSON aesonYAMLOptions

parseExtraDeps :: (MonadError Issue m) => [Text] -> m (Maybe VersionRegistry)
parseExtraDeps [] = pure Nothing
parseExtraDeps entries = do
  parsed <- traverse parseExtraDep entries
  if null parsed then throwError "No valid extra dependencies found" else pure $ Just $ Map.fromList parsed

parseExtraDep :: (MonadError Issue m) => Text -> m (PkgName, Version)
parseExtraDep entry = do
  let (namePart, versionPart) = T.breakOnEnd "-" entry
      segment = T.dropEnd 1 namePart
  when (T.null segment || T.null versionPart)
    $ throwError "Invalid extra-dep format: missing package segment or version part"
  pkgName <- fromEither ("Invalid package name: " <> segment) (parse segment)
  version <- fromEither ("Invalid version: " <> versionPart) (parse versionPart)
  pure (pkgName, version)

syncStackYaml :: ConfigT Status
syncStackYaml = do
  stackYamlPath <- optionsStack <$> askOptions
  rewrite_ stackYamlPath $ const $ do
    BuildEnvironment {..} <- getBuildEnvironment Nothing
    pure
      Stack
        { saveHackageCreds = Just False,
          extraDeps = map format . sort . hkgRefs <$> buildExtraDeps,
          packages = map pkgDirPath buildPkgs,
          compiler = Nothing,
          resolver = buildResolver,
          allowNewer = buildAllowNewer,
          ..
        }

setupStackMatrixEnvironment :: BuildEnvironment -> ConfigT EnvVars
setupStackMatrixEnvironment env = do
  genStackMatrixConfig env
  getStackMatrixEnvVars (buildName env)

getStackMatrixEnvVars :: Name -> ConfigT EnvVars
getStackMatrixEnvVars envName = do
  yamlPath <- stackMatrixPath envName
  pure [("STACK_YAML", yamlPath)]

stackMatrixPath :: Name -> ConfigT FilePath
stackMatrixPath name = liftIO $ makeAbsolute $ ".hwm/matrix/stack-" <> toString name <> ".yaml"

genStackMatrixConfig :: BuildEnvironment -> ConfigT ()
genStackMatrixConfig BuildEnvironment {..} = do
  path <- stackMatrixPath buildName
  liftIO $ createDirectoryIfMissing True ".hwm/matrix/"
  _ <- rewrite_ path $ const $ do
    pure
      Stack
        { saveHackageCreds = Just False,
          extraDeps = map format . sort . hkgRefs <$> buildExtraDeps,
          packages = map (("../../" <>) . pkgDirPath) buildPkgs,
          compiler = Nothing,
          resolver = buildResolver,
          allowNewer = buildAllowNewer,
          ..
        }
  pure ()

scanStackFiles :: (MonadIO m, MonadError Issue m) => Options -> FilePath -> m [(Name, Stack)]
scanStackFiles opts root = do
  let defaultPath = root </> optionsStack opts
  defaultExists <- liftIO $ doesFileExist defaultPath
  variantPaths <- liftIO $ globDir1 (compile "stack-*.yaml") root
  traverse loadEnv ([defaultPath | defaultExists] <> variantPaths)
  where
    loadEnv path = do
      seConfig <- readYaml path
      let stackName = fromMaybe "default" (deriveEnviromentName path)
      pure (stackName, seConfig)

deriveEnviromentName :: FilePath -> Maybe Text
deriveEnviromentName path = slugify <$> T.stripPrefix "stack-" (toText (dropExtension (takeFileName path)))

buildMatrix :: (MonadIO m, MonadError Issue m) => [Pkg] -> [(Name, Stack)] -> m Environments
buildMatrix pkgs (defaultEnv : envs) = do
  environments <- sortOn (profileGhc . snd) <$> traverse (inferBuildEnv pkgs) (defaultEnv : envs)
  pure Environments {envsDefault = fst defaultEnv, envsProfiles = Map.fromList environments, envsStack = Just True, envsNix = Nothing, envsBuilder = Just StackBuilder, envsHie = Nothing}
buildMatrix _ [] = pure $ mkEnvironments latestGHCVersion

inferBuildEnv :: (MonadIO m, MonadError Issue m) => [Pkg] -> (Name, Stack) -> m (Name, EnviromentProfile)
inferBuildEnv allPkgs (name, Stack {extraDeps = deps, ..}) = do
  ghc <- maybe (getSnapshotGHC resolver) (fromEither "GHC Parsing" . parseGHCVersion) compiler
  extraDeps <- parseExtraDeps (fromMaybe [] deps)
  let excludeList = filter ((`notElem` packages) . pkgDirPath) allPkgs
  pure
    ( name,
      (mkEnvironment ghc)
        { profileStack = Just (Enabled StackEnvironment {resolver = Just resolver, ..}),
          profileExclude = if null excludeList then Nothing else Just (map toWorkspaceRef excludeList)
        }
    )