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

module Omni.Agent.Worker where

import Alpha
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Key as AesonKey
import qualified Data.ByteString.Lazy as BSL
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Data.Text.IO as TIO
import qualified Data.Time
import qualified Omni.Agent.Core as Core
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Log as AgentLog
import qualified Omni.Agent.Tools as Tools
import qualified Omni.Fact as Fact
import qualified Omni.Task.Core as TaskCore
import qualified System.Directory as Directory
import qualified System.Environment as Env
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
  if Core.workerQuiet worker
    then putText ("[worker] Starting for " <> Core.workerName worker)
    else do
      AgentLog.init (Core.workerName worker)
      AgentLog.log ("[worker] Starting for " <> Core.workerName worker)
  case maybeTaskId of
    Just tid -> logMsg worker ("[worker] Target task: " <> tid)
    Nothing -> logMsg worker "[worker] No specific task, will pick from ready queue"
  runOnce worker maybeTaskId

-- | Log message respecting quiet mode
logMsg :: Core.Worker -> Text -> IO ()
logMsg worker msg =
  if Core.workerQuiet worker
    then putText msg
    else AgentLog.log msg

-- | Convert key-value pairs to JSON metadata string
toMetadata :: [(Text, Text)] -> Text
toMetadata pairs =
  let obj = Aeson.object [(AesonKey.fromText k, Aeson.String v) | (k, v) <- pairs]
   in TE.decodeUtf8 (BSL.toStrict (Aeson.encode obj))

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
          unless (Core.workerQuiet worker) <| AgentLog.updateActivity ("Task " <> tid <> " not found.")
          logMsg worker ("[worker] Task " <> tid <> " not found.")
        Nothing -> do
          unless (Core.workerQuiet worker) <| AgentLog.updateActivity "No work found."
          logMsg worker "[worker] 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
  let quiet = Core.workerQuiet worker
  let say = logMsg worker

  unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Just tid})
  say ("[worker] Claiming task " <> tid)

  -- Claim task
  TaskCore.logActivity tid TaskCore.Claiming Nothing
  TaskCore.updateTaskStatus tid TaskCore.InProgress []
  say "[worker] Status -> InProgress"

  -- Check if we should use native engine
  useEngine <- shouldUseEngine

  -- Run agent with timing
  startTime <- Data.Time.getCurrentTime
  activityId <- TaskCore.logActivityWithMetrics tid TaskCore.Running Nothing Nothing (Just startTime) Nothing Nothing Nothing

  (exitCode, output, maybeCost) <-
    if useEngine
      then do
        say "[worker] Starting native engine..."
        (code, out, cost) <- runWithEngine repo task
        pure (code, out, Just cost)
      else do
        say "[worker] Starting amp..."
        (code, out) <- runAmp repo task
        pure (code, out, Nothing)

  endTime <- Data.Time.getCurrentTime
  say ("[worker] Agent exited with: " <> tshow exitCode)

  -- Capture metrics - from engine result or agent log
  (threadUrl, costCents) <- case maybeCost of
    Just engineCost -> pure (Nothing, Just engineCost)
    Nothing -> do
      status <- AgentLog.getStatus
      let url = ("https://ampcode.com/threads/" <>) </ AgentLog.statusThread status
      let cost = Just <| floor (AgentLog.statusCredits status * 100)
      pure (url, cost)

  -- Update the activity record with metrics
  TaskCore.updateActivityMetrics activityId threadUrl (Just endTime) costCents Nothing

  case exitCode of
    Exit.ExitSuccess -> do
      TaskCore.logActivity tid TaskCore.Reviewing Nothing
      say "[worker] Running formatters..."
      _ <- runFormatters repo

      -- Try to commit (this runs git hooks which may fail)
      let commitMsg = formatCommitMessage task output
      say "[worker] Attempting commit..."
      commitResult <- tryCommit repo commitMsg

      case commitResult of
        CommitFailed commitErr -> do
          say ("[worker] Commit failed: " <> commitErr)

          -- Save failure context and reopen task for retry
          maybeCtx <- TaskCore.getRetryContext tid
          let attempt = maybe 1 (\c -> TaskCore.retryAttempt c + 1) maybeCtx

          if attempt > 3
            then do
              say "[worker] Task failed 3 times, needs human intervention"
              TaskCore.logActivity tid TaskCore.Failed (Just (toMetadata [("reason", "max_retries_exceeded")]))
              TaskCore.updateTaskStatus tid TaskCore.Open []
            else do
              let currentReason = "attempt " <> tshow attempt <> ": commit_failed: " <> commitErr
              let accumulatedReason = case maybeCtx of
                    Nothing -> currentReason
                    Just ctx -> TaskCore.retryReason ctx <> "\n" <> currentReason
              TaskCore.setRetryContext
                TaskCore.RetryContext
                  { TaskCore.retryTaskId = tid,
                    TaskCore.retryOriginalCommit = "",
                    TaskCore.retryConflictFiles = [],
                    TaskCore.retryAttempt = attempt,
                    TaskCore.retryReason = accumulatedReason,
                    TaskCore.retryNotes = maybeCtx +> TaskCore.retryNotes
                  }
              TaskCore.logActivity tid TaskCore.Retrying (Just (toMetadata [("attempt", tshow attempt)]))
              TaskCore.updateTaskStatus tid TaskCore.Open []
              say ("[worker] Task reopened (attempt " <> tshow attempt <> "/3)")
        NoChanges -> do
          -- No changes = task already implemented, mark as Done
          say "[worker] No changes to commit - task already done"
          TaskCore.clearRetryContext tid
          TaskCore.logActivity tid TaskCore.Completed (Just (toMetadata [("result", "no_changes")]))
          TaskCore.updateTaskStatus tid TaskCore.Done []
          say ("[worker] ✓ Task " <> tid <> " -> Done (no changes)")
          unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Nothing})
        CommitSuccess -> do
          -- Commit succeeded, set to Review
          TaskCore.logActivity tid TaskCore.Completed (Just (toMetadata [("result", "committed")]))
          TaskCore.updateTaskStatus tid TaskCore.Review []
          say ("[worker] ✓ Task " <> tid <> " -> Review")
          unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Nothing})
    Exit.ExitFailure code -> do
      say ("[worker] Amp failed with code " <> tshow code)
      TaskCore.logActivity tid TaskCore.Failed (Just (toMetadata [("exit_code", tshow code)]))
      -- Don't set back to Open here - leave in InProgress for debugging
      say "[worker] Task left in InProgress (amp failure)"

