AgentRun.hs

{-# LANGUAGE TemplateHaskell #-}

-- | Agent runner for blueprints. Loads a blueprint, resolves its
-- variables (with the same precedence chain as @seihou run@), optionally
-- applies its declared @baseModules@, renders the prompt template, and
-- sends the rendered prompt through the configured Baikai provider. See EP-31
-- (docs/plans/31-blueprint-agent-runner.md) for the full design.
module Seihou.CLI.AgentRun
  ( handleAgentRun,
    appliedBlueprintFromOutcome,
    runRenderedAgentPrompt,
  )
where

import Control.Exception (IOException, displayException, try)
import Control.Monad (when)
import Data.FileEmbed (embedFile)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe, maybeToList)
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Text.IO qualified as TIO
import Data.Time (UTCTime)
import Data.Time.Clock (getCurrentTime)
import Seihou.CLI.AgentCompletion
  ( AgentModelConfig (..),
    AgentProvider (..),
    buildAgentCompletionRequest,
    runAgentCompletion,
  )
import Seihou.CLI.AgentLaunch
  ( AgentContext (..),
    BaselineStatus (..),
    formatAvailableModules,
    formatBaselineStatus,
    formatLocalModules,
    formatManifestState,
    formatModuleDhallState,
    formatSeihouProjectState,
    gatherAgentContext,
    substitute,
  )
import Seihou.CLI.AgentLaunchExec (launchConfiguredAgentAddingDirs)
import Seihou.CLI.AppliedBlueprint (recordAppliedBlueprint)
import Seihou.CLI.BlueprintExecution
  ( BlueprintExecutionRequest (..),
    PreparedBlueprintExecution (..),
    prepareBlueprintExecution,
    varValueToText,
  )
import Seihou.CLI.Commands (BlueprintRunOpts (..))
import Seihou.CLI.Shared
  ( deriveNamespace,
    formatVarError,
    logIO,
    toVarNameMap,
    unwrapConfig,
  )
import Seihou.Composition.Instance (ModuleInstance (..), qualifiedName)
import Seihou.Composition.Plan (compileComposedPlan)
import Seihou.Composition.Resolve (loadComposition, resolveWithPrompts)
import Seihou.Core.Context (resolveContext)
import Seihou.Core.Module (defaultSearchPaths, discoverRunnable)
import Seihou.Core.Types
import Seihou.Effect.BaselineStore (pruneBaselines)
import Seihou.Effect.BaselineStoreInterp (runBaselineStore)
import Seihou.Effect.ConfigReader
  ( readContextConfig,
    readGlobalConfig,
    readLocalConfig,
    readNamespaceConfig,
  )
import Seihou.Effect.ConfigReaderInterp (runConfigReader)
import Seihou.Effect.ConsoleInterp (runConsole)
import Seihou.Effect.Filesystem (createDirectoryIfMissing)
import Seihou.Effect.FilesystemInterp (runFilesystem)
import Seihou.Effect.Logger (logError, logInfo, logWarn)
import Seihou.Effect.ManifestStore (readManifest, writeManifest)
import Seihou.Effect.ManifestStoreInterp (runManifestStore)
import Seihou.Engine.Baseline (manifestBaselineRefs, recordGeneratedBaselines)
import Seihou.Engine.Conflict (resolveConflicts)
import Seihou.Engine.Diff (computeDiff)
import Seihou.Engine.Execute (executePlan)
import Seihou.Manifest.Types (currentManifestVersion, emptyManifest)
import Seihou.Prelude
import System.Environment (getEnvironment)
import System.Exit (ExitCode (..), exitFailure, exitWith)
import System.FilePath (takeDirectory, (</>))

-- | The prompt template, embedded at compile time from data/blueprint-prompt.md.
promptTemplate :: Text
promptTemplate = TE.decodeUtf8 $(embedFile "data/blueprint-prompt.md")

