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.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 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

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
  | otherwise = putText (str <| Docopt.usage help)

-- | Run the autonomous loop: work -> review -> repeat
runLoop :: Int -> IO ()
runLoop delaySec = do
  putText "[loop] Starting autonomous jr loop..."
  putText ("[loop] Delay between iterations: " <> tshow delaySec <> "s")
  go
  where
    go = do
      -- First check for tasks to review (prioritize finishing work)
      reviewResult <- reviewPending
      if reviewResult
        then do
          -- Reviewed something, continue loop immediately
          threadDelay (delaySec * 1000000)
          go
        else do
          -- No reviews, check for ready work
          readyTasks <- TaskCore.getReadyTasks
          case readyTasks of
            [] -> do
              putText "[loop] No ready tasks, no pending reviews."
            (task : _) -> do
              putText ""
              putText ("[loop] === Working on: " <> TaskCore.taskId task <> " ===")
              -- Run worker (this blocks until amp completes)
              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 = "."
                      }
              putText "[loop] Starting worker..."
              AgentWorker.start worker (Just (TaskCore.taskId task))
              putText "[loop] Worker finished."

          -- Delay and loop
          putText ("[loop] Sleeping " <> tshow delaySec <> "s...")
          threadDelay (delaySec * 1000000)
          go

    -- Returns True if a task was reviewed, False otherwise
    reviewPending :: IO Bool
    reviewPending = do
      tasks <- TaskCore.loadTasks
      let reviewTasks = filter (\t -> TaskCore.taskStatus t == TaskCore.Review) tasks
      case reviewTasks of
        [] -> pure False
        (t : _) -> do
          putText ""
          putText ("[loop] === Reviewing: " <> TaskCore.taskId t <> " ===")
          tryAutoReview (TaskCore.taskId t)
          pure True

    -- Auto-review that doesn't exit on missing commit
    tryAutoReview :: Text -> IO ()
    tryAutoReview tid = do
      tasks <- TaskCore.loadTasks
      case TaskCore.findTask tid tasks of
        Nothing -> do
          putText ("[review] Task " <> tid <> " not found.")
        Just task -> do
          let grepArg = "--grep=" <> Text.unpack tid
          (code, shaOut, _) <-
            Process.readProcessWithExitCode
              "git"
              ["log", "--pretty=format:%H", "-n", "1", grepArg]
              ""

          if code /= Exit.ExitSuccess || null shaOut
            then do
              putText "[review] No commit found for this task."
              putText "[review] Resetting to Open for retry."
              TaskCore.updateTaskStatus tid TaskCore.Open []
            else do
              let commitSha = case List.lines shaOut of
                    (x : _) -> x
                    [] -> ""

              -- Check for merge conflicts
              conflictResult <- checkMergeConflict commitSha
              case conflictResult of
                Just conflictFiles -> do
                  putText "[review] MERGE CONFLICT DETECTED"
                  traverse_ (\f -> putText ("  - " <> f)) conflictFiles
                  handleConflict tid conflictFiles commitSha
                Nothing -> do
                  autoReview tid task commitSha

-- | Handle merge conflict during review
handleConflict :: Text -> [Text] -> String -> IO ()
handleConflict tid conflictFiles commitSha = do
  maybeCtx <- TaskCore.getRetryContext tid
  let attempt = maybe 1 (\c -> TaskCore.retryAttempt c + 1) maybeCtx

  if attempt > 3
    then do
      putText "[review] Task has failed 3 times. Needs human intervention."
      TaskCore.updateTaskStatus tid TaskCore.Open []
    else do
      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 ("[review] Task " <> tid <> " returned to queue (attempt " <> tshow attempt <> "/3).")

-- | Interactive review command (jr review <task-id>)
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 ==="
          traverse_ (\f -> putText ("  - " <> f)) conflictFiles
          handleConflict tid conflictFiles commitSha
        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 "[review] Running automated review..."
  putText ("[review] Commit: " <> Text.pack (take 8 commitSha))

  -- Determine what to test based on namespace
  let namespace = fromMaybe "." (TaskCore.taskNamespace task)
  let testTarget = Text.unpack namespace

  putText ("[review] 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 "[review] ✓ Tests passed."
      TaskCore.clearRetryContext tid
      TaskCore.updateTaskStatus tid TaskCore.Done []
      putText ("[review] Task " <> tid <> " -> Done")
    Exit.ExitFailure code -> do
      putText ("[review] ✗ Tests failed (exit " <> tshow code <> ")")
      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 "[review] Task has failed 3 times. Needs human intervention."
          TaskCore.updateTaskStatus tid TaskCore.Open []
        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 ("[review] 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"
    ]