-- | Run lint --fix to format and fix lint issues
runFormatters :: FilePath -> IO (Either Text ())
runFormatters repo = do
  let cmd = (Process.proc "lint" ["--fix"]) {Process.cwd = Just repo}
  (code, _, _) <- Process.readCreateProcessWithExitCode cmd ""
  case code of
    Exit.ExitSuccess -> pure (Right ())
    Exit.ExitFailure _ -> pure (Right ()) -- lint --fix may exit non-zero but still fix things

data CommitResult = CommitSuccess | NoChanges | CommitFailed Text
  deriving (Show, Eq)

-- | Try to commit, returning result
tryCommit :: FilePath -> Text -> IO CommitResult
tryCommit repo msg = do
  -- Stage all changes
  let addCmd = (Process.proc "git" ["add", "."]) {Process.cwd = Just repo}
  (addCode, _, addErr) <- Process.readCreateProcessWithExitCode addCmd ""
  case addCode of
    Exit.ExitFailure _ -> pure <| CommitFailed (Text.pack addErr)
    Exit.ExitSuccess -> do
      -- Check for changes
      let checkCmd = (Process.proc "git" ["diff", "--cached", "--quiet"]) {Process.cwd = Just repo}
      (checkCode, _, _) <- Process.readCreateProcessWithExitCode checkCmd ""
      case checkCode of
        Exit.ExitSuccess -> pure NoChanges
        Exit.ExitFailure 1 -> do
          -- There are changes, commit them
          let commitCmd = (Process.proc "git" ["commit", "-m", Text.unpack msg]) {Process.cwd = Just repo}
          (commitCode, _, commitErr) <- Process.readCreateProcessWithExitCode commitCmd ""
          case commitCode of
            Exit.ExitSuccess -> pure CommitSuccess
            Exit.ExitFailure _ -> pure <| CommitFailed (Text.pack commitErr)
        Exit.ExitFailure c -> pure <| CommitFailed ("git diff failed with code " <> tshow c)

