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

-- | HTTP request tools for agent API interactions.
--
-- Provides http_get and http_post tools for making HTTP requests.
-- Supports headers, query params, and JSON body.
--
-- : out omni-agent-tools-http
-- : dep aeson
-- : dep http-conduit
module Omni.Agent.Tools.Http
  ( -- * Tools
    httpGetTool,
    httpPostTool,
    allHttpTools,

    -- * Types
    HttpGetArgs (..),
    HttpPostArgs (..),
    HttpResult (..),

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Key as Key
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString.Lazy as BL
import qualified Data.CaseInsensitive as CI
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Network.HTTP.Client as HTTPClient
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Test as Test
import System.Timeout (timeout)

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Tools.Http"
    [ Test.unit "httpGetTool has correct name" <| do
        Engine.toolName httpGetTool Test.@=? "http_get",
      Test.unit "httpPostTool has correct name" <| do
        Engine.toolName httpPostTool Test.@=? "http_post",
      Test.unit "allHttpTools has 2 tools" <| do
        length allHttpTools Test.@=? 2,
      Test.unit "HttpGetArgs parses correctly" <| do
        let json = Aeson.object ["url" .= ("https://example.com" :: Text)]
        case Aeson.fromJSON json of
          Aeson.Success (args :: HttpGetArgs) -> httpGetUrl args Test.@=? "https://example.com"
          Aeson.Error e -> Test.assertFailure e,
      Test.unit "HttpGetArgs parses with headers" <| do
        let json =
              Aeson.object
                [ "url" .= ("https://api.example.com" :: Text),
                  "headers" .= Aeson.object ["Authorization" .= ("Bearer token" :: Text)]
                ]
        case Aeson.fromJSON json of
          Aeson.Success (args :: HttpGetArgs) -> do
            httpGetUrl args Test.@=? "https://api.example.com"
            isJust (httpGetHeaders args) Test.@=? True
          Aeson.Error e -> Test.assertFailure e,
      Test.unit "HttpPostArgs parses correctly" <| do
        let json =
              Aeson.object
                [ "url" .= ("https://api.example.com" :: Text),
                  "body" .= Aeson.object ["key" .= ("value" :: Text)]
                ]
        case Aeson.fromJSON json of
          Aeson.Success (args :: HttpPostArgs) -> do
            httpPostUrl args Test.@=? "https://api.example.com"
            isJust (httpPostBody args) Test.@=? True
          Aeson.Error e -> Test.assertFailure e,
      Test.unit "HttpResult JSON roundtrip" <| do
        let result =
              HttpResult
                { httpResultStatus = 200,
                  httpResultHeaders = Aeson.object ["Content-Type" .= ("application/json" :: Text)],
                  httpResultBody = "{\"ok\": true}"
                }
        case Aeson.decode (Aeson.encode result) of
          Nothing -> Test.assertFailure "Failed to decode HttpResult"
          Just decoded -> httpResultStatus decoded Test.@=? 200,
      Test.unit "http_get fetches real URL" <| do
        let args = Aeson.object ["url" .= ("https://httpbin.org/get" :: Text)]
        result <- Engine.toolExecute httpGetTool args
        case Aeson.fromJSON result of
          Aeson.Success (r :: HttpResult) -> do
            httpResultStatus r Test.@=? 200
            ("httpbin.org" `Text.isInfixOf` httpResultBody r) Test.@=? True
          Aeson.Error e -> Test.assertFailure e,
      Test.unit "http_post with JSON body" <| do
        let args =
              Aeson.object
                [ "url" .= ("https://httpbin.org/post" :: Text),
                  "body" .= Aeson.object ["test" .= ("value" :: Text)]
                ]
        result <- Engine.toolExecute httpPostTool args
        case Aeson.fromJSON result of
          Aeson.Success (r :: HttpResult) -> do
            httpResultStatus r Test.@=? 200
            ("test" `Text.isInfixOf` httpResultBody r) Test.@=? True
          Aeson.Error e -> Test.assertFailure e
    ]

data HttpGetArgs = HttpGetArgs
  { httpGetUrl :: Text,
    httpGetHeaders :: Maybe Aeson.Object,
    httpGetParams :: Maybe Aeson.Object
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON HttpGetArgs where
  parseJSON =
    Aeson.withObject "HttpGetArgs" <| \v ->
      (HttpGetArgs </ (v .: "url"))
        <*> (v .:? "headers")
        <*> (v .:? "params")

data HttpPostArgs = HttpPostArgs
  { httpPostUrl :: Text,
    httpPostHeaders :: Maybe Aeson.Object,
    httpPostBody :: Maybe Aeson.Value,
    httpPostContentType :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON HttpPostArgs where
  parseJSON =
    Aeson.withObject "HttpPostArgs" <| \v ->
      (HttpPostArgs </ (v .: "url"))
        <*> (v .:? "headers")
        <*> (v .:? "body")
        <*> (v .:? "content_type")

data HttpResult = HttpResult
  { httpResultStatus :: Int,
    httpResultHeaders :: Aeson.Value,
    httpResultBody :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON HttpResult where
  toJSON r =
    Aeson.object
      [ "status" .= httpResultStatus r,
        "headers" .= httpResultHeaders r,
        "body" .= httpResultBody r
      ]

instance Aeson.FromJSON HttpResult where
  parseJSON =
    Aeson.withObject "HttpResult" <| \v ->
      (HttpResult </ (v .: "status"))
        <*> (v .: "headers")
        <*> (v .: "body")

allHttpTools :: [Engine.Tool]
allHttpTools = [httpGetTool, httpPostTool]

httpGetTool :: Engine.Tool
httpGetTool =
  Engine.Tool
    { Engine.toolName = "http_get",
      Engine.toolDescription =
        "Make an HTTP GET request. Returns status code, headers, and response body. "
          <> "Use for fetching data from APIs or web pages.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "url"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The URL to request" :: Text)
                      ],
                  "headers"
                    .= Aeson.object
                      [ "type" .= ("object" :: Text),
                        "description" .= ("Optional headers as key-value pairs" :: Text)
                      ],
                  "params"
                    .= Aeson.object
                      [ "type" .= ("object" :: Text),
                        "description" .= ("Optional query parameters as key-value pairs" :: Text)
                      ]
                ],
            "required" .= (["url"] :: [Text])
          ],
      Engine.toolExecute = executeHttpGet
    }

