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

-- | LLM Agent Engine - Tool protocol and LLM provider abstraction.
--
-- This module provides the core abstractions for building LLM-powered agents:
-- - Tool: Defines tools that agents can use
-- - LLM: OpenAI-compatible chat completions API provider
-- - AgentConfig: Configuration for running agents
--
-- : out omni-agent-engine
-- : dep http-conduit
-- : dep aeson
-- : dep case-insensitive
module Omni.Agent.Engine
  ( Tool (..),
    LLM (..),
    EngineConfig (..),
    AgentConfig (..),
    AgentResult (..),
    Message (..),
    Role (..),
    ToolCall (..),
    FunctionCall (..),
    ToolResult (..),
    ChatCompletionRequest (..),
    ChatCompletionResponse (..),
    Choice (..),
    Usage (..),
    defaultLLM,
    defaultEngineConfig,
    defaultAgentConfig,
    chat,
    runAgent,
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy as BL
import qualified Data.CaseInsensitive as CI
import qualified Data.Map.Strict as Map
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Test as Test

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Engine"
    [ Test.unit "Tool JSON roundtrip" <| do
        let tool =
              Tool
                { toolName = "get_weather",
                  toolDescription = "Get weather for a location",
                  toolJsonSchema = Aeson.object ["type" .= ("object" :: Text), "properties" .= Aeson.object []],
                  toolExecute = \_ -> pure (Aeson.String "sunny")
                }
        let encoded = encodeToolForApi tool
        case Aeson.decode (Aeson.encode encoded) of
          Nothing -> Test.assertFailure "Failed to decode tool"
          Just decoded -> toolName tool Test.@=? toolApiName decoded,
      Test.unit "Message JSON roundtrip" <| do
        let msg = Message User "Hello" Nothing Nothing
        case Aeson.decode (Aeson.encode msg) of
          Nothing -> Test.assertFailure "Failed to decode message"
          Just decoded -> msgContent msg Test.@=? msgContent decoded,
      Test.unit "defaultLLM has correct endpoint" <| do
        llmBaseUrl defaultLLM Test.@=? "https://openrouter.ai/api/v1",
      Test.unit "defaultLLM has OpenRouter headers" <| do
        length (llmExtraHeaders defaultLLM) Test.@=? 2
        llmModel defaultLLM Test.@=? "anthropic/claude-sonnet-4.5",
      Test.unit "defaultAgentConfig has sensible defaults" <| do
        agentMaxIterations defaultAgentConfig Test.@=? 10,
      Test.unit "defaultEngineConfig has no-op callbacks" <| do
        engineOnCost defaultEngineConfig 100 5
        engineOnActivity defaultEngineConfig "test"
        engineOnToolCall defaultEngineConfig "tool" "result"
        True Test.@=? True,
      Test.unit "buildToolMap creates correct map" <| do
        let tool1 =
              Tool
                { toolName = "tool1",
                  toolDescription = "First tool",
                  toolJsonSchema = Aeson.object [],
                  toolExecute = \_ -> pure Aeson.Null
                }
            tool2 =
              Tool
                { toolName = "tool2",
                  toolDescription = "Second tool",
                  toolJsonSchema = Aeson.object [],
                  toolExecute = \_ -> pure Aeson.Null
                }
            toolMap = buildToolMap [tool1, tool2]
        Map.size toolMap Test.@=? 2
        Map.member "tool1" toolMap Test.@=? True
        Map.member "tool2" toolMap Test.@=? True,
      Test.unit "Usage JSON parsing" <| do
        let json = "{\"prompt_tokens\":100,\"completion_tokens\":50,\"total_tokens\":150}"
        case Aeson.decode json of
          Nothing -> Test.assertFailure "Failed to decode usage"
          Just usage -> do
            usagePromptTokens usage Test.@=? 100
            usageCompletionTokens usage Test.@=? 50
            usageTotalTokens usage Test.@=? 150,
      Test.unit "AgentResult JSON roundtrip" <| do
        let result =
              AgentResult
                { resultFinalMessage = "Done",
                  resultToolCallCount = 3,
                  resultIterations = 2,
                  resultTotalCost = 50,
                  resultTotalTokens = 1500
                }
        case Aeson.decode (Aeson.encode result) of
          Nothing -> Test.assertFailure "Failed to decode AgentResult"
          Just decoded -> do
            resultFinalMessage decoded Test.@=? "Done"
            resultToolCallCount decoded Test.@=? 3
            resultIterations decoded Test.@=? 2,
      Test.unit "estimateCost calculates correctly" <| do
        let gpt4oCost = estimateCost "gpt-4o" 1000
            gpt4oMiniCost = estimateCost "gpt-4o-mini" 1000
        (gpt4oCost >= gpt4oMiniCost) Test.@=? True,
      Test.unit "ToolCall JSON roundtrip" <| do
        let tc =
              ToolCall
                { tcId = "call_123",
                  tcType = "function",
                  tcFunction = FunctionCall "read_file" "{\"path\":\"/tmp/test\"}"
                }
        case Aeson.decode (Aeson.encode tc) of
          Nothing -> Test.assertFailure "Failed to decode ToolCall"
          Just decoded -> tcId decoded Test.@=? "call_123",
      Test.unit "FunctionCall JSON roundtrip" <| do
        let fc = FunctionCall "test_func" "{\"arg\":\"value\"}"
        case Aeson.decode (Aeson.encode fc) of
          Nothing -> Test.assertFailure "Failed to decode FunctionCall"
          Just decoded -> do
            fcName decoded Test.@=? "test_func"
            fcArguments decoded Test.@=? "{\"arg\":\"value\"}",
      Test.unit "Role JSON roundtrip for all roles" <| do
        let roles = [System, User, Assistant, ToolRole]
        forM_ roles <| \role ->
          case Aeson.decode (Aeson.encode role) of
            Nothing -> Test.assertFailure ("Failed to decode Role: " <> show role)
            Just decoded -> decoded Test.@=? role
    ]

data Tool = Tool
  { toolName :: Text,
    toolDescription :: Text,
    toolJsonSchema :: Aeson.Value,
    toolExecute :: Aeson.Value -> IO Aeson.Value
  }

data ToolApi = ToolApi
  { toolApiName :: Text,
    toolApiDescription :: Text,
    toolApiParameters :: Aeson.Value
  }
  deriving (Generic)

instance Aeson.ToJSON ToolApi where
  toJSON t =
    Aeson.object
      [ "type" .= ("function" :: Text),
        "function"
          .= Aeson.object
            [ "name" .= toolApiName t,
              "description" .= toolApiDescription t,
              "parameters" .= toolApiParameters t
            ]
      ]

instance Aeson.FromJSON ToolApi where
  parseJSON =
    Aeson.withObject "ToolApi" <| \v -> do
      fn <- v .: "function"
      (ToolApi </ (fn .: "name"))
        <*> (fn .: "description")
        <*> (fn .: "parameters")

encodeToolForApi :: Tool -> ToolApi
encodeToolForApi t =
  ToolApi
    { toolApiName = toolName t,
      toolApiDescription = toolDescription t,
      toolApiParameters = toolJsonSchema t
    }

data LLM = LLM
  { llmBaseUrl :: Text,
    llmApiKey :: Text,
    llmModel :: Text,
    llmExtraHeaders :: [(ByteString, ByteString)]
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON LLM where
  toJSON l =
    Aeson.object
      [ "llmBaseUrl" .= llmBaseUrl l,
        "llmApiKey" .= llmApiKey l,
        "llmModel" .= llmModel l
      ]

instance Aeson.FromJSON LLM where
  parseJSON =
    Aeson.withObject "LLM" <| \v ->
      (LLM </ (v .: "llmBaseUrl"))
        <*> (v .: "llmApiKey")
        <*> (v .: "llmModel")
        <*> pure []

defaultLLM :: LLM
defaultLLM =
  LLM
    { llmBaseUrl = "https://openrouter.ai/api/v1",
      llmApiKey = "",
      llmModel = "anthropic/claude-sonnet-4.5",
      llmExtraHeaders =
        [ ("HTTP-Referer", "https://omni.dev"),
          ("X-Title", "Omni Agent")
        ]
    }

data AgentConfig = AgentConfig
  { agentModel :: Text,
    agentTools :: [Tool],
    agentSystemPrompt :: Text,
    agentMaxIterations :: Int
  }

defaultAgentConfig :: AgentConfig
defaultAgentConfig =
  AgentConfig
    { agentModel = "gpt-4",
      agentTools = [],
      agentSystemPrompt = "You are a helpful assistant.",
      agentMaxIterations = 10
    }

data EngineConfig = EngineConfig
  { engineLLM :: LLM,
    engineOnCost :: Int -> Int -> IO (),
    engineOnActivity :: Text -> IO (),
    engineOnToolCall :: Text -> Text -> IO ()
  }

defaultEngineConfig :: EngineConfig
defaultEngineConfig =
  EngineConfig
    { engineLLM = defaultLLM,
      engineOnCost = \_ _ -> pure (),
      engineOnActivity = \_ -> pure (),
      engineOnToolCall = \_ _ -> pure ()
    }

data AgentResult = AgentResult
  { resultFinalMessage :: Text,
    resultToolCallCount :: Int,
    resultIterations :: Int,
    resultTotalCost :: Int,
    resultTotalTokens :: Int
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON AgentResult

instance Aeson.FromJSON AgentResult

data Role = System | User | Assistant | ToolRole
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Role where
  toJSON System = Aeson.String "system"
  toJSON User = Aeson.String "user"
  toJSON Assistant = Aeson.String "assistant"
  toJSON ToolRole = Aeson.String "tool"

instance Aeson.FromJSON Role where
  parseJSON = Aeson.withText "Role" parseRole
    where
      parseRole "system" = pure System
      parseRole "user" = pure User
      parseRole "assistant" = pure Assistant
      parseRole "tool" = pure ToolRole
      parseRole _ = empty

data Message = Message
  { msgRole :: Role,
    msgContent :: Text,
    msgToolCalls :: Maybe [ToolCall],
    msgToolCallId :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Message where
  toJSON m =
    Aeson.object
      <| catMaybes
        [ Just ("role" .= msgRole m),
          Just ("content" .= msgContent m),
          ("tool_calls" .=) </ msgToolCalls m,
          ("tool_call_id" .=) </ msgToolCallId m
        ]

instance Aeson.FromJSON Message where
  parseJSON =
    Aeson.withObject "Message" <| \v ->
      (Message </ (v .: "role"))
        <*> (v .:? "content" .!= "")
        <*> (v .:? "tool_calls")
        <*> (v .:? "tool_call_id")

data ToolCall = ToolCall
  { tcId :: Text,
    tcType :: Text,
    tcFunction :: FunctionCall
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ToolCall where
  toJSON tc =
    Aeson.object
      [ "id" .= tcId tc,
        "type" .= tcType tc,
        "function" .= tcFunction tc
      ]

instance Aeson.FromJSON ToolCall where
  parseJSON =
    Aeson.withObject "ToolCall" <| \v ->
      (ToolCall </ (v .: "id"))
        <*> (v .:? "type" .!= "function")
        <*> (v .: "function")

data FunctionCall = FunctionCall
  { fcName :: Text,
    fcArguments :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON FunctionCall where
  toJSON fc =
    Aeson.object
      [ "name" .= fcName fc,
        "arguments" .= fcArguments fc
      ]

instance Aeson.FromJSON FunctionCall where
  parseJSON =
    Aeson.withObject "FunctionCall" <| \v ->
      (FunctionCall </ (v .: "name"))
        <*> (v .: "arguments")

data ToolResult = ToolResult
  { trToolCallId :: Text,
    trContent :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ToolResult

instance Aeson.FromJSON ToolResult

data ChatCompletionRequest = ChatCompletionRequest
  { reqModel :: Text,
    reqMessages :: [Message],
    reqTools :: Maybe [ToolApi]
  }
  deriving (Generic)

instance Aeson.ToJSON ChatCompletionRequest where
  toJSON r =
    Aeson.object
      <| catMaybes
        [ Just ("model" .= reqModel r),
          Just ("messages" .= reqMessages r),
          ("tools" .=) </ reqTools r
        ]

data Choice = Choice
  { choiceIndex :: Int,
    choiceMessage :: Message,
    choiceFinishReason :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON Choice where
  parseJSON =
    Aeson.withObject "Choice" <| \v ->
      (Choice </ (v .: "index"))
        <*> (v .: "message")
        <*> (v .:? "finish_reason")

data Usage = Usage
  { usagePromptTokens :: Int,
    usageCompletionTokens :: Int,
    usageTotalTokens :: Int
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON Usage where
  parseJSON =
    Aeson.withObject "Usage" <| \v ->
      (Usage </ (v .: "prompt_tokens"))
        <*> (v .: "completion_tokens")
        <*> (v .: "total_tokens")

data ChatCompletionResponse = ChatCompletionResponse
  { respId :: Text,
    respChoices :: [Choice],
    respModel :: Text,
    respUsage :: Maybe Usage
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON ChatCompletionResponse where
  parseJSON =
    Aeson.withObject "ChatCompletionResponse" <| \v ->
      (ChatCompletionResponse </ (v .: "id"))
        <*> (v .: "choices")
        <*> (v .: "model")
        <*> (v .:? "usage")

data ChatResult = ChatResult
  { chatMessage :: Message,
    chatUsage :: Maybe Usage
  }
  deriving (Show, Eq)

chatWithUsage :: LLM -> [Tool] -> [Message] -> IO (Either Text ChatResult)
chatWithUsage llm tools messages = do
  let url = Text.unpack (llmBaseUrl llm) <> "/chat/completions"
  req0 <- HTTP.parseRequest url
  let toolApis = [encodeToolForApi t | not (null tools), t <- tools]
      body =
        ChatCompletionRequest
          { reqModel = llmModel llm,
            reqMessages = messages,
            reqTools = if null toolApis then Nothing else Just toolApis
          }
      baseReq =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestHeader "Authorization" ["Bearer " <> TE.encodeUtf8 (llmApiKey llm)]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
      req = foldr addHeader baseReq (llmExtraHeaders llm)
      addHeader (name, value) = HTTP.addRequestHeader (CI.mk name) value

  response <- HTTP.httpLBS req
  let status = HTTP.getResponseStatusCode response
  if status >= 200 && status < 300
    then case Aeson.decode (HTTP.getResponseBody response) of
      Just resp ->
        case respChoices resp of
          (c : _) -> pure (Right (ChatResult (choiceMessage c) (respUsage resp)))
          [] -> pure (Left "No choices in response")
      Nothing -> pure (Left "Failed to parse response")
    else pure (Left ("HTTP error: " <> tshow status <> " - " <> TE.decodeUtf8 (BL.toStrict (HTTP.getResponseBody response))))

chat :: LLM -> [Tool] -> [Message] -> IO (Either Text Message)
chat llm tools messages = do
  result <- chatWithUsage llm tools messages
  pure (chatMessage </ result)

runAgent :: EngineConfig -> AgentConfig -> Text -> IO (Either Text AgentResult)
runAgent engineCfg agentCfg userPrompt = do
  let llm =
        (engineLLM engineCfg)
          { llmModel = agentModel agentCfg
          }
      tools = agentTools agentCfg
      toolMap = buildToolMap tools
      systemMsg = Message System (agentSystemPrompt agentCfg) Nothing Nothing
      userMsg = Message User userPrompt Nothing Nothing
      initialMessages = [systemMsg, userMsg]

  engineOnActivity engineCfg "Starting agent loop"
  loop llm tools toolMap initialMessages 0 0 0
  where
    maxIter = agentMaxIterations agentCfg

    loop :: LLM -> [Tool] -> Map.Map Text Tool -> [Message] -> Int -> Int -> Int -> IO (Either Text AgentResult)
    loop llm tools' toolMap msgs iteration totalCalls totalTokens
      | iteration >= maxIter =
          pure
            <| Left
            <| "Max iterations ("
            <> tshow maxIter
            <> ") reached"
      | otherwise = do
          engineOnActivity engineCfg <| "Iteration " <> tshow (iteration + 1)
          result <- chatWithUsage llm tools' msgs
          case result of
            Left err -> pure (Left err)
            Right chatRes -> do
              let msg = chatMessage chatRes
                  tokens = maybe 0 usageTotalTokens (chatUsage chatRes)
                  cost = estimateCost (llmModel llm) tokens
              engineOnCost engineCfg tokens cost
              let newTokens = totalTokens + tokens
              case msgToolCalls msg of
                Nothing -> do
                  engineOnActivity engineCfg "Agent completed"
                  pure
                    <| Right
                    <| AgentResult
                      { resultFinalMessage = msgContent msg,
                        resultToolCallCount = totalCalls,
                        resultIterations = iteration + 1,
                        resultTotalCost = estimateTotalCost (llmModel llm) newTokens,
                        resultTotalTokens = newTokens
                      }
                Just [] -> do
                  engineOnActivity engineCfg "Agent completed (empty tool calls)"
                  pure
                    <| Right
                    <| AgentResult
                      { resultFinalMessage = msgContent msg,
                        resultToolCallCount = totalCalls,
                        resultIterations = iteration + 1,
                        resultTotalCost = estimateTotalCost (llmModel llm) newTokens,
                        resultTotalTokens = newTokens
                      }
                Just tcs -> do
                  toolResults <- executeToolCalls engineCfg toolMap tcs
                  let newMsgs = msgs <> [msg] <> toolResults
                      newCalls = totalCalls + length tcs
                  loop llm tools' toolMap newMsgs (iteration + 1) newCalls newTokens

buildToolMap :: [Tool] -> Map.Map Text Tool
buildToolMap = Map.fromList <. map (\t -> (toolName t, t))

executeToolCalls :: EngineConfig -> Map.Map Text Tool -> [ToolCall] -> IO [Message]
executeToolCalls engineCfg toolMap = traverse executeSingle
  where
    executeSingle tc = do
      let name = fcName (tcFunction tc)
          argsText = fcArguments (tcFunction tc)
          callId = tcId tc
      engineOnActivity engineCfg <| "Executing tool: " <> name
      case Map.lookup name toolMap of
        Nothing -> do
          let errMsg = "Tool not found: " <> name
          engineOnToolCall engineCfg name errMsg
          pure <| Message ToolRole errMsg Nothing (Just callId)
        Just tool -> do
          case Aeson.decode (BL.fromStrict (TE.encodeUtf8 argsText)) of
            Nothing -> do
              let errMsg = "Invalid JSON arguments: " <> argsText
              engineOnToolCall engineCfg name errMsg
              pure <| Message ToolRole errMsg Nothing (Just callId)
            Just args -> do
              resultValue <- toolExecute tool args
              let resultText = TE.decodeUtf8 (BL.toStrict (Aeson.encode resultValue))
                  summary = Text.take 100 resultText
              engineOnToolCall engineCfg name summary
              pure <| Message ToolRole resultText Nothing (Just callId)

estimateCost :: Text -> Int -> Int
estimateCost model tokens
  | "gpt-4o-mini" `Text.isInfixOf` model = tokens * 15 `div` 1000000
  | "gpt-4o" `Text.isInfixOf` model = tokens * 250 `div` 100000
  | "gpt-4" `Text.isInfixOf` model = tokens * 3 `div` 100000
  | "claude" `Text.isInfixOf` model = tokens * 3 `div` 100000
  | otherwise = tokens `div` 100000

estimateTotalCost :: Text -> Int -> Int
estimateTotalCost = estimateCost