{-# 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))

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"

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

      -- Read progress file if it exists
      progressContent <- readProgressFile repo (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 logEventText = TaskCore.insertAgentEvent tid sessionId
          logEventJson eventType value = do
            let contentJson = TE.decodeUtf8 (BSL.toStrict (Aeson.encode value))
            TaskCore.insertAgentEvent tid sessionId eventType contentJson

      -- 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)"
                  logEventJson "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
                  logEventText "ToolCall" (toolName <> ": " <> args),
                Engine.engineOnAssistant = \msg -> do
                  sayLog <| "[assistant] " <> Text.take 200 msg
                  logEventText "Assistant" msg,
                Engine.engineOnToolResult = \toolName success output -> do
                  let statusStr = if success then "ok" else "failed"
                  sayLog <| "[result] " <> toolName <> " (" <> statusStr <> "): " <> Text.take 100 output
                  logEventText "ToolResult" output,
                Engine.engineOnComplete = do
                  sayLog "[engine] Complete"
                  logEventText "Complete" "",
                Engine.engineOnError = \err -> do
                  sayLog <| "[error] " <> err
                  logEventText "Error" err
              }

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

      -- 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. Read AGENTS.md and any existing progress file for this task.\n"
    <> "2. Pick ONE specific change to implement (not everything at once).\n"
    <> "3. Analyze the codebase to understand where to make that change.\n"
    <> "4. Implement ONLY that one change.\n"
    <> "5. BEFORE finishing, you MUST run: bild --test "
    <> ns
    <> "\n"
    <> "6. Fix ALL errors from bild --test (including lint issues).\n"
    <> "7. Keep running bild --test until it passes with no errors.\n"
    <> "8. After tests pass, write progress to: _/llm/"
    <> TaskCore.taskId task
    <> "-progress.md\n"
    <> "9. Do NOT update task status or manage git.\n"
    <> "10. Only exit after bild --test passes and progress is saved.\n\n"
    <> "INCREMENTAL WORKFLOW (IMPORTANT):\n"
    <> "- DO NOT try to implement everything in one go\n"
    <> "- Make ONE focused change, test it, save progress, then stop\n"
    <> "- The task may be run multiple times to complete all changes\n"
    <> "- Each session should leave the code in a clean, testable state\n"
    <> "- If the task is already complete, just verify tests pass and note that in progress\n\n"
    <> "IMPORTANT: The git commit will fail if lint finds issues.\n"
    <> "You must fix all lint suggestions.\n\n"
    <> "BUILD SYSTEM NOTES:\n"
    <> "- Running 'bild --test "
    <> ns
    <> "' automatically tests ALL dependencies of that namespace\n"
    <> "- You do NOT need to run bild --test on individual files - just the main namespace ONCE\n"
    <> "- Once tests pass, do NOT re-run them unless you make more changes\n"
    <> "- The 'lint' command will be run automatically during git commit via hooks\n"
    <> "- You can run 'lint --fix' on changed files if needed, but it's optional\n\n"
    <> "Context:\n"
    <> "- Working directory: "
    <> Text.pack repo
    <> "\n"
    <> "- Namespace: "
    <> ns
    <> "\n"

-- | Read progress file for a task if it exists
readProgressFile :: FilePath -> Text -> IO (Maybe Text)
readProgressFile repo taskId = do
  let progressPath = repo </> "_" </> "llm" </> Text.unpack taskId <> "-progress.md"
  exists <- Directory.doesFileExist progressPath
  if exists
    then Just </ readFile progressPath
    else pure Nothing

-- | Build progress context prompt
buildProgressPrompt :: Maybe Text -> Text
buildProgressPrompt Nothing = ""
buildProgressPrompt (Just progress) =
  "\n\n## PROGRESS FROM PREVIOUS SESSIONS\n\n"
    <> "This task has been worked on before. Here's what has been completed:\n\n"
    <> progress
    <> "\n\n"
    <> "IMPORTANT:\n"
    <> "- Review this progress to understand what's already done\n"
    <> "- Do NOT repeat work that's already completed\n"
    <> "- Pick the NEXT logical step that hasn't been done yet\n"
    <> "- Update the progress file after completing your change\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) <> "]"
       )