handleAgentRun :: Bool -> AgentModelConfig -> BlueprintRunOpts -> IO ()
handleAgentRun debug modelConfig opts = do
  let level = if opts.runBlueprintVerbose then LogVerbose else LogNormal

  -- (a) Discover and validate. discoverRunnable resolves by directory
  -- name (priority: module > recipe > blueprint).
  searchPaths <- defaultSearchPaths
  runnableResult <- discoverRunnable searchPaths opts.runBlueprintName
  (bp, blueprintDir) <- case runnableResult of
    Right (RunnableBlueprint b dir) -> pure (b, dir)
    Right (RunnableModule _ _) ->
      exitErr level $
        "'"
          <> opts.runBlueprintName.unModuleName
          <> "' is a module, not a blueprint. Did you mean 'seihou run "
          <> opts.runBlueprintName.unModuleName
          <> "'?"
    Right (RunnableRecipe _ _) ->
      exitErr level $
        "'"
          <> opts.runBlueprintName.unModuleName
          <> "' is a recipe, not a blueprint. Did you mean 'seihou run "
          <> opts.runBlueprintName.unModuleName
          <> "'?"
    Left err -> exitErr level (renderModuleLoadError err)

  let providerCanMountFiles =
        modelConfig.agentProvider == AgentProviderClaudeCli
          || modelConfig.agentProvider == AgentProviderCodexCli
  -- (b) Resolve variables and prepare the shared prompt/reference/tool state.
  preparedResult <-
    prepareBlueprintExecution
      BlueprintExecutionRequest
        { executionBlueprint = bp,
          executionBlueprintDir = blueprintDir,
          executionVariableOverrides = opts.runBlueprintVars,
          executionNamespaceOverride = opts.runBlueprintNamespace,
          executionContextOverride = opts.runBlueprintContext,
          executionCanMountFiles = providerCanMountFiles,
          executionLogLevel = level
        }
  prepared <- case preparedResult of
    Left errs -> do
      logIO level $ logError "Error resolving blueprint variables:"
      mapM_ (logIO level . logError . ("  " <>) . formatVarError) errs
      exitFailure
    Right result -> pure result
  let resolved = prepared.preparedResolvedVariables
      cliOverrides = Map.fromList [(VarName k, v) | (k, v) <- opts.runBlueprintVars]

  -- (c) Baseline.
  baseline <-
    if opts.runBlueprintNoBaseline
      then pure BaselineSkipped
      else
        if null bp.baseModules
          then pure BaselineEmpty
          else applyBaseline level opts bp.baseModules cliOverrides resolved

  -- (d) Render the system prompt around the prepared shared body.
  ctx <- gatherAgentContext
  let systemPrompt = renderSystemPrompt ctx prepared baseline

  -- (f) Launch.
  launchSucceeded <-
    runRenderedAgentPrompt
      debug
      modelConfig
      prepared.preparedAllowedTools
      prepared.preparedMountedFilesDir
      systemPrompt
      opts.runBlueprintPrompt

  -- (g) Record the applied-blueprint provenance into
  -- .seihou/manifest.json only after a successful provider response. In
  -- debug mode, keep the previous successful dry-launch behavior by recording
  -- after the rendered prompt is printed successfully.
  when launchSucceeded $ do
    now <- getCurrentTime
    let entry = appliedBlueprintFromOutcome bp baseline opts now
        manifestPath = ".seihou" </> "manifest.json"
    writeRes <- recordAppliedBlueprint manifestPath entry
    case writeRes of
      Right () -> pure ()
      Left err ->
        logIO level $
          logError $
            "Warning: agent succeeded but recording the applied-blueprint entry failed: "
              <> err

runRenderedAgentPrompt :: Bool -> AgentModelConfig -> [String] -> Maybe FilePath -> Text -> Maybe Text -> IO Bool
runRenderedAgentPrompt debug modelConfig tools mFilesDir systemPrompt initialPrompt
  | debug = do
      TIO.putStr systemPrompt
      pure True
  | modelConfig.agentProvider == AgentProviderClaudeCli || modelConfig.agentProvider == AgentProviderCodexCli = do
      exitCode <-
        launchConfiguredAgentAddingDirs
          (maybeToList mFilesDir)
          modelConfig
          tools
          debug
          systemPrompt
          initialPrompt
      case exitCode of
        ExitSuccess -> pure True
        ExitFailure _ -> exitWith exitCode
  | otherwise = do
      result <- runAgentCompletion (buildAgentCompletionRequest modelConfig systemPrompt initialPrompt)
      case result of
        Right assistantText -> do
          TIO.putStrLn assistantText
          pure True
        Left err -> do
          TIO.putStrLn $ "Error: " <> err
          exitFailure

