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

-- | Todo tool with due dates and reminders.
--
-- Provides user-scoped todos with optional due dates.
--
-- : out omni-agent-tools-todos
-- : dep aeson
-- : dep sqlite-simple
-- : dep time
module Omni.Agent.Tools.Todos
  ( -- * Tools
    todoAddTool,
    todoListTool,
    todoCompleteTool,
    todoDeleteTool,

    -- * Direct API
    Todo (..),
    createTodo,
    listTodos,
    listPendingTodos,
    listOverdueTodos,
    completeTodo,
    deleteTodo,

    -- * Reminders
    listTodosDueForReminder,
    markReminderSent,
    reminderInterval,

    -- * Database
    initTodosTable,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Text as Text
import Data.Time (LocalTime, NominalDiffTime, TimeZone, UTCTime, addUTCTime, getCurrentTime, localTimeToUTC, minutesToTimeZone, utcToLocalTime)
import Data.Time.Format (defaultTimeLocale, formatTime, parseTimeM)
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.Todos"
    [ Test.unit "todoAddTool has correct schema" <| do
        let tool = todoAddTool "test-user-id"
        Engine.toolName tool Test.@=? "todo_add",
      Test.unit "todoListTool has correct schema" <| do
        let tool = todoListTool "test-user-id"
        Engine.toolName tool Test.@=? "todo_list",
      Test.unit "todoCompleteTool has correct schema" <| do
        let tool = todoCompleteTool "test-user-id"
        Engine.toolName tool Test.@=? "todo_complete",
      Test.unit "todoDeleteTool has correct schema" <| do
        let tool = todoDeleteTool "test-user-id"
        Engine.toolName tool Test.@=? "todo_delete",
      Test.unit "Todo JSON roundtrip" <| do
        now <- getCurrentTime
        let td =
              Todo
                { todoId = 1,
                  todoUserId = "user-123",
                  todoTitle = "Buy milk",
                  todoDueDate = Just now,
                  todoCompleted = False,
                  todoCreatedAt = now,
                  todoLastRemindedAt = Nothing
                }
        case Aeson.decode (Aeson.encode td) of
          Nothing -> Test.assertFailure "Failed to decode Todo"
          Just decoded -> do
            todoTitle decoded Test.@=? "Buy milk"
            todoCompleted decoded Test.@=? False,
      Test.unit "parseDueDate handles various formats" <| do
        isJust (parseDueDate "2024-12-25") Test.@=? True
        isJust (parseDueDate "2024-12-25 14:00") Test.@=? True
    ]

data Todo = Todo
  { todoId :: Int,
    todoUserId :: Text,
    todoTitle :: Text,
    todoDueDate :: Maybe UTCTime,
    todoCompleted :: Bool,
    todoCreatedAt :: UTCTime,
    todoLastRemindedAt :: Maybe UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Todo where
  toJSON td =
    Aeson.object
      [ "id" .= todoId td,
        "user_id" .= todoUserId td,
        "title" .= todoTitle td,
        "due_date" .= todoDueDate td,
        "completed" .= todoCompleted td,
        "created_at" .= todoCreatedAt td,
        "last_reminded_at" .= todoLastRemindedAt td
      ]

instance Aeson.FromJSON Todo where
  parseJSON =
    Aeson.withObject "Todo" <| \v ->
      (Todo </ (v .: "id"))
        <*> (v .: "user_id")
        <*> (v .: "title")
        <*> (v .:? "due_date")
        <*> (v .: "completed")
        <*> (v .: "created_at")
        <*> (v .:? "last_reminded_at")

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

initTodosTable :: SQL.Connection -> IO ()
initTodosTable conn = do
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS todos (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL,\
    \  title TEXT NOT NULL,\
    \  due_date TIMESTAMP,\
    \  completed INTEGER NOT NULL DEFAULT 0,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\
    \  last_reminded_at TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_todos_user ON todos(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_todos_due ON todos(user_id, due_date)"
  migrateTodosTable conn

migrateTodosTable :: SQL.Connection -> IO ()
migrateTodosTable conn = do
  cols <- SQL.query_ conn "PRAGMA table_info(todos)" :: IO [(Int, Text, Text, Int, Maybe Text, Int)]
  let colNames = map (\(_, name, _, _, _, _) -> name) cols
  unless ("last_reminded_at" `elem` colNames) <| do
    SQL.execute_ conn "ALTER TABLE todos ADD COLUMN last_reminded_at TIMESTAMP"

easternTimeZone :: TimeZone
easternTimeZone = minutesToTimeZone (-300)

parseDueDate :: Text -> Maybe UTCTime
parseDueDate txt =
  let s = Text.unpack txt
      parseLocal :: Maybe LocalTime
      parseLocal =
        parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M" s
          <|> parseTimeM True defaultTimeLocale "%Y-%m-%d" s
          <|> parseTimeM True defaultTimeLocale "%Y-%m-%dT%H:%M:%S" s
   in fmap (localTimeToUTC easternTimeZone) parseLocal
        <|> parseTimeM True defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" s

createTodo :: Text -> Text -> Maybe Text -> IO Todo
createTodo uid title maybeDueDateStr = do
  now <- getCurrentTime
  let dueDate = maybeDueDateStr +> parseDueDate
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.execute
      conn
      "INSERT INTO todos (user_id, title, due_date, completed, created_at) VALUES (?, ?, ?, 0, ?)"
      (uid, title, dueDate, now)
    rowId <- SQL.lastInsertRowId conn
    pure
      Todo
        { todoId = fromIntegral rowId,
          todoUserId = uid,
          todoTitle = title,
          todoDueDate = dueDate,
          todoCompleted = False,
          todoCreatedAt = now,
          todoLastRemindedAt = Nothing
        }

listTodos :: Text -> Int -> IO [Todo]
listTodos uid limit =
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.query
      conn
      "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \
      \FROM todos WHERE user_id = ? \
      \ORDER BY completed ASC, due_date ASC NULLS LAST, created_at DESC LIMIT ?"
      (uid, limit)

listPendingTodos :: Text -> Int -> IO [Todo]
listPendingTodos uid limit =
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.query
      conn
      "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \
      \FROM todos WHERE user_id = ? AND completed = 0 \
      \ORDER BY due_date ASC NULLS LAST, created_at DESC LIMIT ?"
      (uid, limit)

listOverdueTodos :: Text -> IO [Todo]
listOverdueTodos uid = do
  now <- getCurrentTime
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.query
      conn
      "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \
      \FROM todos WHERE user_id = ? AND completed = 0 AND due_date < ? \
      \ORDER BY due_date ASC"
      (uid, now)

reminderInterval :: NominalDiffTime
reminderInterval = 24 * 60 * 60

listTodosDueForReminder :: IO [Todo]
listTodosDueForReminder = do
  now <- getCurrentTime
  let cutoff = addUTCTime (negate reminderInterval) now
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.query
      conn
      "SELECT id, user_id, title, due_date, completed, created_at, last_reminded_at \
      \FROM todos \
      \WHERE completed = 0 \
      \  AND due_date IS NOT NULL \
      \  AND due_date < ? \
      \  AND (last_reminded_at IS NULL OR last_reminded_at < ?)"
      (now, cutoff)

markReminderSent :: Int -> IO ()
markReminderSent tid = do
  now <- getCurrentTime
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.execute
      conn
      "UPDATE todos SET last_reminded_at = ? WHERE id = ?"
      (now, tid)

completeTodo :: Text -> Int -> IO Bool
completeTodo uid tid =
  Memory.withMemoryDb <| \conn -> do
    initTodosTable conn
    SQL.execute
      conn
      "UPDATE todos SET completed = 1 WHERE id = ? AND user_id = ?"
      (tid, uid)
    changes <- SQL.changes conn
    pure (changes > 0)

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

todoAddTool :: Text -> Engine.Tool
todoAddTool uid =
  Engine.Tool
    { Engine.toolName = "todo_add",
      Engine.toolDescription =
        "Add a todo item with optional due date. Use for tasks, reminders, "
          <> "or anything the user needs to remember to do. "
          <> "Due date format: 'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM'.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "title"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("What needs to be done" :: Text)
                      ],
                  "due_date"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Optional due date in Eastern time: 'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM'" :: Text)
                      ]
                ],
            "required" .= (["title"] :: [Text])
          ],
      Engine.toolExecute = executeTodoAdd uid
    }

