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

module Omni.Agent.Worker where

import Alpha
import qualified Data.Text as Text
import qualified Data.Text.IO as TIO
import qualified Omni.Agent.Core as Core
import qualified Omni.Agent.Git as Git
import qualified Omni.Agent.Log as AgentLog
import qualified Omni.Task.Core as TaskCore
import qualified System.Directory as Directory
import qualified System.Exit as Exit
import System.FilePath ((</>))
import qualified System.IO as IO
import qualified System.Process as Process

start :: Core.Worker -> Maybe Text -> IO ()
start worker maybeTaskId = do
  AgentLog.init (Core.workerName worker)
  AgentLog.log ("Worker starting for " <> Core.workerName worker)
  runOnce worker maybeTaskId

runOnce :: Core.Worker -> Maybe Text -> IO ()
runOnce worker maybeTaskId = do
  -- Find work
  targetTask <- case maybeTaskId of
    Just tid -> do
      TaskCore.findTask tid </ TaskCore.loadTasks
    Nothing -> do
      readyTasks <- TaskCore.getReadyTasks
      case readyTasks of
        [] -> pure Nothing
        (task : _) -> pure (Just task)

  case targetTask of
    Nothing -> do
      case maybeTaskId of
        Just tid -> do
          AgentLog.updateActivity ("Task " <> tid <> " not found.")
          AgentLog.log ("Task " <> tid <> " not found.")
        Nothing -> do
          AgentLog.updateActivity "No work found."
          AgentLog.log "No ready tasks found."
    Just task -> do
      processTask worker task

processTask :: Core.Worker -> TaskCore.Task -> IO ()
processTask worker task = do
  let repo = Core.workerPath worker
  let tid = TaskCore.taskId task

  AgentLog.update (\s -> s {AgentLog.statusTask = Just tid})
  AgentLog.updateActivity ("Claiming task " <> tid)

  -- Claim task
  TaskCore.updateTaskStatus tid TaskCore.InProgress []

  -- Run Amp
  AgentLog.updateActivity "Running Amp agent..."
  (exitCode, output) <- runAmp repo task

  case exitCode of
    Exit.ExitSuccess -> do
      AgentLog.log "Agent finished successfully"

      -- Update status to Review
      TaskCore.updateTaskStatus tid TaskCore.Review []

      -- Commit changes using Amp output
      let commitMsg = "feat: implement " <> tid <> "\n\n" <> output
      Git.commit repo commitMsg

      -- Submit for review
      AgentLog.updateActivity "Task completed"

      AgentLog.log ("[✓] Task " <> tid <> " completed")
      AgentLog.update (\s -> s {AgentLog.statusTask = Nothing})
    Exit.ExitFailure code -> do
      AgentLog.log ("Agent failed with code " <> tshow code)
      AgentLog.updateActivity "Agent failed, retrying..."
      threadDelay (10 * 1000000) -- Sleep 10s

runAmp :: FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text)
runAmp repo task = do
  let prompt =
        "You are a Worker Agent.\n"
          <> "Your goal is to implement the following task:\n\n"
          <> formatTask task
          <> "\n\nINSTRUCTIONS:\n"
          <> "1. Analyze the codebase (use finder/Grep) to understand where to make changes.\n"
          <> "2. Implement the changes by editing files.\n"
          <> "3. Run tests to verify your work (e.g., 'bild --test Omni/Namespace').\n"
          <> "4. Fix any errors found during testing.\n"
          <> "5. Do NOT update the task status or manage git branches.\n"
          <> "6. When finished and tested, exit.\n\n"
          <> "Context:\n"
          <> "- You are working in '"
          <> Text.pack repo
          <> "'.\n"
          <> "- The task is in namespace '"
          <> fromMaybe "root" (TaskCore.taskNamespace task)
          <> "'.\n"

  let logFile = repo </> "_/llm/amp.log"

  -- Read AGENTS.md
  agentsMd <-
    fmap (fromMaybe "") <| do
      exists <- Directory.doesFileExist (repo </> "AGENTS.md")
      if exists
        then Just </ readFile (repo </> "AGENTS.md")
        else pure Nothing

  let fullPrompt =
        prompt
          <> "\n\nREPOSITORY GUIDELINES (AGENTS.md):\n"
          <> agentsMd

  -- Remove old log file
  exists <- Directory.doesFileExist logFile
  when exists (Directory.removeFile logFile)

  Directory.createDirectoryIfMissing True (repo </> "_/llm")

  -- Assume amp is in PATH
  let args = ["--log-level", "debug", "--log-file", "_/llm/amp.log", "--dangerously-allow-all", "-x", Text.unpack fullPrompt]

  let cp = (Process.proc "amp" args) {Process.cwd = Just repo, Process.std_out = Process.CreatePipe}
  (_, Just hOut, _, ph) <- Process.createProcess cp

  tid <- forkIO <| monitorLog logFile ph

  exitCode <- Process.waitForProcess ph
  output <- TIO.hGetContents hOut
  killThread tid
  pure (exitCode, output)

formatTask :: TaskCore.Task -> Text
formatTask t =
  "Task: "
    <> TaskCore.taskId t
    <> "\n"
    <> "Title:      "
    <> TaskCore.taskTitle t
    <> "\n"
    <> "Type:       "
    <> Text.pack (show (TaskCore.taskType t))
    <> "\n"
    <> "Status:     "
    <> Text.pack (show (TaskCore.taskStatus t))
    <> "\n"
    <> "Priority:   "
    <> Text.pack (show (TaskCore.taskPriority t))
    <> "\n"
    <> maybe "" (\p -> "Parent:     " <> p <> "\n") (TaskCore.taskParent t)
    <> maybe "" (\ns -> "Namespace:  " <> ns <> "\n") (TaskCore.taskNamespace t)
    <> "Created:    "
    <> Text.pack (show (TaskCore.taskCreatedAt t))
    <> "\n"
    <> "Updated:    "
    <> Text.pack (show (TaskCore.taskUpdatedAt t))
    <> "\n"
    <> maybe "" (\d -> "Description:\n" <> d <> "\n\n") (TaskCore.taskDescription t)
    <> (if null (TaskCore.taskDependencies t) then "" else "\nDependencies:\n" <> Text.unlines (map formatDep (TaskCore.taskDependencies t)))
  where
    formatDep dep = "  - " <> TaskCore.depId dep <> " [" <> Text.pack (show (TaskCore.depType dep)) <> "]"

monitorLog :: FilePath -> Process.ProcessHandle -> IO ()
monitorLog path ph = do
  waitForFile path
  IO.withFile path IO.ReadMode <| \h -> do
    IO.hSetBuffering h IO.LineBuffering
    go h
  where
    go h = do
      eof <- IO.hIsEOF h
      if eof
        then do
          mExit <- Process.getProcessExitCode ph
          case mExit of
            Nothing -> do
              threadDelay 100000 -- 0.1s
              go h
            Just _ -> pure ()
        else do
          line <- TIO.hGetLine h
          AgentLog.processLogLine line
          go h

waitForFile :: FilePath -> IO ()
waitForFile path = do
  exists <- Directory.doesFileExist path
  if exists
    then pure ()
    else do
      threadDelay 100000
      waitForFile path