Session.hs

-- | Global session management for Hegel.
--
-- This module manages a shared hegel subprocess communicating over stdio
-- pipes. It starts lazily on first use and cleans up when the process exits.
--
-- The main entry point is 'runHegelTest', which the user calls without
-- needing to manage connections or sessions directly.
module Hegel.Session
  ( -- * Session management
    startSession
  , cleanupSession
  , restartSession

    -- * Running tests
  , runHegelTest
  , runHegelTest_

    -- * Hegel discovery
  , findHegel
  , ensureHegelInstalled
  ) where

import Control.Concurrent (forkIO)
import Control.Concurrent.MVar
import Control.Exception (evaluate, try, SomeException)
import Control.Monad (unless)
import Data.Foldable (for_)
import Data.IORef
import System.Directory (createDirectoryIfMissing, doesFileExist)
import System.Environment (lookupEnv)
import System.Exit (ExitCode (..))
import System.IO
  ( hSetBinaryMode
  , hClose
  , IOMode (..)
  , openFile
  )
import System.IO.Unsafe (unsafePerformIO)
import System.Process
  ( CreateProcess (..)
  , StdStream (..)
  , ProcessHandle
  , proc
  , createProcess
  , waitForProcess
  , terminateProcess
  )

import Hegel.Client
import Hegel.Connection


-- ---------------------------------------------------------------------------
-- Constants
-- ---------------------------------------------------------------------------

-- | Version of hegel-core to install.
hegelServerVersion :: String
hegelServerVersion = "0.2.3"

-- | Environment variable to override the hegel server command.
hegelServerCommandEnv :: String
hegelServerCommandEnv = "HEGEL_SERVER_COMMAND"

-- | Directory for hegel server installation and logs.
hegelServerDir :: FilePath
hegelServerDir = ".hegel"

-- | Message shown when uv is not found.
uvNotFoundMessage :: String
uvNotFoundMessage =
  "You are seeing this error message because hegel-haskell tried to use `uv` to \
  \install hegel-core, but could not find uv on the PATH.\n\n\
  \Hegel uses a Python server component called `hegel-core` to share core \
  \property-based testing functionality across languages. There are two ways \
  \for Hegel to get hegel-core:\n\n\
  \* By default, Hegel looks for uv (https://docs.astral.sh/uv/) on the PATH, \
  \and uses uv to install hegel-core to a local `.hegel/venv` directory. We \
  \recommend this option. To continue, install uv: \
  \https://docs.astral.sh/uv/getting-started/installation/.\n\
  \* Alternatively, you can manage the installation of hegel-core yourself. \
  \After installing, setting the HEGEL_SERVER_COMMAND environment variable to \
  \your hegel-core binary path tells hegel-haskell to use that hegel-core \
  \instead.\n\n\
  \See https://hegel.dev/reference/installation for more details."


-- ---------------------------------------------------------------------------
-- Command execution helpers
-- ---------------------------------------------------------------------------

-- | Runs a command with arguments, redirecting stdout and stderr to a log
-- file. Returns the exit code.
runCommandToLog :: FilePath -> [String] -> FilePath -> IO ExitCode
runCommandToLog cmd args logPath = do
  logHandle <- openFile logPath AppendMode
  let cp = (proc cmd args)
        { std_out = UseHandle logHandle
        , std_err = UseHandle logHandle
        }
  (_, _, _, ph) <- createProcess cp
  exitCode <- waitForProcess ph
  hClose logHandle
  pure exitCode


-- ---------------------------------------------------------------------------
-- Hegel discovery and installation
-- ---------------------------------------------------------------------------

