{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall #-}

module AI where

import Control.Exception.Annotated (checkpointCallStack)
import Control.Lens (re, traversed, view, (&), (.~), (^.), (^?))
import Control.Monad.Catch (MonadCatch)
import Control.Monad.Reader (MonadReader, liftIO)
import Creds
import Data.Aeson (object, (.=))
import Data.Aeson.Lens (key, _Array, _String)
import Data.Functor.Contravariant ((>$<))
import Data.Text (Text)
import Data.Text.Strict.Lens (utf8)
import HTTP (HTTPMetric, Retry (SafeTransient), Trace, logFormatter, postWith, tracerMetrics)
import M (HasKatipTracer, HasOpenrouterMetrics, KatipTracer, katipTracer, openrouterMetrics)
import Network.Wreq (defaults, header, responseBody)
import Tracer (Tracer (..))
import UnliftIO (MonadIO)
import Network.Wreq.Session (Session)
import Data.Generics.Product (typed, HasType)

mkTracerMetrics ::
  ( MonadReader r m,
    HasOpenrouterMetrics r HTTPMetric,
    HasKatipTracer r KatipTracer
  ) =>
  m (Tracer HTTP.Trace)
mkTracerMetrics = do
  l <- view katipTracer
  m <- view openrouterMetrics
  pure $ (logFormatter >$< l) <> HTTP.tracerMetrics m

completions ::
  ( MonadIO m,
    HasType Session r0,
    MonadCatch m,
    MonadReader r0 m,
    HasOpenrouterMetrics r0 HTTPMetric,
    HasKatipTracer r0 KatipTracer
  ) =>
  Creds ->
  Text ->
  m (Maybe Text)
completions c prompt = checkpointCallStack $ do
  let opts = defaults & header "Authorization" .~ ["Bearer " <> c ^. openrouterKey . re utf8]
  let j =
        object
          [ "model" .= ("~deepseek/deepseek-flash-latest" :: Text),
            "messages"
              .= [object ["role" .= ("user" :: Text), "content" .= prompt]]
          ]
  metrics <- mkTracerMetrics
  session <- view $ typed @Session
  r <- liftIO $ postWith metrics session SafeTransient opts "https://openrouter.ai/api/v1/chat/completions" j
  pure $ r ^? responseBody . key "choices" . _Array . traversed . key "message" . key "content" . _String