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

-- | Quick notes tool for agents.
--
-- Provides simple CRUD for tagged notes stored in memory.db.
--
-- : out omni-agent-tools-notes
-- : dep aeson
-- : dep sqlite-simple
module Omni.Agent.Tools.Notes
  ( -- * Tools
    noteAddTool,
    noteListTool,
    noteDeleteTool,

    -- * Direct API
    Note (..),
    createNote,
    listNotes,
    listNotesByTopic,
    deleteNote,

    -- * Database
    initNotesTable,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Text as Text
import Data.Time (UTCTime, getCurrentTime)
import qualified Database.SQLite.Simple as SQL
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Memory as Memory
import qualified Omni.Test as Test

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Tools.Notes"
    [ Test.unit "noteAddTool has correct schema" <| do
        let tool = noteAddTool "test-user-id"
        Engine.toolName tool Test.@=? "note_add",
      Test.unit "noteListTool has correct schema" <| do
        let tool = noteListTool "test-user-id"
        Engine.toolName tool Test.@=? "note_list",
      Test.unit "noteDeleteTool has correct schema" <| do
        let tool = noteDeleteTool "test-user-id"
        Engine.toolName tool Test.@=? "note_delete",
      Test.unit "Note JSON roundtrip" <| do
        now <- getCurrentTime
        let n =
              Note
                { noteId = 1,
                  noteUserId = "user-123",
                  noteTopic = "groceries",
                  noteContent = "Buy milk",
                  noteCreatedAt = now
                }
        case Aeson.decode (Aeson.encode n) of
          Nothing -> Test.assertFailure "Failed to decode Note"
          Just decoded -> do
            noteContent decoded Test.@=? "Buy milk"
            noteTopic decoded Test.@=? "groceries"
    ]

data Note = Note
  { noteId :: Int,
    noteUserId :: Text,
    noteTopic :: Text,
    noteContent :: Text,
    noteCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Note where
  toJSON n =
    Aeson.object
      [ "id" .= noteId n,
        "user_id" .= noteUserId n,
        "topic" .= noteTopic n,
        "content" .= noteContent n,
        "created_at" .= noteCreatedAt n
      ]

instance Aeson.FromJSON Note where
  parseJSON =
    Aeson.withObject "Note" <| \v ->
      (Note </ (v .: "id"))
        <*> (v .: "user_id")
        <*> (v .: "topic")
        <*> (v .: "content")
        <*> (v .: "created_at")

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

initNotesTable :: SQL.Connection -> IO ()
initNotesTable conn = do
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS notes (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL,\
    \  topic TEXT NOT NULL,\
    \  content TEXT NOT NULL,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_notes_user ON notes(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_notes_topic ON notes(user_id, topic)"

createNote :: Text -> Text -> Text -> IO Note
createNote uid topic content = do
  now <- getCurrentTime
  Memory.withMemoryDb <| \conn -> do
    initNotesTable conn
    SQL.execute
      conn
      "INSERT INTO notes (user_id, topic, content, created_at) VALUES (?, ?, ?, ?)"
      (uid, topic, content, now)
    rowId <- SQL.lastInsertRowId conn
    pure
      Note
        { noteId = fromIntegral rowId,
          noteUserId = uid,
          noteTopic = topic,
          noteContent = content,
          noteCreatedAt = now
        }

listNotes :: Text -> Int -> IO [Note]
listNotes uid limit =
  Memory.withMemoryDb <| \conn -> do
    initNotesTable conn
    SQL.query
      conn
      "SELECT id, user_id, topic, content, created_at \
      \FROM notes WHERE user_id = ? \
      \ORDER BY created_at DESC LIMIT ?"
      (uid, limit)

listNotesByTopic :: Text -> Text -> Int -> IO [Note]
listNotesByTopic uid topic limit =
  Memory.withMemoryDb <| \conn -> do
    initNotesTable conn
    SQL.query
      conn
      "SELECT id, user_id, topic, content, created_at \
      \FROM notes WHERE user_id = ? AND topic = ? \
      \ORDER BY created_at DESC LIMIT ?"
      (uid, topic, limit)

deleteNote :: Text -> Int -> IO Bool
deleteNote uid nid =
  Memory.withMemoryDb <| \conn -> do
    initNotesTable conn
    SQL.execute
      conn
      "DELETE FROM notes WHERE id = ? AND user_id = ?"
      (nid, uid)
    changes <- SQL.changes conn
    pure (changes > 0)

noteAddTool :: Text -> Engine.Tool
noteAddTool uid =
  Engine.Tool
    { Engine.toolName = "note_add",
      Engine.toolDescription =
        "Add a quick note on a topic. Use for reminders, lists, ideas, or anything "
          <> "the user wants to jot down. Topics help organize notes (e.g., 'groceries', "
          <> "'ideas', 'todo', 'recipes').",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "topic"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Topic/category for the note (e.g., 'groceries', 'todo')" :: Text)
                      ],
                  "content"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The note content" :: Text)
                      ]
                ],
            "required" .= (["topic", "content"] :: [Text])
          ],
      Engine.toolExecute = executeNoteAdd uid
    }