-- | Project the runner's local state into the persistent
-- 'AppliedBlueprint' shape. Pure so the manifest writer remains a
-- one-liner at the call site and so cross-plan tests can drive it
-- with synthetic inputs.
appliedBlueprintFromOutcome ::
  Blueprint -> BaselineStatus -> BlueprintRunOpts -> UTCTime -> AppliedBlueprint
appliedBlueprintFromOutcome bp baseline opts now =
  AppliedBlueprint
    { name = bp.name,
      blueprintVersion = bp.version,
      appliedAt = now,
      baselineModules = case baseline of
        BaselineApplied entries -> map fst entries
        BaselineEmpty -> []
        BaselineSkipped -> [],
      noBaseline = case baseline of
        BaselineSkipped -> True
        _ -> False,
      userPrompt = opts.runBlueprintPrompt,
      agentSessionId = Nothing
    }

-- | Apply the blueprint's @baseModules@ to the cwd. Mirrors the
-- composition pipeline in @Seihou.CLI.Run.handleRun@: load every
-- declared base module (plus transitive deps), resolve their variables
-- through the same precedence chain (with the blueprint's own resolved
-- vars folded into the CLI override map so the agent's prompt and the
-- base modules see the same values), compile the composed plan,
-- compute the diff, resolve conflicts, execute the plan, and write the
-- resulting manifest. Returns 'BaselineApplied' listing each module's
-- (name, version) for the prompt's "Baseline" section.
applyBaseline ::
  LogLevel ->
  BlueprintRunOpts ->
  [Dependency] ->
  Map VarName Text ->
  Map VarName ResolvedVar ->
  IO BaselineStatus