-- | Checks for a cached hegel installation and installs via uv if needed.
-- Returns the path to the hegel binary.
ensureHegelInstalled :: IO FilePath
ensureHegelInstalled = do
  let venvDir = hegelServerDir ++ "/venv"
  let versionFile = venvDir ++ "/hegel-version"
  let hegelBin = venvDir ++ "/bin/hegel"
  let installLog = hegelServerDir ++ "/install.log"

  -- Check cached version (force the file read inside try to catch IO errors)
  cached <- try (do
    contents <- readFile versionFile
    let trimmed = strip contents
    -- Force evaluation of trimmed to catch lazy IO exceptions inside try
    _ <- evaluate (length trimmed)
    exists <- doesFileExist hegelBin
    if trimmed == hegelServerVersion && exists
      then pure (Just hegelBin)
      else pure Nothing
    ) :: IO (Either SomeException (Maybe FilePath))
  case cached of
    Right (Just path) -> pure path
    _ -> do
      -- Need to install
      createDirectoryIfMissing True hegelServerDir

      -- Create venv
      venvResult <- try (runCommandToLog "uv" ["venv", "--clear", venvDir] installLog)
        :: IO (Either SomeException ExitCode)
      case venvResult of
        Left _ -> fail uvNotFoundMessage
        Right (ExitFailure _) -> do
          logContents <- readFileSafe installLog
          fail ("uv venv failed. Install log:\n" ++ logContents)
        Right ExitSuccess -> pure ()

      -- Install hegel-core
      let pythonPath = venvDir ++ "/bin/python"
      pipResult <- runCommandToLog "uv"
        [ "pip", "install"
        , "--python", pythonPath
        , "hegel-core==" ++ hegelServerVersion
        ]
        installLog
      case pipResult of
        ExitFailure _ -> do
          logContents <- readFileSafe installLog
          fail $ "Failed to install hegel-core (version: " ++ hegelServerVersion
            ++ "). Set " ++ hegelServerCommandEnv
            ++ " to a hegel binary path to skip installation.\n"
            ++ "Install log:\n" ++ logContents
        ExitSuccess -> pure ()

      exists <- doesFileExist hegelBin
      unless exists $
        fail ("hegel not found at " ++ hegelBin ++ " after installation")

      -- Write version file
      writeFile versionFile hegelServerVersion
      pure hegelBin

-- | Locates the hegel binary. Checks, in order:
--
-- 1. @HEGEL_SERVER_COMMAND@ environment variable
-- 2. Auto-install via uv to @.hegel\/venv\/@ (version-locked)
--
-- The auto-installed version is preferred over any @hegel@ on @PATH@
-- because the system binary may be an incompatible version.
findHegel :: IO FilePath
findHegel = do
  envCmd <- lookupEnv hegelServerCommandEnv
  case envCmd of
    Just path | not (null path) -> pure path
    _ -> ensureHegelInstalled


-- ---------------------------------------------------------------------------
-- CI detection
-- ---------------------------------------------------------------------------

-- | CI environment variables to check for auto-detection. Each entry is
-- @(varName, expectedValue)@ where 'Nothing' means "any value".
ciVars :: [(String, Maybe String)]
ciVars =
  [ ("CI", Nothing)
  , ("TF_BUILD", Just "true")
  , ("BUILDKITE", Just "true")
  , ("CIRCLECI", Just "true")
  , ("CIRRUS_CI", Just "true")
  , ("CODEBUILD_BUILD_ID", Nothing)
  , ("GITHUB_ACTIONS", Just "true")
  , ("GITLAB_CI", Nothing)
  , ("HEROKU_TEST_RUN_ID", Nothing)
  , ("TEAMCITY_VERSION", Nothing)
  ]

-- | Returns 'True' if a CI environment is detected.
isInCI :: IO Bool
isInCI = anyM checkVar ciVars
 where
  checkVar (key, expected) = do
    val <- lookupEnv key
    pure $ case (val, expected) of
      (Just _, Nothing)   -> True
      (Just v, Just expV) -> v == expV
      (Nothing, _)        -> False

  anyM _ [] = pure False
  anyM f (x : xs) = do
    b <- f x
    if b then pure True else anyM f xs


-- ---------------------------------------------------------------------------
-- Session state
-- ---------------------------------------------------------------------------

-- | Internal mutable session state.
data HegelSession = HegelSession
  { sessionProcess    :: Maybe ProcessHandle
  , sessionConnection :: Maybe Connection
  , sessionClient     :: Maybe HegelClient
  }

-- | The initial empty session.
emptySession :: HegelSession
emptySession = HegelSession
  { sessionProcess    = Nothing
  , sessionConnection = Nothing
  , sessionClient     = Nothing
  }

