path: root/Omni/Jr.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- : out jr
-- : dep sqlite-simple
-- : dep sqids
-- : dep warp
-- : dep servant-server
-- : dep lucid
-- : dep servant-lucid
module Omni.Jr where

import Alpha
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Omni.Agent.Core as AgentCore
import qualified Omni.Agent.Worker as AgentWorker
import qualified Omni.Cli as Cli
import qualified Omni.Jr.Web as Web
import qualified Omni.Task as Task
import qualified Omni.Task.Core as TaskCore
import qualified Omni.Test as Test
import qualified System.Console.Docopt as Docopt
import qualified System.Directory as Directory
import System.Environment (withArgs)
import qualified System.Environment as Env
import qualified System.Exit as Exit
import System.FilePath (takeFileName)
import qualified System.IO as IO
import qualified System.Process as Process

main :: IO ()
main = Cli.main plan

plan :: Cli.Plan ()
plan =
  Cli.Plan
    { Cli.help = help,
      Cli.move = move,
      Cli.test = test,
      Cli.tidy = \_ -> pure ()
    }

help :: Cli.Docopt
help =
  [Cli.docopt|
jr

Usage:
  jr task [<args>...]
  jr work [<task-id>]
  jr web [--port=PORT]
  jr review [<task-id>] [--auto]
  jr loop [--delay=SECONDS]
  jr merge-driver <ours> <theirs>
  jr test
  jr (-h | --help)

Commands:
  task          Manage tasks
  work          Start a worker agent on a task
  web           Start the web UI server
  review        Review a completed task (show diff, accept/reject)
  loop          Run autonomous work+review loop
  merge-driver  Internal git merge driver

Options:
  -h --help        Show this help
  --port=PORT      Port for web server [default: 8080]
  --auto           Auto-review: accept if tests pass, reject if they fail
  --delay=SECONDS  Delay between loop iterations [default: 5]
|]

move :: Cli.Arguments -> IO ()
move args
  | args `Cli.has` Cli.command "task" = do
      let extraArgs = Cli.getAllArgs args (Cli.argument "args")
      withArgs extraArgs Task.main
  | args `Cli.has` Cli.command "web" = do
      let port = case Cli.getArg args (Cli.longOption "port") of
            Just p -> fromMaybe Web.defaultPort (readMaybe p)
            Nothing -> Web.defaultPort
      Web.run port
  | args `Cli.has` Cli.command "work" = do
      -- Always run in current directory
      let path = "."

      -- Infer name from current directory
      absPath <- Directory.getCurrentDirectory
      let name = Text.pack (takeFileName absPath)

      let worker =
            AgentCore.Worker
              { AgentCore.workerName = name,
                AgentCore.workerPid = Nothing,
                AgentCore.workerStatus = AgentCore.Idle,
                AgentCore.workerPath = path
              }

      let taskId = fmap Text.pack (Cli.getArg args (Cli.argument "task-id"))

      AgentWorker.start worker taskId
  | args `Cli.has` Cli.command "review" = do
      let autoMode = args `Cli.has` Cli.longOption "auto"
      case Cli.getArg args (Cli.argument "task-id") of
        Just tidStr -> reviewTask (Text.pack tidStr) autoMode
        Nothing -> do
          -- Find tasks in Review status
          tasks <- TaskCore.loadTasks
          let reviewTasks = filter (\t -> TaskCore.taskStatus t == TaskCore.Review) tasks
          case reviewTasks of
            [] -> putText "No tasks in Review status."
            (t : _) -> reviewTask (TaskCore.taskId t) autoMode
  | args `Cli.has` Cli.command "loop" = do
      let delay = case Cli.getArg args (Cli.longOption "delay") of
            Just d -> fromMaybe 5 (readMaybe d)
            Nothing -> 5
      runLoop delay
  | args `Cli.has` Cli.command "merge-driver" = mergeDriver args
  | otherwise = putText (str <| Docopt.usage help)

