{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StrictData #-}
{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE RecordWildCards #-}

module Embeddings
  ( generateEmbedding,
    indexEntries,
    search,
    chunkArticle,
    tokensApprox,
    chunkMaxTokens,
    Embedding(..), -- for testing
  ) where

import Control.Exception.Annotated (checkpointCallStack)
import Control.Lens (re, traversed, view, (&), (.~), (^.), (^..), (^?))
import Control.Monad (when)
import Control.Monad.Catch (MonadCatch, MonadThrow, MonadMask)
import Control.Monad.Reader (MonadReader, liftIO)
import Creds (Creds, openrouterKey)
import Data.Aeson (object, (.=))
import Data.Aeson.Lens (key, _Array, _Double)
import Data.Functor.Contravariant ((>$<))
import Data.Generics.Product (HasType, typed)
import Data.List (sortOn)
import Data.Map.Strict qualified as Map
import Data.Ord (Down (..))
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Strict.Lens (utf8)
import Data.Vector.Unboxed qualified as UV
import Database.SQLite.Simple (Only (..), open, close, SQLData (SQLBlob), ResultError (ConversionFailed))
import Database.SQLite.Simple qualified as SQL
import GHC.Stack (HasCallStack)
import HTTP (HTTPMetric, Retry (SafeTransient), Trace, logFormatter, postWith, tracerMetrics)
import Katip qualified as K
import M (HasEmbeddingMetrics, HasKatipTracer, HasWallabagMetrics, embeddingMetrics, katipTracer, KatipTracer)
import Network.Wreq (defaults, header, responseBody)
import Network.Wreq.Session (Session)
import Pipes (lift, (>->))
import Pipes qualified
import Tracer (Tracer (..))
import UnliftIO (MonadIO, MonadUnliftIO)
import Wallabag (Entry, Scan (Everything), content, identifier, title, url)
import qualified Wallabag
import Wallabag.Auth (Auth)
import GHC.Generics (Generic)
import Database.SQLite.Simple.ToField (ToField(..))
import Database.SQLite.Simple.FromField (FromField(..), fieldData, returnError)
import qualified Data.Serialize as Cereal
import qualified Pipes.Safe
import Pipes.Safe (MonadSafe)
import qualified Pipes.Prelude

embeddingUrl :: String
embeddingUrl = "https://openrouter.ai/api/v1/embeddings"

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

generateEmbedding ::
  ( MonadIO m,
    HasType Session r,
    HasType Creds r,
    MonadReader r m,
    HasEmbeddingMetrics r HTTPMetric,
    HasKatipTracer r KatipTracer,
    MonadCatch m,
    K.KatipContext m
  ) =>
  Text ->
  m (Maybe (UV.Vector Float))
generateEmbedding input = checkpointCallStack $ do
  c <- view $ typed @Creds
  let opts = defaults & header "Authorization" .~ ["Bearer " <> c ^. openrouterKey . re utf8]
  let j =
        object
          [ "model" .= ("text-embedding-3-small" :: Text),
            "input" .= input
          ]
  tm <- mkTracerMetrics
  session <- view $ typed @Session
  r <- liftIO $ postWith tm session SafeTransient opts embeddingUrl j
  case r ^? responseBody . key "data" . _Array . traversed . key "embedding" of
    Just embeddingVal -> do
      let floats = embeddingVal ^.. _Array . traversed . _Double
      if null floats
        then do
          K.logFM K.ErrorS $ K.ls $ "Parsing gave empty embedding array from: " <> T.pack (show (r ^. responseBody))
          pure Nothing
        else pure $ Just $ UV.fromList (map realToFrac floats)
    Nothing -> do
      K.logFM K.ErrorS $ K.ls $ "Couldn't find embedding in response: " <> T.pack (show (r ^. responseBody))
      pure Nothing

cosineSimilarity :: UV.Vector Float -> UV.Vector Float -> Float
cosineSimilarity a b = dot / (normA * normB)
  where
    dot = UV.sum $ UV.zipWith (*) a b
    normA = sqrt $ UV.sum $ UV.map (^ (2 :: Int)) a
    normB = sqrt $ UV.sum $ UV.map (^ (2 :: Int)) b

indexEntries ::
  ( MonadIO m,
    HasType Session r,
    HasType Auth r,
    HasType Creds r,
    MonadReader r m,
    K.KatipContext m,
    HasWallabagMetrics r HTTPMetric,
    HasEmbeddingMetrics r HTTPMetric,
    HasKatipTracer r KatipTracer,
    MonadThrow m,
    MonadCatch m,
    HasCallStack,
    MonadUnliftIO m
  ) =>
  FilePath ->
  Bool ->
  m ()
indexEntries dbPath rebuild = do
  liftIO $ do
    conn <- open dbPath
    SQL.execute_ conn
      "CREATE TABLE IF NOT EXISTS embeddings (entry_id INTEGER NOT NULL, chunk_idx INTEGER NOT NULL, url TEXT NOT NULL, title TEXT NOT NULL, embedding TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT (datetime('now')), PRIMARY KEY (entry_id, chunk_idx))"
    when rebuild $ SQL.execute_ conn "DELETE FROM embeddings"
    close conn
  Pipes.runEffect $
    Pipes.for
      (Wallabag.sourceEntries Everything)
      (lift . processEntry dbPath)

chunkMaxTokens :: Int
chunkMaxTokens = 6000

tokensApprox :: Text -> Int
tokensApprox t = T.length t `div` 3

stripTags :: Text -> Text
stripTags t =
  case T.breakOn "<" t of
    (before, rest)
      | T.null rest -> before
      | otherwise ->
          case T.breakOn ">" (T.drop 1 rest) of
            (_, after) -> stripTags before <> stripTags (T.drop 1 after)

normalize :: Text -> Text
normalize = T.unwords . T.words

blockEndTags :: [Text]
blockEndTags =
  ["</p>", "</div>", "</blockquote>", "</section>", "</li>"]
    <> map (\n -> "</h" <> T.pack (show n) <> ">") [1 .. 6 :: Int]

sentinel :: Text
sentinel = "\x00BLOCK\x00"

replaceBlockEnds :: Text -> Text
replaceBlockEnds html = foldr (\tag -> T.replace tag sentinel) html blockEndTags

extractBlocks :: Text -> [Text]
extractBlocks html =
  map normalize
    . map stripTags
    . filter (not . T.null)
    $ T.splitOn sentinel (replaceBlockEnds html)

splitSentences :: Text -> [Text]
splitSentences t =
  filter (not . T.null)
    . map T.strip
    $ T.splitOn sent
    $ T.replace ". " (". " <> sent)
    $ T.replace "? " ("? " <> sent)
    $ T.replace "! " ("! " <> sent)
    $ t
  where
    sent = "\x00SENT\x00"

hardWrap :: Text -> [Text]
hardWrap t
  | T.null t = []
  | otherwise = let (a, b) = T.splitAt (chunkMaxTokens * 3) t in a : hardWrap b

coalesceBlocks :: Int -> [Text] -> [Text]
coalesceBlocks titleSize blocks = reverse (go [] 0 blocks)
  where
    chunkMaxChars = (chunkMaxTokens * 3) - titleSize

    go :: [Text] -> Int -> [Text] -> [Text]
    go acc _ [] = [T.concat (reverse acc) | not (null acc)]
    go acc len (x : xs)
      | T.length x > chunkMaxChars && null acc =
          let sentences = splitSentences x
              splitChunks =
                if sentences == [x]
                  then hardWrap x
                  else sentences
           in go acc len (splitChunks ++ xs)
      | not (null acc) && len + T.length x > chunkMaxChars =
          T.concat (reverse acc) : go [] 0 (x : xs)
      | otherwise =
          go (x : acc) (len + T.length x) xs

chunkArticle :: Text -> Text -> [Text]
chunkArticle title_ html =
  map (prefixTitle title_) (coalesceBlocks titleSize (extractBlocks html))
  where
    prefixTitle t chunk = "Title: " <> t <> "\n\n" <> chunk
    titleSize = T.length ("Title: " <> title_)

processEntry ::
  ( MonadIO m,
    HasType Session r,
    HasType Creds r,
    MonadReader r m,
    HasEmbeddingMetrics r HTTPMetric,
    HasKatipTracer r KatipTracer,
    K.KatipContext m,
    MonadCatch m
  ) =>
  FilePath ->
  Entry ->
  m ()
processEntry dbPath e = do
  conn <- liftIO $ open dbPath
  existing <- liftIO $ (SQL.query conn "SELECT 1 FROM embeddings WHERE entry_id = ?" (Only (e ^. identifier)) :: IO [Only Int])
  if null existing
    then do
      let chunks = chunkArticle (e ^. title) (e ^. content)
      if null chunks
        then do
          K.logFM K.WarningS $ K.ls $ "Skipping (no extractable content): " <> e ^. title
          liftIO $ close conn
        else do
          K.logFM K.InfoS $ K.ls $ "Indexing: " <> e ^. title <> " (" <> T.pack (show (length chunks)) <> " chunks)"
          mapM_ (processChunk conn (e ^. identifier) (e ^. url) (e ^. title)) (zip [0 :: Int ..] chunks)
          liftIO $ close conn
    else do
      K.logFM K.DebugS $ K.ls $ "Skipping (already indexed): " <> e ^. title
      liftIO $ close conn


newtype Embedding = Embedding (UV.Vector Float)
  deriving (Show, Eq)


instance ToField Embedding where
  toField (Embedding vec) = toField $ Cereal.encode $ UV.toList vec

instance FromField Embedding where
  fromField f = do
    case fieldData f of
      SQLBlob bs -> either (returnError ConversionFailed f) (pure . Embedding . UV.fromList) (Cereal.decode bs)
      _ -> returnError ConversionFailed f "Failed to convert to embedding vector"

data EmbeddingRow = EmbeddingRow
  { entryId :: Integer
  , chunkIdx :: Int
  , entryUrl :: Text
  , entryTitle :: Text
  , embedding :: Embedding
  }
  deriving (Generic)

instance SQL.ToRow EmbeddingRow

instance SQL.FromRow EmbeddingRow

processChunk ::
  ( MonadIO m,
    HasType Session r,
    HasType Creds r,
    MonadReader r m,
    HasEmbeddingMetrics r HTTPMetric,
    HasKatipTracer r KatipTracer,
    K.KatipContext m,
    MonadCatch m
  ) =>
  SQL.Connection ->
  Integer ->
  Text ->
  Text ->
  (Int, Text) ->
  m ()
processChunk conn eid url_ title_ (idx, chunk) = do
  mVec <- generateEmbedding chunk
  case mVec of
    Just vec -> do
      let embeddingRow = EmbeddingRow eid idx url_ title_ (Embedding vec)
      liftIO $ SQL.execute conn
        "INSERT OR REPLACE INTO embeddings (entry_id, chunk_idx, url, title, embedding) VALUES (?, ?, ?, ?, ?)"
        embeddingRow
    Nothing ->
      K.logFM K.WarningS $ K.ls $ "Failed to generate embedding for chunk " <> T.pack (show idx) <> " of: " <> title_


search ::
  ( MonadIO m,
    HasType Session r,
    HasType Creds r,
    MonadReader r m,
    HasEmbeddingMetrics r HTTPMetric,
    HasKatipTracer r KatipTracer,
    MonadCatch m,
    K.KatipContext m,
    MonadMask m
  ) =>
  FilePath ->
  Text ->
  Int ->
  m [(Double, Integer, Text, Text)]
search dbPath q limit_ = do
  queryVec <- generateEmbedding q
  case queryVec of
    Nothing -> pure []
    Just qv -> do
      conn <- liftIO $ open dbPath
      results <- Pipes.Safe.runSafeT $ Pipes.runEffect (Pipes.Prelude.fold bestChunk Map.empty topResults $ query_ conn  "SELECT entry_id, chunk_idx, url, title, embedding FROM embeddings" >-> Pipes.Prelude.mapMaybe (scoreRow qv))
      liftIO $ close conn
      pure results
  where
    bestChunk m (s, eid, url_, title_) =
      Map.insertWith
        (\(s1, u, t) (s2, _, _) -> (max s1 s2, u, t))
        eid
        (s, url_, title_)
        m
    sortByScore = sortOn (Down . (\(_, (s, _, _)) -> s)) . Map.toList
    topResults = map (\(eid, (s, url_, title_)) -> (realToFrac s, eid, url_, title_)) . take limit_ . sortByScore


query_ :: (MonadSafe m, SQL.FromRow a) => SQL.Connection -> SQL.Query -> Pipes.Producer' a m ()
query_ conn q = Pipes.Safe.bracket (liftIO openStmt) (liftIO . SQL.closeStatement) eachRow
  where
    openStmt = do
      SQL.openStatement conn q
    eachRow stmt = do
      row <- liftIO $ SQL.nextRow stmt
      maybe (pure ()) (\r -> Pipes.yield r >> eachRow stmt) row


scoreRow ::
  UV.Vector Float ->
  EmbeddingRow ->
  Maybe (Float, Integer, Text, Text)
scoreRow qv EmbeddingRow{embedding = Embedding uv, ..} =
  let sim = cosineSimilarity qv uv
  in Just (sim, entryId, entryUrl, entryTitle)