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

-- | Hledger tools for personal finance queries and transaction entry.
--
-- Provides hledger access for agents via the nix-shell in ~/fund.
--
-- : out omni-agent-tools-hledger
-- : dep aeson
-- : dep process
-- : dep directory
module Omni.Agent.Tools.Hledger
  ( -- * Tools
    hledgerBalanceTool,
    hledgerRegisterTool,
    hledgerAddTool,
    hledgerIncomeStatementTool,
    hledgerBalanceSheetTool,

    -- * All tools (for easy import)
    allHledgerTools,

    -- * Direct API
    queryBalance,
    queryRegister,
    addTransaction,
    incomeStatement,
    balanceSheet,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.IO as TextIO
import Data.Time (getCurrentTime, utcToLocalTime)
import Data.Time.Format (defaultTimeLocale, formatTime)
import Data.Time.LocalTime (getCurrentTimeZone)
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Test as Test
import System.Directory (doesFileExist)
import System.Process (readProcessWithExitCode)

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Tools.Hledger"
    [ Test.unit "hledgerBalanceTool has correct name" <| do
        Engine.toolName hledgerBalanceTool Test.@=? "hledger_balance",
      Test.unit "hledgerRegisterTool has correct name" <| do
        Engine.toolName hledgerRegisterTool Test.@=? "hledger_register",
      Test.unit "hledgerAddTool has correct name" <| do
        Engine.toolName hledgerAddTool Test.@=? "hledger_add",
      Test.unit "hledgerIncomeStatementTool has correct name" <| do
        Engine.toolName hledgerIncomeStatementTool Test.@=? "hledger_income_statement",
      Test.unit "hledgerBalanceSheetTool has correct name" <| do
        Engine.toolName hledgerBalanceSheetTool Test.@=? "hledger_balance_sheet",
      Test.unit "allHledgerTools has 5 tools" <| do
        length allHledgerTools Test.@=? 5
    ]

fundDir :: FilePath
fundDir = "/home/ben/fund"

journalFile :: FilePath
journalFile = fundDir <> "/ledger.journal"

transactionsFile :: FilePath
transactionsFile = fundDir <> "/telegram-transactions.journal"

runHledgerInFund :: [String] -> IO (Either Text Text)
runHledgerInFund args = do
  let fullArgs :: [String]
      fullArgs = ["-f", journalFile] <> args
      hledgerCmd :: String
      hledgerCmd = "hledger " ++ List.unwords fullArgs
      cmd :: String
      cmd = "cd " ++ fundDir ++ " && " ++ hledgerCmd
  result <-
    try <| readProcessWithExitCode "nix-shell" [fundDir ++ "/shell.nix", "--run", cmd] ""
  case result of
    Left (e :: SomeException) ->
      pure (Left ("hledger error: " <> tshow e))
    Right (exitCode, stdoutStr, stderrStr) ->
      case exitCode of
        ExitSuccess -> pure (Right (Text.pack stdoutStr))
        ExitFailure code ->
          pure (Left ("hledger failed (" <> tshow code <> "): " <> Text.pack stderrStr))

allHledgerTools :: [Engine.Tool]
allHledgerTools =
  [ hledgerBalanceTool,
    hledgerRegisterTool,
    hledgerAddTool,
    hledgerIncomeStatementTool,
    hledgerBalanceSheetTool
  ]

queryBalance :: Maybe Text -> Maybe Text -> Maybe Text -> IO (Either Text Text)
queryBalance maybePattern maybePeriod maybeCurrency = do
  let patternArg = maybe [] (\p -> [Text.unpack p]) maybePattern
      periodArg = maybe [] (\p -> ["-p", "'" ++ Text.unpack p ++ "'"]) maybePeriod
      currency = maybe "USD" Text.unpack maybeCurrency
      currencyArg = ["-X", currency]
  runHledgerInFund (["bal", "-1", "--flat"] <> currencyArg <> patternArg <> periodArg)

queryRegister :: Text -> Maybe Int -> Maybe Text -> Maybe Text -> IO (Either Text Text)
queryRegister accountPattern maybeLimit maybeCurrency maybePeriod = do
  let limitArg = maybe ["-n", "10"] (\n -> ["-n", show n]) maybeLimit
      currency = maybe "USD" Text.unpack maybeCurrency
      currencyArg = ["-X", currency]
      periodArg = maybe [] (\p -> ["-p", "'" ++ Text.unpack p ++ "'"]) maybePeriod
  runHledgerInFund (["reg", Text.unpack accountPattern] <> currencyArg <> periodArg <> limitArg)

incomeStatement :: Maybe Text -> Maybe Text -> IO (Either Text Text)
incomeStatement maybePeriod maybeCurrency = do
  let periodArg = maybe ["-p", "thismonth"] (\p -> ["-p", "'" ++ Text.unpack p ++ "'"]) maybePeriod
      currency = maybe "USD" Text.unpack maybeCurrency
      currencyArg = ["-X", currency]
  runHledgerInFund (["is"] <> currencyArg <> periodArg)

balanceSheet :: Maybe Text -> IO (Either Text Text)
balanceSheet maybeCurrency = do
  let currency = maybe "USD" Text.unpack maybeCurrency
      currencyArg = ["-X", currency]
  runHledgerInFund (["bs"] <> currencyArg)

addTransaction :: Text -> Text -> Text -> Text -> Maybe Text -> IO (Either Text Text)
addTransaction description fromAccount toAccount amount maybeDate = do
  now <- getCurrentTime
  tz <- getCurrentTimeZone
  let localTime = utcToLocalTime tz now
      todayStr = formatTime defaultTimeLocale "%Y-%m-%d" localTime
      dateStr = maybe todayStr Text.unpack maybeDate
      transaction =
        Text.unlines
          [ "",
            Text.pack dateStr <> " " <> description,
            "    " <> toAccount <> "    " <> amount,
            "    " <> fromAccount
          ]
  exists <- doesFileExist transactionsFile
  unless exists <| do
    TextIO.writeFile transactionsFile "; Transactions added via Telegram bot\n"
  TextIO.appendFile transactionsFile transaction
  pure (Right ("Transaction added:\n" <> transaction))

hledgerBalanceTool :: Engine.Tool
hledgerBalanceTool =
  Engine.Tool
    { Engine.toolName = "hledger_balance",
      Engine.toolDescription =
        "Query account balances from hledger. "
          <> "Account patterns: 'as' (assets), 'li' (liabilities), 'ex' (expenses), 'in' (income), 'eq' (equity). "
          <> "Can drill down like 'as:me:cash' or 'ex:us:need'. "
          <> "Currency defaults to USD but can be changed (e.g., 'BTC', 'ETH'). "
          <> "Period uses hledger syntax: 'thismonth', 'lastmonth', 'thisyear', '2024', '2024-06', "
          <> "'from 2024-01-01 to 2024-06-30', 'from 2024-06-01'.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "account_pattern"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Account pattern to filter (e.g., 'as:me:cash', 'ex', 'li:us:cred')" :: Text)
                      ],
                  "period"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("hledger period: 'thismonth', 'lastmonth', '2024', '2024-06', 'from 2024-01-01 to 2024-06-30'" :: Text)
                      ],
                  "currency"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Currency to display values in (default: 'USD'). Examples: 'BTC', 'ETH', 'EUR'" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeBalance
    }