mergeDriver :: Cli.Arguments -> IO ()
mergeDriver args = do
  ours <- getArgOrExit args (Cli.argument "ours")
  theirs <- getArgOrExit args (Cli.argument "theirs")

  -- Set TASK_DB_PATH to ours (the file git provided as the current version)
  -- Since we are no longer using git-tracked tasks.jsonl, this merge driver logic needs rethinking.
  -- If the merge driver is called, it means git found a conflict in .tasks/tasks.jsonl, but we deleted it.
  -- So this code might be dead, or we might be dealing with legacy files during a rebase.
  -- For now, we'll keep the logic but be aware it's likely unused.
  Env.setEnv "TASK_DB_PATH" ours
  -- TaskCore.importTasks theirs
  -- Task.importTasks theirs

  -- We need to call task import via CLI or library.
  -- Omni.Task.importTasks IS NOT EXPOSED.
  -- But we can call Task.main
  withArgs ["import", "-i", theirs] Task.main
  Exit.exitSuccess

getArgOrExit :: Cli.Arguments -> Docopt.Option -> IO String
getArgOrExit args opt =
  case Cli.getArg args opt of
    Just val -> pure val
    Nothing -> do
      putText <| "Error: Missing required argument " <> Text.pack (show opt)
      Exit.exitFailure

-- | Run the autonomous loop: work -> review -> repeat
runLoop :: Int -> IO ()
runLoop delaySec = do
  putText "Starting autonomous jr loop..."
  putText ("Delay between iterations: " <> tshow delaySec <> "s")
  go
  where
    go = do
      -- Check for ready work
      readyTasks <- TaskCore.getReadyTasks
      case readyTasks of
        [] -> do
          putText "\nNo ready tasks. Checking for tasks to review..."
          reviewPending
        (task : _) -> do
          putText ("\n=== Working on: " <> TaskCore.taskId task <> " ===")
          -- Run worker
          absPath <- Directory.getCurrentDirectory
          let name = Text.pack (takeFileName absPath)
          let worker =
                AgentCore.Worker
                  { AgentCore.workerName = name,
                    AgentCore.workerPid = Nothing,
                    AgentCore.workerStatus = AgentCore.Idle,
                    AgentCore.workerPath = "."
                  }
          AgentWorker.start worker (Just (TaskCore.taskId task))
          -- After work, check for review
          reviewPending

      -- Delay and loop
      putText ("\nSleeping " <> tshow delaySec <> "s...")
      threadDelay (delaySec * 1000000)
      go

    reviewPending = do
      tasks <- TaskCore.loadTasks
      let reviewTasks = filter (\t -> TaskCore.taskStatus t == TaskCore.Review) tasks
      case reviewTasks of
        [] -> putText "No tasks pending review."
        (t : _) -> do
          putText ("\n=== Auto-reviewing: " <> TaskCore.taskId t <> " ===")
          reviewTask (TaskCore.taskId t) True