executeTodoAdd :: Text -> Aeson.Value -> IO Aeson.Value
executeTodoAdd uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: TodoAddArgs) -> do
      td <- createTodo uid (taTitle args) (taDueDate args)
      let dueDateMsg = case todoDueDate td of
            Just d ->
              let localTime = utcToLocalTime easternTimeZone d
               in " (due: " <> Text.pack (formatTime defaultTimeLocale "%Y-%m-%d %H:%M ET" localTime) <> ")"
            Nothing -> ""
      pure
        ( Aeson.object
            [ "success" .= True,
              "todo_id" .= todoId td,
              "message" .= ("Added todo: " <> todoTitle td <> dueDateMsg)
            ]
        )

data TodoAddArgs = TodoAddArgs
  { taTitle :: Text,
    taDueDate :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON TodoAddArgs where
  parseJSON =
    Aeson.withObject "TodoAddArgs" <| \v ->
      (TodoAddArgs </ (v .: "title"))
        <*> (v .:? "due_date")

todoListTool :: Text -> Engine.Tool
todoListTool uid =
  Engine.Tool
    { Engine.toolName = "todo_list",
      Engine.toolDescription =
        "List todos. By default shows pending (incomplete) todos. "
          <> "Can show all todos or just overdue ones.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "filter"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Filter: 'pending' (default), 'all', or 'overdue'" :: Text)
                      ],
                  "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Max todos to return (default: 20)" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeTodoList uid
    }