applyBaseline level opts baseModules cliOverridesIn resolvedBlueprintVars = do
  searchPaths <- defaultSearchPaths
  (primary, additionals) <- case baseModules of
    d : rs -> pure (d.depModule, map (.depModule) rs)
    [] -> exitErr level "internal error: applyBaseline called with empty baseModules"
  compositionResult <- loadComposition searchPaths primary additionals
  modulesInOrder <- case compositionResult of
    Left err -> do
      logIO level $ logError $ "Baseline error: " <> renderModuleLoadError err
      exitFailure
    Right ms -> pure ms

  -- Fold the blueprint's resolved vars into the CLI override map for
  -- the base modules. CLI overrides (already present in cliOverridesIn)
  -- win over blueprint values, mirroring 'seihou run' semantics.
  let blueprintAsOverrides =
        Map.fromList
          [(vn, varValueToText rv.value) | (vn, rv) <- Map.toList resolvedBlueprintVars]
      cliOverrides = Map.union cliOverridesIn blueprintAsOverrides

  envPairs <- getEnvironment
  let envVars = Map.fromList [(T.pack k, T.pack v) | (k, v) <- envPairs]
      namespace = fromMaybe (deriveNamespace primary) opts.runBlueprintNamespace
  context <- resolveContext opts.runBlueprintContext envVars
  let contextName = fromMaybe "" context

  baseResolveResult <- runEff $ runConfigReader $ runConsole $ do
    localCfg <- readLocalConfig >>= unwrapConfig level
    nsCfg <- readNamespaceConfig namespace >>= unwrapConfig level
    ctxCfg <- readContextConfig contextName >>= unwrapConfig level
    gCfg <- readGlobalConfig >>= unwrapConfig level
    resolveWithPrompts
      modulesInOrder
      cliOverrides
      envVars
      namespace
      contextName
      (toVarNameMap localCfg)
      (toVarNameMap nsCfg)
      (toVarNameMap ctxCfg)
      (toVarNameMap gCfg)

  baseResolved <- case baseResolveResult of
    Left errs -> do
      logIO level $ logError "Error resolving baseline variables:"
      mapM_ (logIO level . logError . ("  " <>) . formatVarError) errs
      exitFailure
    Right r -> pure r

  -- Compile the plan.
  let quads =
        [ (inst, m, dir, Map.map (.value) (baseResolved Map.! inst))
        | (inst, m, dir) <- modulesInOrder
        ]
  planResult <- compileComposedPlan quads
  (ops, _warnings, ownerMap) <- case planResult of
    Left errs -> do
      logIO level $ logError "Errors compiling baseline plan:"
      mapM_ (logIO level . logError . ("  " <>)) errs
      exitFailure
    Right r -> pure r

  -- Read the manifest, compute the diff, resolve conflicts, execute,
  -- write the manifest. Mirrors Seihou.CLI.Run.handleRun.
  now <- getCurrentTime
  let manifestPath = ".seihou" </> "manifest.json"
      baselineDir = ".seihou" </> "baselines"
      planned =
        [(dest, content, primary, Nothing) | WriteFileOp dest content _ <- ops]
          ++ [(dest, content, mName, Just pOp) | PatchFileOp dest content pOp _ mName <- ops]

  existingRes <- runEff $ runFilesystem $ runManifestStore manifestPath $ do
    createDirectoryIfMissing True (takeDirectory manifestPath)
    readManifest
  manifest <- case existingRes of
    Left err -> do
      logIO level $ logError $ "Error reading manifest: " <> err
      exitFailure
    Right m -> pure (fromMaybe (emptyManifest now) m)

  diff <- runEff $ runFilesystem $ runManifestStore manifestPath $ do
    let composedNames =
          Set.fromList $
            concatMap (\(inst, _, _) -> [inst.instanceModule, qualifiedName inst]) modulesInOrder
    computeDiff manifest composedNames planned

  resolutions <-
    runEff $ runConsole $ resolveConflicts opts.runBlueprintForce diff.conflicts
  case resolutions of
    Nothing -> do
      logIO level $ logError "Baseline conflicts detected (use --force to overwrite):"
      mapM_ (\c -> logIO level (logError ("  ! " <> T.pack c.path))) diff.conflicts
      exitFailure
    Just conflictResolved -> do
      let keepRecords =
            Map.fromList
              [ ( c.path,
                  case Map.lookup c.path manifest.files of
                    Just existing -> existing {hash = c.diskHash, generatedAt = now}
                    Nothing ->
                      FileRecord
                        { hash = c.diskHash,
                          moduleName = c.moduleName,
                          strategy = Template,
                          generatedAt = now,
                          baseline = Nothing,
                          applicationIds = mempty
                        }
                )
              | (c, KeepCurrent) <- conflictResolved
              ]
          skipPaths = [c.path | (c, Skip) <- conflictResolved]
          excludePaths = Set.fromList (Map.keys keepRecords ++ skipPaths)
          opsForExec = filter (not . opTargetsPath excludePaths) ops

      generationAttempt <-
        try @IOException $
          runEff $
            runFilesystem $
              runBaselineStore baselineDir $
                runManifestStore manifestPath $ do
                  recs <- executePlan "" opsForExec ownerMap primary now
                  baselineResult <- recordGeneratedBaselines "" recs
                  case baselineResult of
                    Left err -> pure (Left err)
                    Right baselineRecords -> do
                      let orphanedPaths = map (.path) diff.orphaned
                          cleanedFiles = foldr Map.delete manifest.files orphanedPaths
                          allModuleEntries = updateAllModules manifest.modules modulesInOrder now
                          allResolvedVals =
                            Map.unions [Map.map (.value) vs | vs <- Map.elems baseResolved]
                          newManifest =
                            Manifest
                              { version = currentManifestVersion,
                                genAt = now,
                                modules = allModuleEntries,
                                vars = Map.union (Map.map varValueToText allResolvedVals) manifest.vars,
                                files = Map.unions [baselineRecords, keepRecords, cleanedFiles],
                                applications = manifest.applications,
                                recipe = manifest.recipe,
                                blueprint = manifest.blueprint,
                                blueprintMigrations = manifest.blueprintMigrations
                              }
                      writeManifest newManifest
                      pure (Right newManifest)

      newManifest <- case generationAttempt of
        Left err -> do
          logIO level $ logError $ "Error applying baseline files or storing generated baselines: " <> T.pack (displayException err)
          exitFailure
        Right (Left err) -> do
          logIO level $ logError $ "Error storing generated baselines: " <> T.pack (show err)
          exitFailure
        Right (Right saved) -> pure saved

      pruneAttempt <-
        try @IOException $
          runEff $
            runFilesystem $
              runBaselineStore baselineDir $
                pruneBaselines (manifestBaselineRefs newManifest)
      case pruneAttempt of
        Left err ->
          logIO level $ logWarn $ "Warning: could not prune generated baselines: " <> T.pack (displayException err)
        Right _ -> pure ()

      let nNew = length diff.new
          nMod = length diff.modified
          nUnch = length diff.unchanged
      logIO level $
        logInfo $
          "Baseline applied: "
            <> T.pack (show nNew)
            <> " new, "
            <> T.pack (show nMod)
            <> " modified, "
            <> T.pack (show nUnch)
            <> " unchanged."
      pure $
        BaselineApplied
          [(m.name, m.version) | (_, m, _) <- modulesInOrder]