executeHttpGet :: Aeson.Value -> IO Aeson.Value
executeHttpGet v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure <| mkError ("Invalid arguments: " <> Text.pack e)
    Aeson.Success args -> do
      let urlWithParams = case httpGetParams args of
            Nothing -> httpGetUrl args
            Just params ->
              let paramList = [(k, v') | (k, v') <- KeyMap.toList params]
                  paramStr = Text.intercalate "&" [Key.toText k <> "=" <> valueToText v' | (k, v') <- paramList]
               in if Text.null paramStr
                    then httpGetUrl args
                    else httpGetUrl args <> "?" <> paramStr
      doHttpRequest "GET" urlWithParams (httpGetHeaders args) Nothing

httpPostTool :: Engine.Tool
httpPostTool =
  Engine.Tool
    { Engine.toolName = "http_post",
      Engine.toolDescription =
        "Make an HTTP POST request. Returns status code, headers, and response body. "
          <> "Use for submitting data to APIs or forms.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "url"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The URL to request" :: Text)
                      ],
                  "headers"
                    .= Aeson.object
                      [ "type" .= ("object" :: Text),
                        "description" .= ("Optional headers as key-value pairs" :: Text)
                      ],
                  "body"
                    .= Aeson.object
                      [ "type" .= ("object" :: Text),
                        "description" .= ("Optional JSON body (object or string)" :: Text)
                      ],
                  "content_type"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Content type (default: application/json)" :: Text)
                      ]
                ],
            "required" .= (["url"] :: [Text])
          ],
      Engine.toolExecute = executeHttpPost
    }

executeHttpPost :: Aeson.Value -> IO Aeson.Value
executeHttpPost v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure <| mkError ("Invalid arguments: " <> Text.pack e)
    Aeson.Success args -> do
      let contentType = fromMaybe "application/json" (httpPostContentType args)
          body = case httpPostBody args of
            Nothing -> Nothing
            Just b -> Just (contentType, BL.toStrict (Aeson.encode b))
      doHttpRequest "POST" (httpPostUrl args) (httpPostHeaders args) body

doHttpRequest ::
  ByteString ->
  Text ->
  Maybe Aeson.Object ->
  Maybe (Text, ByteString) ->
  IO Aeson.Value
doHttpRequest method url mHeaders mBody = do
  let timeoutMicros = 30 * 1000000
  result <-
    try <| do
      req0 <- HTTP.parseRequest (Text.unpack url)
      let req1 =
            HTTP.setRequestMethod method
              <| HTTP.setRequestHeader "User-Agent" ["OmniAgent/1.0"]
              <| HTTP.setRequestResponseTimeout (HTTPClient.responseTimeoutMicro timeoutMicros)
              <| req0
          req2 = case mHeaders of
            Nothing -> req1
            Just hdrs -> foldr addHeader req1 (KeyMap.toList hdrs)
          req3 = case mBody of
            Nothing -> req2
            Just (ct, bodyBytes) ->
              HTTP.setRequestHeader "Content-Type" [TE.encodeUtf8 ct]
                <| HTTP.setRequestBodyLBS (BL.fromStrict bodyBytes)
                <| req2
      mResp <- timeout timeoutMicros (HTTP.httpLBS req3)
      case mResp of
        Nothing -> pure (Left "Request timed out after 30 seconds")
        Just resp -> pure (Right resp)
  case result of
    Left (e :: SomeException) -> pure <| mkError ("Request failed: " <> tshow e)
    Right (Left err) -> pure <| mkError err
    Right (Right response) -> do
      let status = HTTP.getResponseStatusCode response
          respHeaders = HTTP.getResponseHeaders response
          headerObj =
            Aeson.object
              [ Key.fromText (TE.decodeUtf8 (CI.original k)) .= TE.decodeUtf8 v
                | (k, v) <- respHeaders
              ]
          body = TE.decodeUtf8With (\_ _ -> Just '?') (BL.toStrict (HTTP.getResponseBody response))
      pure
        <| Aeson.toJSON
        <| HttpResult
          { httpResultStatus = status,
            httpResultHeaders = headerObj,
            httpResultBody = body
          }
  where
    addHeader :: (Aeson.Key, Aeson.Value) -> HTTP.Request -> HTTP.Request
    addHeader (k, v) req =
      let headerName = CI.mk (TE.encodeUtf8 (Key.toText k))
          headerValue = TE.encodeUtf8 (valueToText v)
       in HTTP.addRequestHeader headerName headerValue req

valueToText :: Aeson.Value -> Text
valueToText (Aeson.String s) = s
valueToText (Aeson.Number n) = tshow n
valueToText (Aeson.Bool b) = if b then "true" else "false"
valueToText Aeson.Null = ""
valueToText other = TE.decodeUtf8 (BL.toStrict (Aeson.encode other))

mkError :: Text -> Aeson.Value
mkError err =
  Aeson.object
    [ "status" .= (-1 :: Int),
      "headers" .= Aeson.object [],
      "body" .= err
    ]