Nix.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoImplicitPrelude #-}

module HWM.Integrations.Toolchain.Nix (syncNixFile) where

import qualified Data.Set as Set
import qualified Data.Text as T
import HWM.Core.Common (Name)
import HWM.Core.Formatting (Status (..), format, toCamelCase)
import HWM.Core.Options (Options (..))
import HWM.Core.Pkg (Pkg (..))
import HWM.Core.Sync (SyncMode (..))
import HWM.Core.Version (Era (eraNixpkgs), formatNixGhc, selectEra)
import HWM.Domain.Config (Config (..))
import HWM.Domain.ConfigT (ConfigT, Env (..))
import HWM.Domain.Environments (BuildEnvironment (..), Environments (envsDefault), TargetModes (..), getBuildEnvironment, getBuildEnvironments)
import HWM.Domain.Release (getArtifactEnvironments)
import HWM.Runtime.Files (syncFile)
import Relude
import System.Directory (doesFileExist)

enabled :: SyncMode -> Bool
enabled mode = mode /= SyncModeIgnore

syncNixFile :: SyncMode -> ConfigT Status
syncNixFile SyncModeSync = do
  Config {..} <- asks config
  ops <- asks options
  benv <- getBuildEnvironment (Just $ envsDefault cfgEnvironments)
  benvs <- filter (enabled . targetNix . buildTargets) <$> getBuildEnvironments
  releasePkgs <- maybe (pure []) getArtifactEnvironments cfgRelease
  syncFile (optionsNix ops) (deriveFlakeNix (Context releasePkgs cfgName (map (cfgName,) benvs)) (cfgName, benv) (map (cfgName,) benvs))
syncNixFile SyncModeCheck = do
  nixPath <- optionsNix <$> asks options
  exists <- liftIO $ doesFileExist nixPath
  pure $ if exists then Checked else Warning
syncNixFile SyncModeIgnore = pure Ignored

genName :: (Text, Text) -> Text
genName (name, subName) = toCamelCase (name <> T.toTitle subName <> "WorkspacePackages")

genNixName :: BCOntext -> Text
genNixName (name, BuildEnvironment {..}) = genName (name, format buildName)

type BCOntext = (Name, BuildEnvironment)

systems :: [Text]
systems =
  [ "x86_64-linux",
    "aarch64-linux",
    "x86_64-darwin",
    "aarch64-darwin"
  ]

type ReleasePkg = (Pkg, [BuildEnvironment])

type Release = [ReleasePkg]

data Context = Context
  { release :: Release,
    projectName :: Text,
    environments :: [BCOntext]
  }

deriveFlakeNix :: Context -> BCOntext -> [BCOntext] -> Text
deriveFlakeNix context ctx@(projectName, benv) ctxs =
  T.unlines
    $ braces
      ( [ "description = \"A Haskell " <> projectName <> " workspace generated by HWM(Haskell Workspace Manager)\";",
          "inputs = {",
          "  nixpkgs.url = \"github:NixOS/nixpkgs/" <> eraNixpkgs (selectEra (buildGHC benv)) <> "\";",
          "};",
          "outputs = { self, nixpkgs }:"
        ]
          <> letBlock
            ( [ "supportedSystems = [ " <> T.intercalate " " (map show systems) <> " ];",
                "forAllSystems = nixpkgs.lib.genAttrs supportedSystems;"
              ]
                -- OVERLAY
                <> genOverlay context
            )
            -- PACKAGES
            ( forAllSystems "packages" (nixPkgs <> artifactEngine) (genPackages context ctx ctxs)
                -- DEVSHELLS
                <> forAllSystems "devShells" nixPkgs (genDevShell True ctx <> concatMap (genDevShell False) ctxs)
                -- CHECKS
                <> ["checks = forAllSystems (system: self.packages.${system});"]
            )
            False
      )

-- OVERLAY
genOverlay :: Context -> [Text]
genOverlay ctx@Context {..} =
  scoped
    "haskellOverlay = final: prev:"
    (concatMap rendergOverlayItem environments <> concatMap (rendergOverlayStatic ctx) staticEnvs)
  where
    staticEnvs = Set.toList $ Set.fromList (concatMap snd release)