runAmp :: FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text)
runAmp repo task = do
  -- Check for retry context
  maybeRetry <- TaskCore.getRetryContext (TaskCore.taskId task)

  let ns = fromMaybe "." (TaskCore.taskNamespace task)
  let basePrompt =
        "You are a Worker Agent.\n"
          <> "Your goal is to implement the following task:\n\n"
          <> formatTask task
          <> "\n\nCRITICAL INSTRUCTIONS:\n"
          <> "1. Analyze the codebase to understand where to make changes.\n"
          <> "2. Implement the changes by editing files.\n"
          <> "3. BEFORE finishing, you MUST run: bild --test "
          <> ns
          <> "\n"
          <> "4. Fix ALL errors from bild --test (including hlint suggestions).\n"
          <> "5. Keep running bild --test until it passes with no errors.\n"
          <> "6. Do NOT update task status or manage git.\n"
          <> "7. Only exit after bild --test passes.\n\n"
          <> "IMPORTANT: The git commit will fail if hlint finds issues.\n"
          <> "You must fix hlint suggestions like:\n"
          <> "- 'Use list comprehension' -> use [x | cond] instead of if/else\n"
          <> "- 'Avoid lambda' -> use function composition\n"
          <> "- 'Redundant bracket' -> remove unnecessary parens\n\n"
          <> "Context:\n"
          <> "- Working directory: "
          <> Text.pack repo
          <> "\n"
          <> "- Namespace: "
          <> ns
          <> "\n"

  -- Add retry context if present
  let retryPrompt = case maybeRetry of
        Nothing -> ""
        Just ctx ->
          "\n\n## RETRY CONTEXT (IMPORTANT)\n\n"
            <> "This task was previously attempted but failed. Attempt: "
            <> tshow (TaskCore.retryAttempt ctx)
            <> "/3\n"
            <> "Reason: "
            <> TaskCore.retryReason ctx
            <> "\n\n"
            <> ( if null (TaskCore.retryConflictFiles ctx)
                   then ""
                   else
                     "Conflicting files from previous attempt:\n"
                       <> Text.unlines (map ("  - " <>) (TaskCore.retryConflictFiles ctx))
                       <> "\n"
               )
            <> "Original commit: "
            <> TaskCore.retryOriginalCommit ctx
            <> "\n\n"
            <> maybe "" (\notes -> "## HUMAN NOTES/GUIDANCE\n\n" <> notes <> "\n\n") (TaskCore.retryNotes ctx)
            <> "INSTRUCTIONS FOR RETRY:\n"
            <> "- The codebase has changed since your last attempt\n"
            <> "- Re-implement this task on top of the CURRENT codebase\n"
            <> "- If there were merge conflicts, the conflicting files may have been modified by others\n"
            <> "- Review the current state of those files before making changes\n"

  let prompt = basePrompt <> retryPrompt

  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

  -- Get relevant facts from the knowledge base
  relevantFacts <- getRelevantFacts task
  let factsSection = formatFacts relevantFacts

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

  -- 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 = ["--try-opus", "--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)

-- | Check if we should use native engine instead of amp subprocess
shouldUseEngine :: IO Bool
shouldUseEngine = do
  env <- Env.lookupEnv "JR_USE_ENGINE"
  pure <| env == Just "1"

-- | Run task using native Engine instead of amp subprocess
-- Returns (ExitCode, output text, cost in cents)
runWithEngine :: FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text, Int)
runWithEngine repo task = do
  -- Read API key from environment
  maybeApiKey <- Env.lookupEnv "OPENROUTER_API_KEY"
  case maybeApiKey of
    Nothing -> pure (Exit.ExitFailure 1, "OPENROUTER_API_KEY not set", 0)
    Just apiKey -> do
      -- Check for retry context
      maybeRetry <- TaskCore.getRetryContext (TaskCore.taskId task)

      -- Build the full prompt (same as runAmp)
      let ns = fromMaybe "." (TaskCore.taskNamespace task)
      let basePrompt = buildBasePrompt task ns repo

      -- Add retry context if present
      let retryPrompt = buildRetryPrompt maybeRetry

      let prompt = basePrompt <> retryPrompt

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

      -- Get relevant facts from the knowledge base
      relevantFacts <- getRelevantFacts task
      let factsSection = formatFacts relevantFacts

      -- Build system prompt
      let systemPrompt =
            prompt
              <> "\n\nREPOSITORY GUIDELINES (AGENTS.md):\n"
              <> agentsMd
              <> factsSection

      -- Build user prompt from task comments
      let userPrompt = formatTask task

      -- Select model based on task complexity (simple heuristic)
      let model = selectModel task

      -- Build Engine config with callbacks
      totalCostRef <- newIORef (0 :: Int)
      let engineCfg =
            Engine.EngineConfig
              { Engine.engineLLM =
                  Engine.defaultLLM
                    { Engine.llmApiKey = Text.pack apiKey
                    },
                Engine.engineOnCost = \tokens cost -> do
                  modifyIORef' totalCostRef (+ cost)
                  AgentLog.log <| "Cost: " <> tshow cost <> " cents (" <> tshow tokens <> " tokens)",
                Engine.engineOnActivity = \activity -> do
                  AgentLog.log <| "[engine] " <> activity,
                Engine.engineOnToolCall = \toolName result -> do
                  AgentLog.log <| "[tool] " <> toolName <> ": " <> Text.take 100 result
              }

      -- Build Agent config
      let agentCfg =
            Engine.AgentConfig
              { Engine.agentModel = model,
                Engine.agentTools = Tools.allTools,
                Engine.agentSystemPrompt = systemPrompt,
                Engine.agentMaxIterations = 20
              }

      -- Run the agent
      result <- Engine.runAgent engineCfg agentCfg userPrompt
      totalCost <- readIORef totalCostRef

      case result of
        Left err -> pure (Exit.ExitFailure 1, "Engine error: " <> err, totalCost)
        Right agentResult -> do
          let output = Engine.resultFinalMessage agentResult
          pure (Exit.ExitSuccess, output, totalCost)