-- | Stitch the system-prompt template together. Each block in
-- @blueprint-prompt.md@ has a @{{key}}@ placeholder filled here.
renderSystemPrompt :: AgentContext -> PreparedBlueprintExecution -> BaselineStatus -> Text
renderSystemPrompt ctx prepared baseline =
  let bp = prepared.preparedBlueprint
   in substitute
        [ ("cwd", ctx.cwd),
          ("seihou_project_state", formatSeihouProjectState ctx),
          ("manifest_state", formatManifestState ctx),
          ("module_dhall_state", formatModuleDhallState ctx),
          ("local_modules", formatLocalModules ctx),
          ("available_modules", formatAvailableModules ctx),
          ("blueprint_name", bp.name.unModuleName),
          ("blueprint_version", fromMaybe "(unspecified)" bp.version),
          ("blueprint_description", fromMaybe "(no description)" bp.description),
          ("baseline_status", formatBaselineStatus baseline),
          ("reference_files", prepared.preparedReferenceFiles),
          ("reference_files_dir", prepared.preparedReferenceFilesAccess),
          ("user_prompt", prepared.preparedSharedPrompt)
        ]
        promptTemplate

-- | Whether an operation targets a file in the given path set. Local
-- copy of @Seihou.CLI.Run.opTargetsPath@.
opTargetsPath :: Set FilePath -> Operation -> Bool
opTargetsPath paths (WriteFileOp dest _ _) = Set.member dest paths
opTargetsPath paths (PatchFileOp dest _ _ _ _) = Set.member dest paths
opTargetsPath _ _ = False

-- | Merge freshly-composed module instances into the manifest's
-- applied-modules list. Local copy of @Seihou.CLI.Run.updateAllModules@.
updateAllModules ::
  [AppliedModule] ->
  [(ModuleInstance, Module, FilePath)] ->
  UTCTime ->
  [AppliedModule]
updateAllModules existing modulesInOrder now =
  let composedKeys =
        Set.fromList
          [ (inst.instanceModule, inst.instanceParentVars)
          | (inst, _, _) <- modulesInOrder
          ]
      filtered =
        filter (\am -> not (Set.member (am.name, am.parentVars) composedKeys)) existing
      new =
        [ AppliedModule
            { name = inst.instanceModule,
              parentVars = inst.instanceParentVars,
              source = dir,
              moduleVersion = m.version,
              appliedAt = now,
              removal = m.removal
            }
        | (inst, m, dir) <- modulesInOrder
        ]
   in filtered ++ new

-- | Print an error message and exit with code 1.
exitErr :: LogLevel -> Text -> IO a
exitErr level msg = do
  logIO level (logError msg)
  exitFailure

-- | Render a 'ModuleLoadError' for display.
renderModuleLoadError :: ModuleLoadError -> Text
renderModuleLoadError = \case
  ModuleNotFound name searched ->
    "Module '"
      <> name.unModuleName
      <> "' not found. Searched in:\n"
      <> T.intercalate "\n" (map (("  " <>) . T.pack) searched)
  DhallEvalError name msg ->
    "Failed to evaluate '" <> name.unModuleName <> "': " <> msg
  DhallDecodeError name msg ->
    "Failed to decode '" <> name.unModuleName <> "': " <> msg
  ValidationError name msgs ->
    "Validation failed for '"
      <> name.unModuleName
      <> "':\n"
      <> T.intercalate "\n" (map ("  " <>) msgs)
  CircularDependency names ->
    "Circular dependency detected: "
      <> T.intercalate " -> " (map (.unModuleName) names)
  MissingSourceFile name path ->
    "Missing source file in '"
      <> name.unModuleName
      <> "': "
      <> T.pack path
  RegistryEvalError path msg ->
    "Failed to evaluate registry at '" <> path <> "': " <> msg