rendergOverlayItem :: BCOntext -> [Text]
rendergOverlayItem (name, BuildEnvironment {..}) =
  overlayFun
    (genName (name, buildName))
    ["haskell", "packages", formatNixGhc buildGHC]
    (map (renderPackageDef renderPackageBody) buildPkgs)

genStaticEnvName :: Text -> BuildEnvironment -> Text
genStaticEnvName projectName buildEnv = genName (projectName, format (buildName buildEnv) <> "-static")

rendergOverlayStatic :: Context -> BuildEnvironment -> [Text]
rendergOverlayStatic Context {..} env@BuildEnvironment {..} =
  overlayFun
    (genStaticEnvName projectName env)
    ["pkgsStatic", "haskell", "packages", formatNixGhc buildGHC]
    (map (renderPackageDef (stripExecutables . renderPackageBody)) buildPkgs)

overlayFun :: Text -> [Text] -> [Text] -> [Text]
overlayFun name extend = fun name (concatName $ ["prev"] <> extend <> ["extend"]) "hfinal: hprev:"

stripExecutables :: (Semigroup a, IsString a) => a -> a
stripExecutables x = "prev.haskell.lib.justStaticExecutables ( " <> x <> ")"

renderPackageDef :: (Pkg -> Text) -> Pkg -> Text
renderPackageDef body pkg = format (pkgName pkg) <> " = " <> body pkg <> ";"

renderPackageBody :: Pkg -> Text
renderPackageBody pkg = "hfinal.callCabal2nix \"" <> format (pkgName pkg) <> "\" ./" <> format (pkgDirPath pkg) <> " {}"

-- PACKAGES
genPackages :: Context -> BCOntext -> [BCOntext] -> [Text]
genPackages context ctx@(projectName, defaultEnv) allEnvs =
  map (\pkg -> "default = pkgs." <> genNixName ctx <> "." <> format (pkgName pkg) <> ";") defaultPkg
    <> map (\pkg -> format (pkgName pkg) <> " = pkgs." <> genNixName ctx <> "." <> format (pkgName pkg) <> ";") (buildPkgs defaultEnv)
    <> concatMap genEnvPackages allEnvs
    <> concatMap allEnvPackages allEnvs
    <> concatMap (genReleasePkgs context) (release context)
  where
    defaultPkg = filter ((projectName ==) . format . pkgName) (buildPkgs defaultEnv)

genPackage :: Text -> Text -> Pkg -> Text
genPackage overlay envName pkg = format (pkgName pkg) <> "-" <> envName <> " = pkgs." <> overlay <> "." <> format (pkgName pkg) <> ";"

genEnvPackages :: BCOntext -> [Text]
genEnvPackages ctx@(_, env) = map (genPackage (genNixName ctx) (toCamelCase (format (buildName env)))) (buildPkgs env)

allEnvPackages :: BCOntext -> [Text]
allEnvPackages ctx@(_, env) =
  genSymlinkJoin
    ("env-" <> toCamelCase (format (buildName env)) <> "-all")
    (format (buildName env) <> "-workspace")
    (map (\pkg -> "pkgs." <> genNixName ctx <> "." <> format (pkgName pkg)) (buildPkgs env))

genReleasePkgs :: Context -> ReleasePkg -> [Text]
genReleasePkgs ctx (pkg, envs) = map (genReleasePkg ctx pkg) envs

genReleasePkg :: Context -> Pkg -> BuildEnvironment -> Text
genReleasePkg ctx pkg env = name <> "-release = " <> genReleasePkgBody ctx pkg env <> ";"
  where
    name = format (pkgName pkg) <> "-" <> buildName env

genReleasePkgBody :: Context -> Pkg -> BuildEnvironment -> Text
genReleasePkgBody Context {..} pkg env = "mkReleaseArtifact \"" <> format (pkgName pkg) <> "\" pkgs." <> genNixName (projectName, env) <> "." <> format (pkgName pkg) <> " pkgs." <> genStaticEnvName projectName env <> "." <> format (pkgName pkg)