executeBalance :: Aeson.Value -> IO Aeson.Value
executeBalance v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: BalanceArgs) -> do
      result <- queryBalance (baPattern args) (baPeriod args) (baCurrency args)
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right output ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "balances" .= output
                ]
            )

data BalanceArgs = BalanceArgs
  { baPattern :: Maybe Text,
    baPeriod :: Maybe Text,
    baCurrency :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON BalanceArgs where
  parseJSON =
    Aeson.withObject "BalanceArgs" <| \v ->
      (BalanceArgs </ (v .:? "account_pattern"))
        <*> (v .:? "period")
        <*> (v .:? "currency")

hledgerRegisterTool :: Engine.Tool
hledgerRegisterTool =
  Engine.Tool
    { Engine.toolName = "hledger_register",
      Engine.toolDescription =
        "Show recent transactions for an account. "
          <> "Useful for seeing transaction history and checking recent spending. "
          <> "Currency defaults to USD.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "account_pattern"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Account pattern to show transactions for (e.g., 'ex:us:need:grocery')" :: Text)
                      ],
                  "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Max transactions to show (default: 10)" :: Text)
                      ],
                  "currency"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Currency to display values in (default: 'USD')" :: Text)
                      ],
                  "period"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("hledger period: 'thismonth', 'lastmonth', '2024', '2024-06', 'from 2024-06-01 to 2024-12-31'" :: Text)
                      ]
                ],
            "required" .= (["account_pattern"] :: [Text])
          ],
      Engine.toolExecute = executeRegister
    }

executeRegister :: Aeson.Value -> IO Aeson.Value
executeRegister v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: RegisterArgs) -> do
      result <- queryRegister (raPattern args) (raLimit args) (raCurrency args) (raPeriod args)
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right output ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "transactions" .= output
                ]
            )

