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

-- | Subagent system for spawning specialized agents.
--
-- Enables the orchestrator (Ava) to delegate focused tasks to specialized
-- subagents that run with their own tool sets and resource limits.
--
-- Key features:
-- - Role-based tool selection (WebCrawler, CodeReviewer, etc.)
-- - Per-subagent resource limits (timeout, cost, tokens)
-- - Structured result format with confidence scores
-- - No sub-subagent spawning (hierarchical control)
--
-- : out omni-agent-subagent
-- : dep aeson
-- : dep async
module Omni.Agent.Subagent
  ( -- * Types
    SubagentRole (..),
    SubagentConfig (..),
    SubagentResult (..),
    SubagentStatus (..),
    SubagentCallbacks (..),

    -- * Execution
    runSubagent,
    runSubagentWithCallbacks,

    -- * Tool
    spawnSubagentTool,

    -- * Role-specific tools
    SubagentApiKeys (..),
    toolsForRole,
    modelForRole,
    systemPromptForRole,

    -- * Defaults
    defaultSubagentConfig,
    defaultCallbacks,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Text as Text
import qualified Data.Time.Clock as Clock
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Provider as Provider
import qualified Omni.Agent.Tools as Tools
import qualified Omni.Agent.Tools.WebReader as WebReader
import qualified Omni.Agent.Tools.WebSearch as WebSearch
import qualified Omni.Test as Test

main :: IO ()
main = Test.run test

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Subagent"
    [ Test.unit "SubagentRole JSON roundtrip" <| do
        let roles = [WebCrawler, CodeReviewer, DataExtractor, Researcher]
        forM_ roles <| \role ->
          case Aeson.decode (Aeson.encode role) of
            Nothing -> Test.assertFailure ("Failed to decode role: " <> show role)
            Just decoded -> decoded Test.@=? role,
      Test.unit "SubagentConfig JSON roundtrip" <| do
        let cfg = defaultSubagentConfig WebCrawler "test task"
        case Aeson.decode (Aeson.encode cfg) of
          Nothing -> Test.assertFailure "Failed to decode SubagentConfig"
          Just decoded -> subagentTask decoded Test.@=? "test task",
      Test.unit "SubagentResult JSON roundtrip" <| do
        let result =
              SubagentResult
                { subagentOutput = Aeson.object ["data" .= ("test" :: Text)],
                  subagentSummary = "Test summary",
                  subagentConfidence = 0.85,
                  subagentTokensUsed = 1000,
                  subagentCostCents = 0.5,
                  subagentDuration = 30,
                  subagentIterations = 3,
                  subagentStatus = SubagentSuccess
                }
        case Aeson.decode (Aeson.encode result) of
          Nothing -> Test.assertFailure "Failed to decode SubagentResult"
          Just decoded -> subagentSummary decoded Test.@=? "Test summary",
      Test.unit "SubagentStatus JSON roundtrip" <| do
        let statuses =
              [ SubagentSuccess,
                SubagentTimeout,
                SubagentCostExceeded,
                SubagentError "test error"
              ]
        forM_ statuses <| \status ->
          case Aeson.decode (Aeson.encode status) of
            Nothing -> Test.assertFailure ("Failed to decode status: " <> show status)
            Just decoded -> decoded Test.@=? status,
      Test.unit "toolsForRole WebCrawler has web tools" <| do
        let keys = SubagentApiKeys "test-openrouter-key" (Just "test-kagi-key")
        let tools = toolsForRole WebCrawler keys
        let names = map Engine.toolName tools
        ("web_search" `elem` names) Test.@=? True
        ("read_webpages" `elem` names) Test.@=? True,
      Test.unit "toolsForRole CodeReviewer has code tools" <| do
        let keys = SubagentApiKeys "test-openrouter-key" Nothing
        let tools = toolsForRole CodeReviewer keys
        let names = map Engine.toolName tools
        ("read_file" `elem` names) Test.@=? True
        ("search_codebase" `elem` names) Test.@=? True,
      Test.unit "modelForRole returns appropriate models" <| do
        modelForRole WebCrawler Test.@=? "anthropic/claude-3-haiku"
        modelForRole CodeReviewer Test.@=? "anthropic/claude-sonnet-4"
        modelForRole Researcher Test.@=? "anthropic/claude-sonnet-4",
      Test.unit "defaultSubagentConfig has sensible defaults" <| do
        let cfg = defaultSubagentConfig WebCrawler "task"
        subagentTimeout cfg Test.@=? 600
        subagentMaxCost cfg Test.@=? 100.0
        subagentMaxTokens cfg Test.@=? 200000
        subagentMaxIterations cfg Test.@=? 20,
      Test.unit "spawnSubagentTool has correct name" <| do
        let keys = SubagentApiKeys "test-openrouter-key" (Just "test-kagi-key")
        let tool = spawnSubagentTool keys
        Engine.toolName tool Test.@=? "spawn_subagent"
    ]

data SubagentRole
  = WebCrawler
  | CodeReviewer
  | DataExtractor
  | Researcher
  | CustomRole Text
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentRole where
  toJSON WebCrawler = Aeson.String "web_crawler"
  toJSON CodeReviewer = Aeson.String "code_reviewer"
  toJSON DataExtractor = Aeson.String "data_extractor"
  toJSON Researcher = Aeson.String "researcher"
  toJSON (CustomRole name) = Aeson.String name

instance Aeson.FromJSON SubagentRole where
  parseJSON = Aeson.withText "SubagentRole" parseRole
    where
      parseRole "web_crawler" = pure WebCrawler
      parseRole "code_reviewer" = pure CodeReviewer
      parseRole "data_extractor" = pure DataExtractor
      parseRole "researcher" = pure Researcher
      parseRole name = pure (CustomRole name)

data SubagentConfig = SubagentConfig
  { subagentRole :: SubagentRole,
    subagentTask :: Text,
    subagentModel :: Maybe Text,
    subagentTimeout :: Int,
    subagentMaxCost :: Double,
    subagentMaxTokens :: Int,
    subagentMaxIterations :: Int,
    subagentExtendedThinking :: Bool,
    subagentContext :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentConfig where
  toJSON c =
    Aeson.object
      <| catMaybes
        [ Just ("role" .= subagentRole c),
          Just ("task" .= subagentTask c),
          ("model" .=) </ subagentModel c,
          Just ("timeout" .= subagentTimeout c),
          Just ("max_cost_cents" .= subagentMaxCost c),
          Just ("max_tokens" .= subagentMaxTokens c),
          Just ("max_iterations" .= subagentMaxIterations c),
          Just ("extended_thinking" .= subagentExtendedThinking c),
          ("context" .=) </ subagentContext c
        ]

instance Aeson.FromJSON SubagentConfig where
  parseJSON =
    Aeson.withObject "SubagentConfig" <| \v ->
      (SubagentConfig </ (v .: "role"))
        <*> (v .: "task")
        <*> (v .:? "model")
        <*> (v .:? "timeout" .!= 600)
        <*> (v .:? "max_cost_cents" .!= 50.0)
        <*> (v .:? "max_tokens" .!= 100000)
        <*> (v .:? "max_iterations" .!= 20)
        <*> (v .:? "extended_thinking" .!= False)
        <*> (v .:? "context")

data SubagentResult = SubagentResult
  { subagentOutput :: Aeson.Value,
    subagentSummary :: Text,
    subagentConfidence :: Double,
    subagentTokensUsed :: Int,
    subagentCostCents :: Double,
    subagentDuration :: Int,
    subagentIterations :: Int,
    subagentStatus :: SubagentStatus
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentResult

instance Aeson.FromJSON SubagentResult

data SubagentStatus
  = SubagentSuccess
  | SubagentTimeout
  | SubagentCostExceeded
  | SubagentError Text
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentStatus where
  toJSON SubagentSuccess = Aeson.String "success"
  toJSON SubagentTimeout = Aeson.String "timeout"
  toJSON SubagentCostExceeded = Aeson.String "cost_exceeded"
  toJSON (SubagentError msg) = Aeson.object ["error" .= msg]

instance Aeson.FromJSON SubagentStatus where
  parseJSON (Aeson.String "success") = pure SubagentSuccess
  parseJSON (Aeson.String "timeout") = pure SubagentTimeout
  parseJSON (Aeson.String "cost_exceeded") = pure SubagentCostExceeded
  parseJSON (Aeson.Object v) = SubagentError </ (v .: "error")
  parseJSON _ = empty

data SubagentCallbacks = SubagentCallbacks
  { onSubagentStart :: Text -> IO (),
    onSubagentActivity :: Text -> IO (),
    onSubagentToolCall :: Text -> Text -> IO (),
    onSubagentComplete :: SubagentResult -> IO ()
  }

defaultCallbacks :: SubagentCallbacks
defaultCallbacks =
  SubagentCallbacks
    { onSubagentStart = \_ -> pure (),
      onSubagentActivity = \_ -> pure (),
      onSubagentToolCall = \_ _ -> pure (),
      onSubagentComplete = \_ -> pure ()
    }

defaultSubagentConfig :: SubagentRole -> Text -> SubagentConfig
defaultSubagentConfig role task =
  SubagentConfig
    { subagentRole = role,
      subagentTask = task,
      subagentModel = Nothing,
      subagentTimeout = 600,
      subagentMaxCost = 100.0,
      subagentMaxTokens = 200000,
      subagentMaxIterations = 20,
      subagentExtendedThinking = False,
      subagentContext = Nothing
    }

modelForRole :: SubagentRole -> Text
modelForRole WebCrawler = "anthropic/claude-3-haiku"
modelForRole CodeReviewer = "anthropic/claude-sonnet-4"
modelForRole DataExtractor = "anthropic/claude-3-haiku"
modelForRole Researcher = "anthropic/claude-sonnet-4"
modelForRole (CustomRole _) = "anthropic/claude-sonnet-4"

data SubagentApiKeys = SubagentApiKeys
  { subagentOpenRouterKey :: Text,
    subagentKagiKey :: Maybe Text
  }
  deriving (Show, Eq)

toolsForRole :: SubagentRole -> SubagentApiKeys -> [Engine.Tool]
toolsForRole WebCrawler keys =
  let webSearchTools = case subagentKagiKey keys of
        Just kagiKey -> [WebSearch.webSearchTool kagiKey]
        Nothing -> []
   in webSearchTools
        <> [ WebReader.webReaderTool (subagentOpenRouterKey keys),
             Tools.searchCodebaseTool
           ]
toolsForRole CodeReviewer _keys =
  [ Tools.readFileTool,
    Tools.searchCodebaseTool,
    Tools.searchAndReadTool,
    Tools.runBashTool
  ]
toolsForRole DataExtractor keys =
  [ WebReader.webReaderTool (subagentOpenRouterKey keys),
    Tools.readFileTool,
    Tools.searchCodebaseTool
  ]
toolsForRole Researcher keys =
  let webSearchTools = case subagentKagiKey keys of
        Just kagiKey -> [WebSearch.webSearchTool kagiKey]
        Nothing -> []
   in webSearchTools
        <> [ WebReader.webReaderTool (subagentOpenRouterKey keys),
             Tools.readFileTool,
             Tools.searchCodebaseTool,
             Tools.searchAndReadTool
           ]
toolsForRole (CustomRole _) keys = toolsForRole Researcher keys

systemPromptForRole :: SubagentRole -> Text -> Maybe Text -> Text
systemPromptForRole role task maybeContext =
  Text.unlines
    [ "You are a specialized " <> roleDescription role <> " subagent working on a focused task.",
      "",
      "## Your Task",
      task,
      "",
      maybe "" (\ctx -> "## Context from Orchestrator\n" <> ctx <> "\n") maybeContext,
      "## Guidelines",
      "1. Be EFFICIENT with context - extract only key facts, don't save full page contents",
      "2. Summarize findings as you go rather than accumulating raw data",
      "3. Limit web page reads to 3-5 most relevant sources",
      "4. Work iteratively: search → skim results → read best 2-3 → synthesize",
      "5. ALWAYS cite sources - every claim needs a URL",
      "6. Stop when you have sufficient information - don't over-research",
      "",
      "## Output Format",
      "Return findings as a list of structured insights:",
      "",
      "```json",
      "{",
      "  \"summary\": \"Brief overall summary (1-2 sentences)\",",
      "  \"confidence\": 0.85,",
      "  \"findings\": [",
      "    {",
      "      \"claim\": \"The key insight or fact discovered\",",
      "      \"source_url\": \"https://example.com/page\",",
      "      \"quote\": \"Relevant excerpt supporting the claim\",",
      "      \"source_name\": \"Example Site\"",
      "    }",
      "  ],",
      "  \"caveats\": \"Any limitations or uncertainties\"",
      "}",
      "```"
    ]
  where
    roleDescription :: SubagentRole -> Text
    roleDescription WebCrawler = "web research"
    roleDescription CodeReviewer = "code review"
    roleDescription DataExtractor = "data extraction"
    roleDescription Researcher = "research"
    roleDescription (CustomRole name) = name

runSubagent :: SubagentApiKeys -> SubagentConfig -> IO SubagentResult
runSubagent keys config = runSubagentWithCallbacks keys config defaultCallbacks

runSubagentWithCallbacks :: SubagentApiKeys -> SubagentConfig -> SubagentCallbacks -> IO SubagentResult
runSubagentWithCallbacks keys config callbacks = do
  startTime <- Clock.getCurrentTime

  let role = subagentRole config
  let model = fromMaybe (modelForRole role) (subagentModel config)
  let tools = toolsForRole role keys
  let systemPrompt = systemPromptForRole role (subagentTask config) (subagentContext config)

  onSubagentStart callbacks ("Starting " <> tshow role <> " subagent...")

  let provider = Provider.defaultOpenRouter (subagentOpenRouterKey keys) model

  let guardrails =
        Engine.Guardrails
          { Engine.guardrailMaxCostCents = subagentMaxCost config,
            Engine.guardrailMaxTokens = subagentMaxTokens config,
            Engine.guardrailMaxDuplicateToolCalls = 20,
            Engine.guardrailMaxTestFailures = 3,
            Engine.guardrailMaxEditFailures = 5
          }

  let agentConfig =
        Engine.AgentConfig
          { Engine.agentModel = model,
            Engine.agentTools = tools,
            Engine.agentSystemPrompt = systemPrompt,
            Engine.agentMaxIterations = subagentMaxIterations config,
            Engine.agentGuardrails = guardrails
          }

  let engineConfig =
        Engine.EngineConfig
          { Engine.engineLLM = Engine.defaultLLM,
            Engine.engineOnCost = \_ _ -> pure (),
            Engine.engineOnActivity = onSubagentActivity callbacks,
            Engine.engineOnToolCall = onSubagentToolCall callbacks,
            Engine.engineOnAssistant = \_ -> pure (),
            Engine.engineOnToolResult = \_ _ _ -> pure (),
            Engine.engineOnComplete = pure (),
            Engine.engineOnError = \_ -> pure (),
            Engine.engineOnGuardrail = \_ -> pure ()
          }

  let timeoutMicros = subagentTimeout config * 1000000

  resultOrTimeout <-
    race
      (threadDelay timeoutMicros)
      (Engine.runAgentWithProvider engineConfig provider agentConfig (subagentTask config))

  endTime <- Clock.getCurrentTime
  let durationSecs = round (Clock.diffUTCTime endTime startTime)

  let result = case resultOrTimeout of
        Left () ->
          SubagentResult
            { subagentOutput = Aeson.object ["error" .= ("Timeout after " <> tshow (subagentTimeout config) <> " seconds" :: Text)],
              subagentSummary = "Subagent timed out",
              subagentConfidence = 0.0,
              subagentTokensUsed = 0,
              subagentCostCents = 0.0,
              subagentDuration = durationSecs,
              subagentIterations = 0,
              subagentStatus = SubagentTimeout
            }
        Right (Left err) ->
          let status = if "cost" `Text.isInfixOf` Text.toLower err then SubagentCostExceeded else SubagentError err
           in SubagentResult
                { subagentOutput = Aeson.object ["error" .= err],
                  subagentSummary = "Subagent failed: " <> err,
                  subagentConfidence = 0.0,
                  subagentTokensUsed = 0,
                  subagentCostCents = 0.0,
                  subagentDuration = durationSecs,
                  subagentIterations = 0,
                  subagentStatus = status
                }
        Right (Right agentResult) ->
          SubagentResult
            { subagentOutput = Aeson.object ["response" .= Engine.resultFinalMessage agentResult],
              subagentSummary = truncateSummary (Engine.resultFinalMessage agentResult),
              subagentConfidence = 0.8,
              subagentTokensUsed = Engine.resultTotalTokens agentResult,
              subagentCostCents = Engine.resultTotalCost agentResult,
              subagentDuration = durationSecs,
              subagentIterations = Engine.resultIterations agentResult,
              subagentStatus = SubagentSuccess
            }

  onSubagentComplete callbacks result
  pure result
  where
    truncateSummary :: Text -> Text
    truncateSummary txt =
      let firstLine = Text.takeWhile (/= '\n') txt
       in if Text.length firstLine > 200
            then Text.take 197 firstLine <> "..."
            else firstLine

spawnSubagentTool :: SubagentApiKeys -> Engine.Tool
spawnSubagentTool keys =
  Engine.Tool
    { Engine.toolName = "spawn_subagent",
      Engine.toolDescription =
        "Spawn a specialized subagent for a focused task. "
          <> "Use for tasks that benefit from deep exploration, parallel execution, "
          <> "or specialized tools. The subagent will iterate until task completion "
          <> "or resource limits are reached. "
          <> "Available roles: web_crawler (fast web research), code_reviewer (thorough code analysis), "
          <> "data_extractor (structured data extraction), researcher (general research).",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "role"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "enum" .= (["web_crawler", "code_reviewer", "data_extractor", "researcher"] :: [Text]),
                        "description" .= ("Subagent role determining tools and model" :: Text)
                      ],
                  "task"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The specific task for the subagent to accomplish" :: Text)
                      ],
                  "context"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Additional context to help the subagent understand the goal" :: Text)
                      ],
                  "model"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Override the default model for this role" :: Text)
                      ],
                  "timeout"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Timeout in seconds (default: 600)" :: Text)
                      ],
                  "max_cost_cents"
                    .= Aeson.object
                      [ "type" .= ("number" :: Text),
                        "description" .= ("Maximum cost in cents (default: 50)" :: Text)
                      ]
                ],
            "required" .= (["role", "task"] :: [Text])
          ],
      Engine.toolExecute = executeSpawnSubagent keys
    }

executeSpawnSubagent :: SubagentApiKeys -> Aeson.Value -> IO Aeson.Value
executeSpawnSubagent keys v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure <| Aeson.object ["error" .= ("Invalid arguments: " <> Text.pack e)]
    Aeson.Success config -> do
      result <- runSubagent keys config
      pure (Aeson.toJSON result)