-- genReleaseGroup :: BCOntext -> [Pkg] -> [Text]
-- genReleaseGroup ctx pkgs =
--   genSymlinkJoin
--     "release"
--     "release-artifacts"
--     (map (\pkg -> "(" <> genReleasePkgBody ctx pkg <> ")") pkgs)

-- DEVSHELL
genDevShell :: Bool -> BCOntext -> [Text]
genDevShell _ (_, BuildEnvironment {buildPkgs = []}) = [] -- Handle empty workspace
genDevShell isDefault ctx@(_, BuildEnvironment {buildTargets = TargetModes {..}, ..}) =
  [ name <> " = " <> devShellPackageName ctx <> ".shellFor {",
    "  packages = p: [ " <> renderPackageList buildPkgs <> " ];",
    "  buildInputs = with " <> devShellPackageName ctx <> "; ["
  ]
    <> map ("    " <>) libs
    <> [ "  ];",
         "};"
       ]
  where
    name = if isDefault then "default" else toCamelCase (format buildName)
    libs = ["cabal-install", "hlint"] <> ["stack" | enabled targetStack] <> ["haskell-language-server" | enabled targetHie]
    renderPackageList = T.intercalate " " . map (\pkg -> "p." <> format (pkgName pkg))

devShellPackageName :: BCOntext -> Text
devShellPackageName ctx = "pkgs." <> genNixName ctx

artifactEngine :: [Text]
artifactEngine =
  [ "mkReleaseArtifact = pkgName: basePkg: staticPkg:",
    "  if pkgs.stdenv.hostPlatform.isLinux then",
    "    # LINUX: Return the fully static (musl), zero-dependency executable",
    "    staticPkg",
    "  else if pkgs.stdenv.hostPlatform.isDarwin then",
    "    # MACOS: We \"Bundle\" the dynamic dependencies since strict static is hard on Mac.",
    "    pkgs.runCommand \"${pkgName}-macos-bundle\" {",
    "      nativeBuildInputs = [ pkgs.macdylibbundler pkgs.darwin.autoSignDarwinBinariesHook ];",
    "    } ''",
    "      mkdir -p $out/bin",
    "",
    "      # 1. Copy the fast, dynamically linked binary (stripping docs/libs first)",
    "      cp ${pkgs.haskell.lib.justStaticExecutables basePkg}/bin/${pkgName} $out/bin/${pkgName}",
    "      chmod 755 $out/bin/${pkgName}",
    "",
    "      # 2. Run dylibbundler to pull all Nix store .dylibs into the local folder",
    "      dylibbundler -b --no-codesign -x $out/bin/${pkgName} -d $out/bin -p '@executable_path'",
    "",
    "      # 3. Re-sign the binary so Apple Silicon (M1/M2/M3) allows it to run natively",
    "      signDarwinBinariesInAllOutputs",
    "    ''",
    "  else",
    "    # Fallback for unexpected systems",
    "    basePkg;"
  ]

-- SYNTAX

concatName :: [Text] -> Text
concatName = T.intercalate "."

nixPkgs :: [Text]
nixPkgs = ["pkgs = import nixpkgs { inherit system; overlays = [ haskellOverlay ]; };"]

forAllSystems :: Text -> [Text] -> [Text] -> [Text]
forAllSystems system h body = [system <> " = forAllSystems (system:"] <> letBlock h body True

letBlock :: [Text] -> [Text] -> Bool -> [Text]
letBlock h body end = ["let"] <> indent h <> ["in", "{"] <> indent body <> if end then ["});"] else ["};"]

braces :: [Text] -> [Text]
braces body = ["{"] <> indent body <> ["}"]

scoped :: Text -> [Text] -> [Text]
scoped name body = [name <> " {"] <> indent body <> ["};"]

fun :: Text -> Text -> Text -> [Text] -> [Text]
fun varName name args body = [varName <> " = " <> name <> " (" <> args <> " {"] <> indent body <> ["});"]

genSymlinkJoin :: Text -> Text -> [Text] -> [Text]
genSymlinkJoin varName name pkgs =
  scoped
    (varName <> " = pkgs.symlinkJoin")
    (["name = \"" <> name <> "\";", "paths = ["] <> indent pkgs <> ["];"])

indent :: [Text] -> [Text]
indent = map ("  " <>)