reviewTask :: Text -> Bool -> IO ()
reviewTask tid autoMode = do
  tasks <- TaskCore.loadTasks
  case TaskCore.findTask tid tasks of
    Nothing -> do
      putText ("Task " <> tid <> " not found.")
      Exit.exitFailure
    Just task -> do
      unless autoMode <| TaskCore.showTaskDetailed task

      let grepArg = "--grep=" <> Text.unpack tid
      (code, shaOut, _) <-
        Process.readProcessWithExitCode
          "git"
          ["log", "--pretty=format:%H", "-n", "1", grepArg]
          ""

      when (code /= Exit.ExitSuccess || null shaOut) <| do
        putText "\nNo commit found for this task."
        putText "The worker may not have completed yet, or the commit message doesn't include the task ID."
        Exit.exitFailure

      let commitSha = case List.lines shaOut of
            (x : _) -> x
            [] -> ""

      -- Check for merge conflicts before showing diff
      conflictResult <- checkMergeConflict commitSha
      case conflictResult of
        Just conflictFiles -> do
          putText "\n=== MERGE CONFLICT DETECTED ==="
          putText "This commit cannot be cleanly applied to live."
          putText "Conflicting files:"
          traverse_ (\f -> putText ("  - " <> f)) conflictFiles
          putText ""
          putText "Kicking back to coder with context..."

          -- Get current retry count
          maybeCtx <- TaskCore.getRetryContext tid
          let attempt = maybe 1 (\c -> TaskCore.retryAttempt c + 1) maybeCtx

          if attempt > 3
            then do
              putText "\nTask has failed 3 times. Marking as NeedsHuman."
              -- For now, just mark as Open with a note (no NeedsHuman status yet)
              TaskCore.updateTaskStatus tid TaskCore.Open []
              putText ("Task " <> tid <> " needs human intervention (3 failed attempts).")
            else do
              -- Save retry context
              TaskCore.setRetryContext
                TaskCore.RetryContext
                  { TaskCore.retryTaskId = tid,
                    TaskCore.retryOriginalCommit = Text.pack commitSha,
                    TaskCore.retryConflictFiles = conflictFiles,
                    TaskCore.retryAttempt = attempt,
                    TaskCore.retryReason = "merge_conflict"
                  }
              TaskCore.updateTaskStatus tid TaskCore.Open []
              putText ("Task " <> tid <> " returned to queue (attempt " <> tshow attempt <> "/3).")
        Nothing -> do
          if autoMode
            then autoReview tid task commitSha
            else interactiveReview tid commitSha

-- | Auto-review: run tests on namespace, accept if pass, reject if fail
autoReview :: Text -> TaskCore.Task -> String -> IO ()
autoReview tid task commitSha = do
  putText "Running automated review..."

  -- Determine what to test based on namespace
  -- Keep .hs suffix since bild expects it
  let namespace = fromMaybe "." (TaskCore.taskNamespace task)
  let testTarget = Text.unpack namespace

  putText ("Testing: " <> Text.pack testTarget)

  -- Run bild --test on the namespace
  (testCode, testOut, testErr) <-
    Process.readProcessWithExitCode
      "bild"
      ["--test", testTarget]
      ""

  case testCode of
    Exit.ExitSuccess -> do
      putText "✓ Tests passed. Accepting task."
      TaskCore.clearRetryContext tid
      TaskCore.updateTaskStatus tid TaskCore.Done []
      putText ("Task " <> tid <> " marked as Done.")
    Exit.ExitFailure _ -> do
      putText "✗ Tests failed. Rejecting task."
      let reason = "Test failure:\n" <> Text.pack testOut <> Text.pack testErr

      maybeCtx <- TaskCore.getRetryContext tid
      let attempt = maybe 1 (\ctx -> TaskCore.retryAttempt ctx + 1) maybeCtx

      if attempt > 3
        then do
          putText "\nTask has failed 3 times. Marking as NeedsHuman."
          TaskCore.updateTaskStatus tid TaskCore.Open []
          putText ("Task " <> tid <> " needs human intervention (3 failed attempts).")
        else do
          TaskCore.setRetryContext
            TaskCore.RetryContext
              { TaskCore.retryTaskId = tid,
                TaskCore.retryOriginalCommit = Text.pack commitSha,
                TaskCore.retryConflictFiles = [],
                TaskCore.retryAttempt = attempt,
                TaskCore.retryReason = reason
              }
          TaskCore.updateTaskStatus tid TaskCore.Open []
          putText ("Task " <> tid <> " reopened (attempt " <> tshow attempt <> "/3).")