executeNoteAdd :: Text -> Aeson.Value -> IO Aeson.Value
executeNoteAdd uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: NoteAddArgs) -> do
      newNote <- createNote uid (naTopic args) (naContent args)
      pure
        ( Aeson.object
            [ "success" .= True,
              "note_id" .= noteId newNote,
              "message" .= ("Added note to '" <> noteTopic newNote <> "': " <> noteContent newNote)
            ]
        )

data NoteAddArgs = NoteAddArgs
  { naTopic :: Text,
    naContent :: Text
  }
  deriving (Generic)

instance Aeson.FromJSON NoteAddArgs where
  parseJSON =
    Aeson.withObject "NoteAddArgs" <| \v ->
      (NoteAddArgs </ (v .: "topic"))
        <*> (v .: "content")

noteListTool :: Text -> Engine.Tool
noteListTool uid =
  Engine.Tool
    { Engine.toolName = "note_list",
      Engine.toolDescription =
        "List notes, optionally filtered by topic. Use to show the user their "
          <> "saved notes or check what's on a specific list.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "topic"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Filter by topic (optional, omit to list all)" :: Text)
                      ],
                  "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Max notes to return (default: 20)" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeNoteList uid
    }

executeNoteList :: Text -> Aeson.Value -> IO Aeson.Value
executeNoteList uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: NoteListArgs) -> do
      let lim = min 50 (max 1 (nlLimit args))
      notes <- case nlTopic args of
        Just topic -> listNotesByTopic uid topic lim
        Nothing -> listNotes uid lim
      pure
        ( Aeson.object
            [ "success" .= True,
              "count" .= length notes,
              "notes" .= formatNotesForLLM notes
            ]
        )

formatNotesForLLM :: [Note] -> Text
formatNotesForLLM [] = "No notes found."
formatNotesForLLM notes =
  Text.unlines (map formatNote notes)
  where
    formatNote n =
      "[" <> noteTopic n <> "] " <> noteContent n <> " (id: " <> tshow (noteId n) <> ")"

data NoteListArgs = NoteListArgs
  { nlTopic :: Maybe Text,
    nlLimit :: Int
  }
  deriving (Generic)

instance Aeson.FromJSON NoteListArgs where
  parseJSON =
    Aeson.withObject "NoteListArgs" <| \v ->
      (NoteListArgs </ (v .:? "topic"))
        <*> (v .:? "limit" .!= 20)

noteDeleteTool :: Text -> Engine.Tool
noteDeleteTool uid =
  Engine.Tool
    { Engine.toolName = "note_delete",
      Engine.toolDescription =
        "Delete a note by its ID. Use after the user says they've completed an item "
          <> "or no longer need a note.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "note_id"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("The ID of the note to delete" :: Text)
                      ]
                ],
            "required" .= (["note_id"] :: [Text])
          ],
      Engine.toolExecute = executeNoteDelete uid
    }

executeNoteDelete :: Text -> Aeson.Value -> IO Aeson.Value
executeNoteDelete uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: NoteDeleteArgs) -> do
      deleted <- deleteNote uid (ndNoteId args)
      if deleted
        then
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message" .= ("Note deleted" :: Text)
                ]
            )
        else
          pure
            ( Aeson.object
                [ "success" .= False,
                  "error" .= ("Note not found or already deleted" :: Text)
                ]
            )

newtype NoteDeleteArgs = NoteDeleteArgs
  { ndNoteId :: Int
  }
  deriving (Generic)

instance Aeson.FromJSON NoteDeleteArgs where
  parseJSON =
    Aeson.withObject "NoteDeleteArgs" <| \v ->
      NoteDeleteArgs </ (v .: "note_id")