{-# 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 Data.IORef (modifyIORef', newIORef, readIORef)
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
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.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))

-- | Format guardrail result for logging
formatGuardrailResult :: Engine.GuardrailResult -> Text
formatGuardrailResult Engine.GuardrailOk = "OK"
formatGuardrailResult (Engine.GuardrailCostExceeded actual limit) =
  "Cost exceeded: " <> tshow actual <> " cents (limit: " <> tshow limit <> ")"
formatGuardrailResult (Engine.GuardrailTokensExceeded actual limit) =
  "Tokens exceeded: " <> tshow actual <> " (limit: " <> tshow limit <> ")"
formatGuardrailResult (Engine.GuardrailDuplicateToolCalls tool count) =
  "Duplicate tool calls: " <> tool <> " called " <> tshow count <> " times"
formatGuardrailResult (Engine.GuardrailTestFailures count) =
  "Test failures: " <> tshow count <> " failures"

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.updateTaskStatusWithActor tid TaskCore.InProgress [] TaskCore.Junior
  say "[worker] Status -> InProgress"

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

  say "[worker] Starting engine..."
  (exitCode, output, costCents) <- runWithEngine worker repo task

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

  -- Update the activity record with metrics (convert Double to Int by rounding)
  TaskCore.updateActivityMetrics activityId Nothing (Just endTime) (Just (round 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.updateTaskStatusWithActor tid TaskCore.Open [] TaskCore.Junior
            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.updateTaskStatusWithActor tid TaskCore.Open [] TaskCore.Junior
              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.updateTaskStatusWithActor tid TaskCore.Done [] TaskCore.Junior
          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.updateTaskStatusWithActor tid TaskCore.Review [] TaskCore.Junior
          say ("[worker] ✓ Task " <> tid <> " -> Review")
          unless quiet <| AgentLog.update (\s -> s {AgentLog.statusTask = Nothing})
    Exit.ExitFailure code -> do
      say ("[worker] Engine 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 (engine 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)

-- | Run task using native Engine
-- Returns (ExitCode, output text, cost in cents)
runWithEngine :: Core.Worker -> FilePath -> TaskCore.Task -> IO (Exit.ExitCode, Text, Double)
runWithEngine worker 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)

      -- Get progress from database (checkpoint events from previous sessions)
      progressContent <- TaskCore.getProgressSummary (TaskCore.taskId task)

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

      -- Add progress context if present
      let progressPrompt = buildProgressPrompt progressContent

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

      let prompt = basePrompt <> progressPrompt <> 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

      -- Generate session ID for event logging
      sessionId <- TaskCore.generateSessionId
      let tid = TaskCore.taskId task

      -- Helper to log events to DB
      -- For text content, store as-is; for structured data, JSON-encode
      let logJuniorEvent eventType content = TaskCore.insertAgentEvent tid sessionId eventType content TaskCore.Junior
          logJuniorJson eventType value = do
            let contentJson = TE.decodeUtf8 (BSL.toStrict (Aeson.encode value))
            TaskCore.insertAgentEvent tid sessionId eventType contentJson TaskCore.Junior
          logSystemEvent eventType content = TaskCore.insertAgentEvent tid sessionId eventType content TaskCore.System

      -- Build Engine config with callbacks
      totalCostRef <- newIORef (0 :: Double)
      let quiet = Core.workerQuiet worker
          sayLog msg = if quiet then putText msg else AgentLog.log msg
          engineCfg =
            Engine.EngineConfig
              { Engine.engineLLM =
                  Engine.defaultLLM
                    { Engine.llmApiKey = Text.pack apiKey
                    },
                Engine.engineOnCost = \tokens cost -> do
                  modifyIORef' totalCostRef (+ cost)
                  sayLog <| "Cost: " <> tshow cost <> " cents (" <> tshow tokens <> " tokens)"
                  logJuniorJson "Cost" (Aeson.object [("tokens", Aeson.toJSON tokens), ("cents", Aeson.toJSON cost)]),
                Engine.engineOnActivity = \activity -> do
                  sayLog <| "[engine] " <> activity,
                Engine.engineOnToolCall = \toolName args -> do
                  sayLog <| "[tool] " <> toolName
                  logJuniorEvent "ToolCall" (toolName <> ": " <> args),
                Engine.engineOnAssistant = \msg -> do
                  sayLog <| "[assistant] " <> Text.take 200 msg
                  logJuniorEvent "Assistant" msg,
                Engine.engineOnToolResult = \toolName success output -> do
                  let statusStr = if success then "ok" else "failed"
                  sayLog <| "[result] " <> toolName <> " (" <> statusStr <> "): " <> Text.take 100 output
                  logJuniorEvent "ToolResult" output,
                Engine.engineOnComplete = do
                  sayLog "[engine] Complete"
                  logJuniorEvent "Complete" "",
                Engine.engineOnError = \err -> do
                  sayLog <| "[error] " <> err
                  logJuniorEvent "Error" err,
                Engine.engineOnGuardrail = \guardrailResult -> do
                  let guardrailMsg = formatGuardrailResult guardrailResult
                      contentJson = TE.decodeUtf8 (BSL.toStrict (Aeson.encode guardrailResult))
                  sayLog <| "[guardrail] " <> guardrailMsg
                  logSystemEvent "Guardrail" contentJson
              }

      -- Build Agent config with guardrails
      let guardrails =
            Engine.Guardrails
              { Engine.guardrailMaxCostCents = 200.0,
                Engine.guardrailMaxTokens = 1000000,
                Engine.guardrailMaxDuplicateToolCalls = 5,
                Engine.guardrailMaxTestFailures = 3
              }
          agentCfg =
            Engine.AgentConfig
              { Engine.agentModel = model,
                Engine.agentTools = Tools.allTools,
                Engine.agentSystemPrompt = systemPrompt,
                Engine.agentMaxIterations = 100,
                Engine.agentGuardrails = guardrails
              }

      -- 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 an autonomous Worker Agent.\n"
    <> "Your goal is to implement the following task:\n\n"
    <> formatTask task
    <> "\n\nCRITICAL INSTRUCTIONS:\n"
    <> "1. Read AGENTS.md first to understand the codebase conventions.\n"
    <> "2. Complete ONE logical change (e.g., update schema + call sites + tests).\n"
    <> "3. Run 'bild --test "
    <> ns
    <> "' ONCE after implementing.\n"
    <> "4. If tests pass, you are DONE - stop immediately.\n"
    <> "5. If tests fail, fix the issue and run tests again.\n"
    <> "6. If tests fail 3 times on the same issue, STOP - the task will be marked for human review.\n"
    <> "7. Do NOT update task status or manage git - the worker handles that.\n\n"
    <> "AUTONOMOUS OPERATION (NO HUMAN IN LOOP):\n"
    <> "- You are running autonomously without human intervention\n"
    <> "- There is NO human to ask questions or get clarification from\n"
    <> "- Make reasonable decisions based on the task description\n"
    <> "- If something is truly ambiguous, implement the most straightforward interpretation\n"
    <> "- Guardrails will stop you if you exceed cost/token budgets or make repeated mistakes\n\n"
    <> "BUILD SYSTEM NOTES:\n"
    <> "- 'bild --test "
    <> ns
    <> "' tests ALL dependencies transitively - run it ONCE, not per-file\n"
    <> "- Do NOT run bild --test on individual files separately\n"
    <> "- Once tests pass, STOP - do not continue adding features or re-running tests\n"
    <> "- Use 'lint --fix' for formatting issues (not hlint directly)\n\n"
    <> "EFFICIENCY REQUIREMENTS:\n"
    <> "- Do not repeat the same action multiple times\n"
    <> "- Do not re-run passing tests\n"
    <> "- Do not test files individually when namespace test covers them\n"
    <> "- Aim to complete the task in under 50 tool calls\n\n"
    <> "Context:\n"
    <> "- Working directory: "
    <> Text.pack repo
    <> "\n"
    <> "- Namespace: "
    <> ns
    <> "\n"

-- | Build progress context prompt
buildProgressPrompt :: Maybe Text -> Text
buildProgressPrompt Nothing = ""
buildProgressPrompt (Just progress) =
  "\n\n## PROGRESS FROM PREVIOUS SESSIONS (from database)\n\n"
    <> "This task has been worked on before. Here are the checkpoint notes:\n\n"
    <> progress
    <> "\n\n"
    <> "IMPORTANT:\n"
    <> "- Review these checkpoints to understand what's already done\n"
    <> "- Do NOT repeat work that's already completed\n"
    <> "- If the task appears complete, verify tests pass and exit\n\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 (1-5 scale)
-- Uses OpenRouter model identifiers for Claude models
selectModel :: TaskCore.Task -> Text
selectModel task = selectModelByComplexity (TaskCore.taskComplexity task)

-- | Select model based on complexity level
selectModelByComplexity :: Maybe Int -> Text
selectModelByComplexity Nothing = "anthropic/claude-sonnet-4.5"
selectModelByComplexity (Just 1) = "anthropic/claude-haiku-4.5"
selectModelByComplexity (Just 2) = "anthropic/claude-haiku-4.5"
selectModelByComplexity (Just 3) = "anthropic/claude-sonnet-4.5"
selectModelByComplexity (Just 4) = "anthropic/claude-sonnet-4.5"
selectModelByComplexity (Just 5) = "anthropic/claude-opus-4.5"
selectModelByComplexity (Just _) = "anthropic/claude-sonnet-4.5"

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 agentOutput =
  let tid = TaskCore.taskId task
      subject = cleanSubject (TaskCore.taskTitle task)
      body = cleanBody agentOutput
   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) <> "]"
       )