-- | Build the base prompt for the agent
buildBasePrompt :: TaskCore.Task -> Text -> FilePath -> Text
buildBasePrompt task ns repo =
  "You are a Worker Agent.\n"
    <> "Your goal is to implement the following task:\n\n"
    <> formatTask task
    <> "\n\nCRITICAL INSTRUCTIONS:\n"
    <> "1. Analyze the codebase to understand where to make changes.\n"
    <> "2. Implement the changes by editing files.\n"
    <> "3. BEFORE finishing, you MUST run: bild --test "
    <> ns
    <> "\n"
    <> "4. Fix ALL errors from bild --test (including hlint suggestions).\n"
    <> "5. Keep running bild --test until it passes with no errors.\n"
    <> "6. Do NOT update task status or manage git.\n"
    <> "7. Only exit after bild --test passes.\n\n"
    <> "IMPORTANT: The git commit will fail if hlint finds issues.\n"
    <> "You must fix hlint suggestions like:\n"
    <> "- 'Use list comprehension' -> use [x | cond] instead of if/else\n"
    <> "- 'Avoid lambda' -> use function composition\n"
    <> "- 'Redundant bracket' -> remove unnecessary parens\n\n"
    <> "Context:\n"
    <> "- Working directory: "
    <> Text.pack repo
    <> "\n"
    <> "- Namespace: "
    <> ns
    <> "\n"

