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

-- | Tool trace storage for Ava.
--
-- Records tool execution traces for debugging and analytics.
--
-- : out omni-ava-trace
-- : dep aeson
-- : dep sqlite-simple
-- : dep uuid
module Omni.Ava.Trace
  ( TraceRecord (..),
    initTraceDb,
    insertTrace,
    getTrace,
    cleanupOldTraces,
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.Text as Text
import qualified Data.UUID as UUID
import qualified Data.UUID.V4 as UUID
import qualified Database.SQLite.Simple as SQL
import qualified Database.SQLite.Simple.ToField as SQL
import qualified Omni.Test as Test

-- | Initialize the tool_traces table if it doesn't exist.
initTraceDb :: SQL.Connection -> IO ()
initTraceDb conn = do
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS tool_traces (\
    \  id TEXT PRIMARY KEY,\
    \  created_at TEXT NOT NULL,\
    \  tool_name TEXT NOT NULL,\
    \  input TEXT NOT NULL,\
    \  output TEXT NOT NULL,\
    \  duration_ms INTEGER NOT NULL,\
    \  user_id TEXT,\
    \  chat_id TEXT\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_traces_created ON tool_traces(created_at)"

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

test :: Test.Tree
test =
  Test.group
    "Omni.Ava.Trace"
    [ Test.unit "TraceRecord JSON roundtrip" <| do
        let tr =
              TraceRecord
                { trcId = "trace-123",
                  trcCreatedAt = "2024-01-15T10:30:00Z",
                  trcToolName = "web_search",
                  trcInput = "{\"query\":\"test\"}",
                  trcOutput = "{\"results\":[]}",
                  trcDurationMs = 150,
                  trcUserId = Just "user-456",
                  trcChatId = Just "chat-789"
                }
        case Aeson.decode (Aeson.encode tr) of
          Nothing -> Test.assertFailure "Failed to decode TraceRecord"
          Just decoded -> do
            trcToolName decoded Test.@=? "web_search"
            trcDurationMs decoded Test.@=? 150
    ]

data TraceRecord = TraceRecord
  { trcId :: Text,
    trcCreatedAt :: Text,
    trcToolName :: Text,
    trcInput :: Text,
    trcOutput :: Text,
    trcDurationMs :: Int,
    trcUserId :: Maybe Text,
    trcChatId :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TraceRecord where
  toJSON tr =
    Aeson.object
      [ "id" .= trcId tr,
        "created_at" .= trcCreatedAt tr,
        "tool_name" .= trcToolName tr,
        "input" .= trcInput tr,
        "output" .= trcOutput tr,
        "duration_ms" .= trcDurationMs tr,
        "user_id" .= trcUserId tr,
        "chat_id" .= trcChatId tr
      ]

instance Aeson.FromJSON TraceRecord where
  parseJSON =
    Aeson.withObject "TraceRecord" <| \v ->
      (TraceRecord </ (v Aeson..: "id"))
        <*> (v Aeson..: "created_at")
        <*> (v Aeson..: "tool_name")
        <*> (v Aeson..: "input")
        <*> (v Aeson..: "output")
        <*> (v Aeson..: "duration_ms")
        <*> (v Aeson..:? "user_id")
        <*> (v Aeson..:? "chat_id")

instance SQL.FromRow TraceRecord where
  fromRow =
    (TraceRecord </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field

instance SQL.ToRow TraceRecord where
  toRow tr =
    [ SQL.toField (trcId tr),
      SQL.toField (trcCreatedAt tr),
      SQL.toField (trcToolName tr),
      SQL.toField (trcInput tr),
      SQL.toField (trcOutput tr),
      SQL.toField (trcDurationMs tr),
      SQL.toField (trcUserId tr),
      SQL.toField (trcChatId tr)
    ]

insertTrace :: SQL.Connection -> TraceRecord -> IO Text
insertTrace conn tr = do
  tid <-
    if Text.null (trcId tr)
      then (Text.pack <. UUID.toString) </ UUID.nextRandom
      else pure (trcId tr)
  let trWithId = tr {trcId = tid}
  SQL.execute
    conn
    "INSERT INTO tool_traces (id, created_at, tool_name, input, output, duration_ms, user_id, chat_id) \
    \VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
    trWithId
  pure tid

getTrace :: SQL.Connection -> Text -> IO (Maybe TraceRecord)
getTrace conn tid = do
  results <-
    SQL.query
      conn
      "SELECT id, created_at, tool_name, input, output, duration_ms, user_id, chat_id \
      \FROM tool_traces WHERE id = ?"
      (SQL.Only tid)
  pure (listToMaybe results)

cleanupOldTraces :: SQL.Connection -> IO Int
cleanupOldTraces conn = do
  SQL.execute_
    conn
    "DELETE FROM tool_traces WHERE created_at < datetime('now', '-7 days')"
  SQL.changes conn