data RegisterArgs = RegisterArgs
  { raPattern :: Text,
    raLimit :: Maybe Int,
    raCurrency :: Maybe Text,
    raPeriod :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON RegisterArgs where
  parseJSON =
    Aeson.withObject "RegisterArgs" <| \v ->
      (RegisterArgs </ (v .: "account_pattern"))
        <*> (v .:? "limit")
        <*> (v .:? "currency")
        <*> (v .:? "period")

hledgerAddTool :: Engine.Tool
hledgerAddTool =
  Engine.Tool
    { Engine.toolName = "hledger_add",
      Engine.toolDescription =
        "Add a new transaction to the ledger. "
          <> "Use for recording expenses like 'I spent $30 at the barber'. "
          <> "Account naming: ex:me:want (personal discretionary), ex:us:need (shared necessities), "
          <> "as:me:cash:checking (bank account), li:us:cred:chase (credit card). "
          <> "Common expense accounts: ex:us:need:grocery, ex:us:need:utilities, ex:me:want:dining, ex:me:want:grooming.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "description"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Transaction description (e.g., 'Haircut at Joe's Barber')" :: Text)
                      ],
                  "from_account"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Account paying (e.g., 'as:me:cash:checking', 'li:us:cred:chase')" :: Text)
                      ],
                  "to_account"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Account receiving (e.g., 'ex:me:want:grooming')" :: Text)
                      ],
                  "amount"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Amount with currency (e.g., '$30.00', '30 USD')" :: Text)
                      ],
                  "date"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Transaction date YYYY-MM-DD (default: today)" :: Text)
                      ]
                ],
            "required" .= (["description", "from_account", "to_account", "amount"] :: [Text])
          ],
      Engine.toolExecute = executeAdd
    }

executeAdd :: Aeson.Value -> IO Aeson.Value
executeAdd v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: AddArgs) -> do
      result <-
        addTransaction
          (aaDescription args)
          (aaFromAccount args)
          (aaToAccount args)
          (aaAmount args)
          (aaDate args)
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right msg ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message" .= msg
                ]
            )

data AddArgs = AddArgs
  { aaDescription :: Text,
    aaFromAccount :: Text,
    aaToAccount :: Text,
    aaAmount :: Text,
    aaDate :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON AddArgs where
  parseJSON =
    Aeson.withObject "AddArgs" <| \v ->
      (AddArgs </ (v .: "description"))
        <*> (v .: "from_account")
        <*> (v .: "to_account")
        <*> (v .: "amount")
        <*> (v .:? "date")

hledgerIncomeStatementTool :: Engine.Tool
hledgerIncomeStatementTool =
  Engine.Tool
    { Engine.toolName = "hledger_income_statement",
      Engine.toolDescription =
        "Show income statement (income vs expenses) for a period. "
          <> "Good for seeing 'how much did I spend this month' or 'what's my net income'. "
          <> "Currency defaults to USD. "
          <> "Period uses hledger syntax: 'thismonth', 'lastmonth', '2024', 'from 2024-01-01 to 2024-06-30'.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "period"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("hledger period (default: 'thismonth'): 'lastmonth', '2024', '2024-06', 'from 2024-01-01 to 2024-06-30'" :: Text)
                      ],
                  "currency"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Currency to display values in (default: 'USD')" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeIncomeStatement
    }

executeIncomeStatement :: Aeson.Value -> IO Aeson.Value
executeIncomeStatement v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: IncomeStatementArgs) -> do
      result <- incomeStatement (isaPeriod args) (isaCurrency args)
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right output ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "income_statement" .= output
                ]
            )

data IncomeStatementArgs = IncomeStatementArgs
  { isaPeriod :: Maybe Text,
    isaCurrency :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON IncomeStatementArgs where
  parseJSON =
    Aeson.withObject "IncomeStatementArgs" <| \v ->
      (IncomeStatementArgs </ (v .:? "period"))
        <*> (v .:? "currency")

hledgerBalanceSheetTool :: Engine.Tool
hledgerBalanceSheetTool =
  Engine.Tool
    { Engine.toolName = "hledger_balance_sheet",
      Engine.toolDescription =
        "Show current balance sheet (assets, liabilities, net worth). "
          <> "Good for seeing 'what's my net worth' or 'how much do I have'. "
          <> "Currency defaults to USD.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "currency"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Currency to display values in (default: 'USD')" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeBalanceSheet
    }

executeBalanceSheet :: Aeson.Value -> IO Aeson.Value
executeBalanceSheet v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: BalanceSheetArgs) -> do
      result <- balanceSheet (bsCurrency args)
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right output ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "balance_sheet" .= output
                ]
            )

newtype BalanceSheetArgs = BalanceSheetArgs
  { bsCurrency :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON BalanceSheetArgs where
  parseJSON =
    Aeson.withObject "BalanceSheetArgs" <| \v ->
      BalanceSheetArgs </ (v .:? "currency")