-- | Build retry context prompt
buildRetryPrompt :: Maybe TaskCore.RetryContext -> Text
buildRetryPrompt Nothing = ""
buildRetryPrompt (Just ctx) =
  "\n\n## RETRY CONTEXT (IMPORTANT)\n\n"
    <> "This task was previously attempted but failed. Attempt: "
    <> tshow (TaskCore.retryAttempt ctx)
    <> "/3\n"
    <> "Reason: "
    <> TaskCore.retryReason ctx
    <> "\n\n"
    <> ( if null (TaskCore.retryConflictFiles ctx)
           then ""
           else
             "Conflicting files from previous attempt:\n"
               <> Text.unlines (map ("  - " <>) (TaskCore.retryConflictFiles ctx))
               <> "\n"
       )
    <> "Original commit: "
    <> TaskCore.retryOriginalCommit ctx
    <> "\n\n"
    <> maybe "" (\notes -> "## HUMAN NOTES/GUIDANCE\n\n" <> notes <> "\n\n") (TaskCore.retryNotes ctx)
    <> "INSTRUCTIONS FOR RETRY:\n"
    <> "- The codebase has changed since your last attempt\n"
    <> "- Re-implement this task on top of the CURRENT codebase\n"
    <> "- If there were merge conflicts, the conflicting files may have been modified by others\n"
    <> "- Review the current state of those files before making changes\n"

-- | Select model based on task complexity
-- Currently always uses claude-sonnet-4, but can be extended for model selection
selectModel :: TaskCore.Task -> Text
selectModel _ = "anthropic/claude-sonnet-4-20250514"

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"
    <> (if Text.null (TaskCore.taskDescription t) then "" else "Description:\n" <> TaskCore.taskDescription t <> "\n\n")
    <> formatDeps (TaskCore.taskDependencies t)
    <> formatComments (TaskCore.taskComments t)
  where
    formatDeps [] = ""
    formatDeps deps = "\nDependencies:\n" <> Text.unlines (map formatDep deps)
    formatDep dep = "  - " <> TaskCore.depId dep <> " [" <> Text.pack (show (TaskCore.depType dep)) <> "]"
    formatComments [] = ""
    formatComments cs = "\nComments/Notes:\n" <> Text.unlines (map formatComment cs)
    formatComment c = "  [" <> Text.pack (show (TaskCore.commentCreatedAt c)) <> "] " <> TaskCore.commentText c

formatCommitMessage :: TaskCore.Task -> Text -> Text
formatCommitMessage task ampOutput =
  let tid = TaskCore.taskId task
      subject = cleanSubject (TaskCore.taskTitle task)
      body = cleanBody ampOutput
   in if Text.null body
        then subject <> "\n\nTask-Id: " <> tid
        else subject <> "\n\n" <> body <> "\n\nTask-Id: " <> tid
  where
    cleanSubject s =
      let trailingPunct = ['.', ':', '!', '?', ',', ';', ' ', '-']
          stripped = Text.dropWhileEnd (`elem` trailingPunct) s
          truncated = Text.take 72 stripped
          noPunct = Text.dropWhileEnd (`elem` trailingPunct) truncated
          capitalized = case Text.uncons noPunct of
            Just (c, rest) -> Text.cons (toUpper c) rest
            Nothing -> noPunct
       in capitalized

    cleanBody :: Text -> Text
    cleanBody output =
      let stripped = Text.strip output
       in if Text.null stripped
            then ""
            else
              let lns = Text.lines stripped
                  cleaned = [Text.take 72 ln | ln <- lns]
               in Text.intercalate "\n" cleaned

-- | Get facts relevant to a task based on namespace/project
getRelevantFacts :: TaskCore.Task -> IO [TaskCore.Fact]
getRelevantFacts task = do
  let namespace = fromMaybe "Omni" (TaskCore.taskNamespace task)
  projectFacts <- Fact.getFactsByProject namespace
  let sorted = List.sortBy (comparing (Down <. TaskCore.factConfidence)) projectFacts
  pure (take 10 sorted)

-- | Format facts for inclusion in the prompt
formatFacts :: [TaskCore.Fact] -> Text
formatFacts [] = ""
formatFacts facts =
  Text.unlines
    [ "\n\nKNOWLEDGE BASE FACTS:",
      "(These are learned patterns/conventions from previous work)",
      ""
    ]
    <> Text.unlines (map formatFact facts)

-- | Format a single fact for the prompt
formatFact :: TaskCore.Fact -> Text
formatFact f =
  "- "
    <> TaskCore.factContent f
    <> ( if null (TaskCore.factRelatedFiles f)
           then ""
           else " [" <> Text.intercalate ", " (TaskCore.factRelatedFiles f) <> "]"
       )

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