executeTodoList :: Text -> Aeson.Value -> IO Aeson.Value
executeTodoList uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: TodoListArgs) -> do
      let lim = min 50 (max 1 (tlLimit args))
      todos <- case tlFilter args of
        "all" -> listTodos uid lim
        "overdue" -> listOverdueTodos uid
        _ -> listPendingTodos uid lim
      pure
        ( Aeson.object
            [ "success" .= True,
              "count" .= length todos,
              "todos" .= formatTodosForLLM todos
            ]
        )

formatTodosForLLM :: [Todo] -> Text
formatTodosForLLM [] = "No todos found."
formatTodosForLLM todos =
  Text.unlines (map formatTodo todos)
  where
    formatTodo td =
      let status = if todoCompleted td then "[x]" else "[ ]"
          dueStr = case todoDueDate td of
            Just d ->
              let localTime = utcToLocalTime easternTimeZone d
               in " (due: " <> Text.pack (formatTime defaultTimeLocale "%Y-%m-%d %H:%M ET" localTime) <> ")"
            Nothing -> ""
       in status <> " " <> todoTitle td <> dueStr <> " (id: " <> tshow (todoId td) <> ")"

data TodoListArgs = TodoListArgs
  { tlFilter :: Text,
    tlLimit :: Int
  }
  deriving (Generic)

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

todoCompleteTool :: Text -> Engine.Tool
todoCompleteTool uid =
  Engine.Tool
    { Engine.toolName = "todo_complete",
      Engine.toolDescription =
        "Mark a todo as completed. Use when the user says they finished something.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "todo_id"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("The ID of the todo to complete" :: Text)
                      ]
                ],
            "required" .= (["todo_id"] :: [Text])
          ],
      Engine.toolExecute = executeTodoComplete uid
    }

executeTodoComplete :: Text -> Aeson.Value -> IO Aeson.Value
executeTodoComplete uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: TodoCompleteArgs) -> do
      completed <- completeTodo uid (tcTodoId args)
      if completed
        then
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message" .= ("Todo marked as complete" :: Text)
                ]
            )
        else
          pure
            ( Aeson.object
                [ "success" .= False,
                  "error" .= ("Todo not found" :: Text)
                ]
            )

newtype TodoCompleteArgs = TodoCompleteArgs
  { tcTodoId :: Int
  }
  deriving (Generic)

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

todoDeleteTool :: Text -> Engine.Tool
todoDeleteTool uid =
  Engine.Tool
    { Engine.toolName = "todo_delete",
      Engine.toolDescription =
        "Delete a todo permanently. Use when a todo is no longer needed.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "todo_id"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("The ID of the todo to delete" :: Text)
                      ]
                ],
            "required" .= (["todo_id"] :: [Text])
          ],
      Engine.toolExecute = executeTodoDelete uid
    }

executeTodoDelete :: Text -> Aeson.Value -> IO Aeson.Value
executeTodoDelete uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: TodoDeleteArgs) -> do
      deleted <- deleteTodo uid (tdTodoId args)
      if deleted
        then
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message" .= ("Todo deleted" :: Text)
                ]
            )
        else
          pure
            ( Aeson.object
                [ "success" .= False,
                  "error" .= ("Todo not found" :: Text)
                ]
            )

newtype TodoDeleteArgs = TodoDeleteArgs
  { tdTodoId :: Int
  }
  deriving (Generic)

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