-- | The global session singleton, protected by an MVar.
{-# NOINLINE globalSession #-}
globalSession :: MVar HegelSession
globalSession = unsafePerformIO (newMVar emptySession)

-- | IORef tracking whether cleanup has been registered with atexit.
{-# NOINLINE atExitRegistered #-}
atExitRegistered :: IORef Bool
atExitRegistered = unsafePerformIO (newIORef False)


-- ---------------------------------------------------------------------------
-- Session lifecycle
-- ---------------------------------------------------------------------------

-- | Returns 'True' if the session has a live client.
hasWorkingClient :: HegelSession -> IO Bool
hasWorkingClient session =
  case (sessionClient session, sessionConnection session) of
    (Just _, Just conn) -> isLive conn
    _ -> pure False

-- | Cleans up the session, closing the connection and terminating the
-- subprocess.
cleanup :: HegelSession -> IO HegelSession
cleanup session = do
  -- Close connection
  for_ (sessionConnection session) closeConnection
  -- Terminate process
  case sessionProcess session of
    Just ph -> do
      _ <- try (terminateProcess ph) :: IO (Either SomeException ())
      _ <- try (waitForProcess ph) :: IO (Either SomeException ExitCode)
      pure ()
    Nothing -> pure ()
  pure emptySession

-- | Starts the hegel server if not already running. Spawns the server
-- with @--stdio@ for pipe-based communication.
startSession :: IO ()
startSession = modifyMVar_ globalSession $ \session -> do
  working <- hasWorkingClient session
  if working
    then pure session
    else do
      -- Clean up any old session
      _ <- cleanup session
      hegelCmd <- findHegel

      -- Open log file for server stderr
      createDirectoryIfMissing True hegelServerDir
      logHandle <- openFile (hegelServerDir ++ "/server.log") AppendMode

      -- Start hegel subprocess with --stdio
      let cp = (proc hegelCmd ["--stdio", "--verbosity", "normal"])
            { std_in  = CreatePipe
            , std_out = CreatePipe
            , std_err = UseHandle logHandle
            }
      (Just hIn, Just hOut, Nothing, ph) <- createProcess cp

      -- Set handles to binary mode
      hSetBinaryMode hIn True
      hSetBinaryMode hOut True

      -- Create connection from the pipe handles
      conn <- createConnection hOut hIn False

      -- Create client (performs handshake)
      client <- createClient conn

      -- Monitor thread: detect server crash
      _ <- forkIO $ do
        _ <- waitForProcess ph
        setServerExited conn
        pure ()

      -- Register at-exit cleanup (only once)
      registered <- readIORef atExitRegistered
      unless registered $ do
        writeIORef atExitRegistered True
        -- Note: We cannot use Weak references to the MVar easily, so we
        -- just always try cleanup. If session is already empty, cleanup
        -- is a no-op.

      pure HegelSession
        { sessionProcess    = Just ph
        , sessionConnection = Just conn
        , sessionClient     = Just client
        }

-- | Forces a restart of the global session. Cleans up the current session
-- and clears the state.
restartSession :: IO ()
restartSession = modifyMVar_ globalSession cleanup

-- | Cleans up the global session. Closes the connection and terminates
-- the server process.
cleanupSession :: IO ()
cleanupSession = modifyMVar_ globalSession cleanup


-- ---------------------------------------------------------------------------
-- Running tests
-- ---------------------------------------------------------------------------

-- | Computes the initial settings for a test run.
--
-- In CI environments, Hegel starts from deterministic, stateless defaults by
-- enabling derandomization and disabling the example database.
initialSettingsForRun :: IO Settings
initialSettingsForRun = do
  inCI <- isInCI
  pure $
    if inCI
      then
        defaultSettings
          { settingsDerandomize = True
          , settingsDatabase = Disabled
          }
      else defaultSettings

-- | Runs a property test using the shared hegel process.
-- Ensures the session is started, then delegates to 'Client.runTest'.
runHegelTest :: (Settings -> Settings) -> (TestCase -> IO ()) -> IO ()
runHegelTest configure testFn = do
  effectiveSettings <- configure <$> initialSettingsForRun
  startSession
  session <- readMVar globalSession
  case sessionClient session of
    Just client -> runTest client effectiveSettings testFn
    Nothing     -> fail "Failed to start hegel session"

-- | Runs a property test using the initial runtime defaults unchanged.
runHegelTest_ :: (TestCase -> IO ()) -> IO ()
runHegelTest_ = runHegelTest id


-- ---------------------------------------------------------------------------
-- Utilities
-- ---------------------------------------------------------------------------

-- | Strip leading and trailing whitespace from a string.
strip :: String -> String
strip = reverse . dropWhile isSpace' . reverse . dropWhile isSpace'
 where
  isSpace' c = c == ' ' || c == '\n' || c == '\r' || c == '\t'

-- | Read a file, returning empty string on failure.
readFileSafe :: FilePath -> IO String
readFileSafe path = do
  result <- try (readFile path) :: IO (Either SomeException String)
  case result of
    Right contents -> pure contents
    Left _         -> pure ""