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

-- | Web search tool using Kagi Search API.
--
-- Provides web search capabilities for agents.
--
-- : out omni-agent-tools-websearch
-- : dep aeson
-- : dep http-conduit
module Omni.Agent.Tools.WebSearch
  ( -- * Tool
    webSearchTool,

    -- * Direct API
    kagiSearch,
    SearchResult (..),

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Network.HTTP.Simple as HTTP
import qualified Network.HTTP.Types.URI as URI
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Test as Test

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Tools.WebSearch"
    [ Test.unit "SearchResult JSON parsing" <| do
        let json =
              Aeson.object
                [ "t" .= (0 :: Int),
                  "url" .= ("https://example.com" :: Text),
                  "title" .= ("Example Title" :: Text),
                  "snippet" .= ("This is a snippet" :: Text)
                ]
        case parseSearchResult json of
          Nothing -> Test.assertFailure "Failed to parse search result"
          Just sr -> do
            srUrl sr Test.@=? "https://example.com"
            srTitle sr Test.@=? "Example Title"
            srSnippet sr Test.@=? Just "This is a snippet",
      Test.unit "webSearchTool has correct schema" <| do
        let tool = webSearchTool "test-key"
        Engine.toolName tool Test.@=? "web_search",
      Test.unit "formatResultsForLLM formats correctly" <| do
        let results =
              [ SearchResult "https://a.com" "Title A" (Just "Snippet A") Nothing,
                SearchResult "https://b.com" "Title B" Nothing Nothing
              ]
            formatted = formatResultsForLLM results
        ("Title A" `Text.isInfixOf` formatted) Test.@=? True
        ("https://a.com" `Text.isInfixOf` formatted) Test.@=? True
    ]

data SearchResult = SearchResult
  { srUrl :: Text,
    srTitle :: Text,
    srSnippet :: Maybe Text,
    srPublished :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SearchResult where
  toJSON r =
    Aeson.object
      [ "url" .= srUrl r,
        "title" .= srTitle r,
        "snippet" .= srSnippet r,
        "published" .= srPublished r
      ]

parseSearchResult :: Aeson.Value -> Maybe SearchResult
parseSearchResult val = do
  Aeson.Object obj <- pure val
  t <- case KeyMap.lookup "t" obj of
    Just (Aeson.Number n) -> Just (round n :: Int)
    _ -> Nothing
  guard (t == 0)
  url <- case KeyMap.lookup "url" obj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  title <- case KeyMap.lookup "title" obj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  let snippet = case KeyMap.lookup "snippet" obj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
      published = case KeyMap.lookup "published" obj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
  pure SearchResult {srUrl = url, srTitle = title, srSnippet = snippet, srPublished = published}

kagiSearch :: Text -> Text -> Int -> IO (Either Text [SearchResult])
kagiSearch apiKey query limit = do
  let encodedQuery = TE.decodeUtf8 (URI.urlEncode False (TE.encodeUtf8 query))
      url = "https://kagi.com/api/v0/search?q=" <> Text.unpack encodedQuery <> "&limit=" <> show limit
  result <-
    try <| do
      req0 <- HTTP.parseRequest url
      let req =
            HTTP.setRequestMethod "GET"
              <| HTTP.setRequestHeader "Authorization" ["Bot " <> TE.encodeUtf8 apiKey]
              <| req0
      HTTP.httpLBS req
  case result of
    Left (e :: SomeException) ->
      pure (Left ("Kagi API error: " <> tshow e))
    Right response -> do
      let status = HTTP.getResponseStatusCode response
      if status >= 200 && status < 300
        then case Aeson.decode (HTTP.getResponseBody response) of
          Just (Aeson.Object obj) -> case KeyMap.lookup "data" obj of
            Just (Aeson.Array arr) ->
              pure (Right (mapMaybe parseSearchResult (toList arr)))
            _ -> pure (Left "No data in response")
          _ -> pure (Left "Failed to parse Kagi response")
        else case Aeson.decode (HTTP.getResponseBody response) of
          Just (Aeson.Object obj) -> case KeyMap.lookup "error" obj of
            Just errArr -> pure (Left ("Kagi error: " <> tshow errArr))
            _ -> pure (Left ("Kagi HTTP error: " <> tshow status))
          _ -> pure (Left ("Kagi HTTP error: " <> tshow status))

formatResultsForLLM :: [SearchResult] -> Text
formatResultsForLLM [] = "No results found."
formatResultsForLLM results =
  Text.unlines (zipWith formatResult [1 ..] results)
  where
    formatResult :: Int -> SearchResult -> Text
    formatResult n r =
      tshow n
        <> ". "
        <> srTitle r
        <> "\n   "
        <> srUrl r
        <> maybe "" (\s -> "\n   " <> Text.take 200 s) (srSnippet r)

webSearchTool :: Text -> Engine.Tool
webSearchTool apiKey =
  Engine.Tool
    { Engine.toolName = "web_search",
      Engine.toolDescription =
        "Search the web using Kagi. Use this to find current information, "
          <> "verify facts, look up documentation, or research topics. "
          <> "Returns titles, URLs, and snippets from search results.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "query"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The search query" :: Text)
                      ],
                  "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Max results to return (default: 5, max: 10)" :: Text)
                      ]
                ],
            "required" .= (["query"] :: [Text])
          ],
      Engine.toolExecute = executeWebSearch apiKey
    }

executeWebSearch :: Text -> Aeson.Value -> IO Aeson.Value
executeWebSearch apiKey v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: WebSearchArgs) -> do
      let lim = min 10 (max 1 (wsLimit args))
      result <- kagiSearch apiKey (wsQuery args) lim
      case result of
        Left err ->
          pure (Aeson.object ["error" .= err])
        Right results ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "count" .= length results,
                  "results" .= formatResultsForLLM results
                ]
            )

data WebSearchArgs = WebSearchArgs
  { wsQuery :: Text,
    wsLimit :: Int
  }
  deriving (Generic)

instance Aeson.FromJSON WebSearchArgs where
  parseJSON =
    Aeson.withObject "WebSearchArgs" <| \v ->
      (WebSearchArgs </ (v Aeson..: "query"))
        <*> (v Aeson..:? "limit" Aeson..!= 5)