-- | Interactive review with user prompts
interactiveReview :: Text -> String -> IO ()
interactiveReview tid commitSha = do
  putText "\n=== Diff for this task ===\n"
  _ <- Process.rawSystem "git" ["show", commitSha]

  putText "\n[a]ccept / [r]eject / [s]kip? "
  IO.hFlush IO.stdout
  choice <- getLine

  case Text.toLower choice of
    c
      | "a" `Text.isPrefixOf` c -> do
          TaskCore.clearRetryContext tid
          TaskCore.updateTaskStatus tid TaskCore.Done []
          putText ("Task " <> tid <> " marked as Done.")
      | "r" `Text.isPrefixOf` c -> do
          putText "Enter rejection reason: "
          IO.hFlush IO.stdout
          reason <- getLine
          -- Save rejection as retry context
          maybeCtx <- TaskCore.getRetryContext tid
          let attempt = maybe 1 (\ctx -> TaskCore.retryAttempt ctx + 1) maybeCtx
          TaskCore.setRetryContext
            TaskCore.RetryContext
              { TaskCore.retryTaskId = tid,
                TaskCore.retryOriginalCommit = Text.pack commitSha,
                TaskCore.retryConflictFiles = [],
                TaskCore.retryAttempt = attempt,
                TaskCore.retryReason = "rejected: " <> reason
              }
          TaskCore.updateTaskStatus tid TaskCore.Open []
          putText ("Task " <> tid <> " reopened (attempt " <> tshow attempt <> "/3).")
      | otherwise -> putText "Skipped; no status change."

-- | Check if a commit can be cleanly cherry-picked onto live
-- Returns Nothing if clean, Just [conflicting files] if conflict
checkMergeConflict :: String -> IO (Maybe [Text])
checkMergeConflict commitSha = do
  -- Save current state
  (_, _, _) <- Process.readProcessWithExitCode "git" ["branch", "--show-current"] ""
  (_, origHead, _) <- Process.readProcessWithExitCode "git" ["rev-parse", "HEAD"] ""

  -- Try cherry-pick
  (cpCode, _, cpErr) <-
    Process.readProcessWithExitCode
      "git"
      ["cherry-pick", "--no-commit", commitSha]
      ""

  -- Always abort/reset regardless of result
  _ <- Process.readProcessWithExitCode "git" ["cherry-pick", "--abort"] ""
  _ <- Process.readProcessWithExitCode "git" ["reset", "--hard", List.head (List.lines origHead)] ""

  case cpCode of
    Exit.ExitSuccess -> pure Nothing
    Exit.ExitFailure _ -> do
      -- Parse conflict files from error message
      let errLines = Text.lines (Text.pack cpErr)
          conflictLines = filter (Text.isPrefixOf "CONFLICT") errLines
          -- Extract file names (rough parsing)
          files = mapMaybe extractConflictFile conflictLines
      pure (Just (if null files then ["(unknown files)"] else files))

extractConflictFile :: Text -> Maybe Text
extractConflictFile line =
  -- CONFLICT (content): Merge conflict in path/to/file.hs
  case Text.breakOn "Merge conflict in " line of
    (_, rest)
      | not (Text.null rest) -> Just (Text.strip (Text.drop 18 rest))
    _ -> case Text.breakOn "in " line of
      (_, rest)
        | not (Text.null rest) -> Just (Text.strip (Text.drop 3 rest))
      _ -> Nothing

test :: Test.Tree
test =
  Test.group
    "Omni.Jr"
    [ Test.unit "can run tests" <| True Test.@?= True,
      Test.unit "can parse work command" <| do
        let result = Docopt.parseArgs help ["work"]
        case result of
          Left err -> Test.assertFailure <| "Failed to parse 'work': " <> show err
          Right args -> args `Cli.has` Cli.command "work" Test.@?= True,
      Test.unit "can parse work command with task id" <| do
        let result = Docopt.parseArgs help ["work", "t-123"]
        case result of
          Left err -> Test.assertFailure <| "Failed to parse 'work t-123': " <> show err
          Right args -> do
            args `Cli.has` Cli.command "work" Test.@?= True
            Cli.getArg args (Cli.argument "task-id") Test.@?= Just "t-123"
    ]