{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}

module Haskoin.Store.Logic
  ( ImportException (..),
    MonadImport,
    initBest,
    revertBlock,
    importBlock,
    newMempoolTx,
    deleteUnconfirmedTx,
  )
where

import Control.Monad
  ( forM,
    forM_,
    guard,
    unless,
    void,
    when,
    zipWithM_,
    (<=<),
  )
import Control.Monad.Except (MonadError, throwError)
import Control.Monad.Logger
  ( MonadLoggerIO (..),
    logDebugS,
    logErrorS,
  )
import qualified Data.ByteString as B
import Data.Either (rights)
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import qualified Data.IntMap.Strict as I
import Data.List (nub)
import Data.Maybe
  ( catMaybes,
    fromMaybe,
    isJust,
    isNothing,
  )
import Data.Serialize (encode)
import Data.String.Conversions (cs)
import Data.Word (Word32, Word64)
import Haskoin
  ( Address,
    Block (..),
    BlockHash,
    BlockHeader (..),
    BlockNode (..),
    Network (..),
    OutPoint (..),
    Tx (..),
    TxHash,
    TxIn (..),
    TxOut (..),
    blockHashToHex,
    computeSubsidy,
    eitherToMaybe,
    genesisBlock,
    genesisNode,
    headerHash,
    isGenesis,
    nullOutPoint,
    scriptToAddressBS,
    txHash,
    txHashToHex,
  )
import Haskoin.Store.Common
import Haskoin.Store.Data
  ( Balance (..),
    BlockData (..),
    BlockRef (..),
    Prev (..),
    Spender (..),
    TxData (..),
    TxRef (..),
    UnixTime,
    Unspent (..),
    confirmed,
    isCoinbaseTx,
    txDataFee,
  )
import UnliftIO (Exception)

type MonadImport m =
  ( MonadError ImportException m,
    MonadLoggerIO m,
    StoreReadBase m,
    StoreWrite m
  )

data ImportException
  = PrevBlockNotBest
  | Orphan
  | UnexpectedCoinbase
  | BestBlockNotFound
  | BlockNotBest
  | TxNotFound
  | DoubleSpend
  | TxConfirmed
  | InsufficientFunds
  | DuplicatePrevOutput
  | TxSpent
  deriving (Eq, Ord, Exception)

instance Show ImportException where
  show PrevBlockNotBest = "Previous block not best"
  show Orphan = "Orphan"
  show UnexpectedCoinbase = "Unexpected coinbase"
  show BestBlockNotFound = "Best block not found"
  show BlockNotBest = "Block not best"
  show TxNotFound = "Transaction not found"
  show DoubleSpend = "Double spend"
  show TxConfirmed = "Transaction confirmed"
  show InsufficientFunds = "Insufficient funds"
  show DuplicatePrevOutput = "Duplicate previous output"
  show TxSpent = "Transaction is spent"

initBest :: MonadImport m => m ()
initBest = do
  $(logDebugS) "BlockStore" "Initializing best block"
  net <- getNetwork
  m <- getBestBlock
  when (isNothing m) . void $ do
    $(logDebugS) "BlockStore" "Importing Genesis block"
    importBlock (genesisBlock net) (genesisNode net)

newMempoolTx :: MonadImport m => Tx -> UnixTime -> m Bool
newMempoolTx tx w =
  getActiveTxData (txHash tx) >>= \case
    Just _ ->
      return False
    Nothing -> do
      freeOutputs True True tx
      rbf <- isRBF (MemRef w) tx
      checkNewTx tx
      importTx (MemRef w) w rbf tx
      return True

bestBlockData :: MonadImport m => m BlockData
bestBlockData = do
  h <-
    getBestBlock >>= \case
      Nothing -> do
        $(logErrorS) "BlockStore" "Best block unknown"
        throwError BestBlockNotFound
      Just h -> return h
  getBlock h >>= \case
    Nothing -> do
      $(logErrorS) "BlockStore" "Best block not found"
      throwError BestBlockNotFound
    Just b -> return b

revertBlock :: MonadImport m => BlockHash -> m ()
revertBlock bh = do
  bd <-
    bestBlockData >>= \b ->
      if headerHash (blockDataHeader b) == bh
        then return b
        else do
          $(logErrorS) "BlockStore" $
            "Cannot revert non-head block: " <> blockHashToHex bh
          throwError BlockNotBest
  $(logDebugS) "BlockStore" $
    "Obtained block data for " <> blockHashToHex bh
  tds <- mapM getImportTxData (blockDataTxs bd)
  $(logDebugS) "BlockStore" $
    "Obtained import tx data for block " <> blockHashToHex bh
  setBest (prevBlock (blockDataHeader bd))
  $(logDebugS) "BlockStore" $
    "Set parent as best block "
      <> blockHashToHex (prevBlock (blockDataHeader bd))
  insertBlock bd {blockDataMainChain = False}
  $(logDebugS) "BlockStore" $
    "Updated as not in main chain: " <> blockHashToHex bh
  forM_ (tail tds) unConfirmTx
  $(logDebugS) "BlockStore" $
    "Unconfirmed " <> cs (show (length tds)) <> " transactions"
  deleteConfirmedTx (txHash (txData (head tds)))
  $(logDebugS) "BlockStore" $
    "Deleted coinbase: " <> txHashToHex (txHash (txData (head tds)))

checkNewBlock :: MonadImport m => Block -> BlockNode -> m ()
checkNewBlock b n =
  getBestBlock >>= \case
    Nothing
      | isGenesis n -> return ()
      | otherwise -> do
        $(logErrorS) "BlockStore" $
          "Cannot import non-genesis block: "
            <> blockHashToHex (headerHash (blockHeader b))
        throwError BestBlockNotFound
    Just h
      | prevBlock (blockHeader b) == h -> return ()
      | otherwise -> do
        $(logErrorS) "BlockStore" $
          "Block does not build on head: "
            <> blockHashToHex (headerHash (blockHeader b))
        throwError PrevBlockNotBest

importOrConfirm :: MonadImport m => BlockNode -> [Tx] -> m [TxData]
importOrConfirm bn txns = do
  mapM_ (freeOutputs True False . snd) (reverse txs)
  mapM (uncurry action) txs
  where
    txs = sortTxs txns
    br i = BlockRef {blockRefHeight = nodeHeight bn, blockRefPos = i}
    bn_time = fromIntegral . blockTimestamp $ nodeHeader bn
    action i tx =
      testPresent tx >>= \case
        False -> import_it i tx
        True -> confirm_it i tx
    confirm_it i tx =
      getActiveTxData (txHash tx) >>= \case
        Just t -> do
          $(logDebugS) "BlockStore" $
            "Confirming tx: "
              <> txHashToHex (txHash tx)
          confirmTx t (br i)
        Nothing -> do
          $(logErrorS) "BlockStore" $
            "Cannot find tx to confirm: "
              <> txHashToHex (txHash tx)
          throwError TxNotFound
    import_it i tx = do
      $(logDebugS) "BlockStore" $
        "Importing tx: " <> txHashToHex (txHash tx)
      importTx (br i) bn_time False tx

importBlock :: MonadImport m => Block -> BlockNode -> m (BlockData, [TxData])
importBlock b n = do
  $(logDebugS) "BlockStore" $
    "Checking new block: "
      <> blockHashToHex (headerHash (nodeHeader n))
  checkNewBlock b n
  $(logDebugS) "BlockStore" "Passed check"
  net <- getNetwork
  let subsidy = computeSubsidy net (nodeHeight n)
  bs <- getBlocksAtHeight (nodeHeight n)
  $(logDebugS) "BlockStore" $
    "Inserting block entries for: "
      <> blockHashToHex (headerHash (nodeHeader n))
  setBlocksAtHeight
    (nub (headerHash (nodeHeader n) : bs))
    (nodeHeight n)
  setBest (headerHash (nodeHeader n))
  tds <- importOrConfirm n (blockTxns b)
  let bd =
        BlockData
          { blockDataHeight = nodeHeight n,
            blockDataMainChain = True,
            blockDataWork = nodeWork n,
            blockDataHeader = nodeHeader n,
            blockDataSize = fromIntegral (B.length (encode b)),
            blockDataTxs = map txHash (blockTxns b),
            blockDataWeight = if getSegWit net then w else 0,
            blockDataSubsidy = subsidy,
            blockDataFees = sum $ map txDataFee tds,
            blockDataOutputs = ts_out_val
          }
  insertBlock bd
  $(logDebugS) "BlockStore" $
    "Finished importing block: "
      <> blockHashToHex (headerHash (nodeHeader n))
  return (bd, tds)
  where
    cb_out_val =
      sum $ map outValue $ txOut $ head $ blockTxns b
    ts_out_val =
      sum $ map (sum . map outValue . txOut) $ tail $ blockTxns b
    w =
      let f t = t {txWitness = []}
          b' = b {blockTxns = map f (blockTxns b)}
          x = B.length (encode b)
          s = B.length (encode b')
       in fromIntegral $ s * 3 + x

checkNewTx :: MonadImport m => Tx -> m ()
checkNewTx tx = do
  when (unique_inputs < length (txIn tx)) $ do
    $(logErrorS) "BlockStore" $
      "Transaction spends same output twice: "
        <> txHashToHex (txHash tx)
    throwError DuplicatePrevOutput
  us <- getUnspentOutputs tx
  when (any isNothing us) $ do
    $(logErrorS) "BlockStore" $
      "Orphan: " <> txHashToHex (txHash tx)
    throwError Orphan
  when (isCoinbaseTx tx) $ do
    $(logErrorS) "BlockStore" $
      "Coinbase cannot be imported into mempool: "
        <> txHashToHex (txHash tx)
    throwError UnexpectedCoinbase
  when (length (prevOuts tx) > length us) $ do
    $(logErrorS) "BlockStore" $
      "Orphan: " <> txHashToHex (txHash tx)
    throwError Orphan
  when (outputs > unspents us) $ do
    $(logErrorS) "BlockStore" $
      "Insufficient funds for tx: " <> txHashToHex (txHash tx)
    throwError InsufficientFunds
  where
    unspents = sum . map unspentAmount . catMaybes
    outputs = sum (map outValue (txOut tx))
    unique_inputs = length (nub' (map prevOutput (txIn tx)))

getUnspentOutputs :: StoreReadBase m => Tx -> m [Maybe Unspent]
getUnspentOutputs tx = mapM getUnspent (prevOuts tx)

prepareTxData :: Bool -> BlockRef -> Word64 -> Tx -> [Unspent] -> TxData
prepareTxData rbf br tt tx us =
  TxData
    { txDataBlock = br,
      txData = tx,
      txDataPrevs = ps,
      txDataDeleted = False,
      txDataRBF = rbf,
      txDataTime = tt,
      txDataSpenders = I.empty
    }
  where
    mkprv u = Prev (unspentScript u) (unspentAmount u)
    ps = I.fromList $ zip [0 ..] $ map mkprv us

importTx ::
  MonadImport m =>
  BlockRef ->
  -- | unix time
  Word64 ->
  -- | RBF
  Bool ->
  Tx ->
  m TxData
importTx br tt rbf tx = do
  mus <- getUnspentOutputs tx
  us <- forM mus $ \case
    Nothing -> do
      $(logErrorS) "BlockStore" $
        "Attempted to import a tx missing UTXO: "
          <> txHashToHex (txHash tx)
      throwError Orphan
    Just u -> return u
  let td = prepareTxData rbf br tt tx us
  commitAddTx td
  return td

unConfirmTx :: MonadImport m => TxData -> m TxData
unConfirmTx t = confTx t Nothing

confirmTx :: MonadImport m => TxData -> BlockRef -> m TxData
confirmTx t br = confTx t (Just br)

replaceAddressTx :: MonadImport m => TxData -> BlockRef -> m ()
replaceAddressTx t new = forM_ (txDataAddresses t) $ \a -> do
  deleteAddrTx
    a
    TxRef
      { txRefBlock = txDataBlock t,
        txRefHash = txHash (txData t)
      }
  insertAddrTx
    a
    TxRef
      { txRefBlock = new,
        txRefHash = txHash (txData t)
      }

adjustAddressOutput ::
  MonadImport m =>
  OutPoint ->
  TxOut ->
  BlockRef ->
  BlockRef ->
  m ()
adjustAddressOutput op o old new = do
  let pk = scriptOutput o
  getUnspent op >>= \case
    Nothing -> return ()
    Just u -> do
      unless (unspentBlock u == old) $
        error $ "Existing unspent block bad for output: " <> show op
      replace_unspent pk
  where
    replace_unspent pk = do
      let ma = eitherToMaybe (scriptToAddressBS pk)
      deleteUnspent op
      insertUnspent
        Unspent
          { unspentBlock = new,
            unspentPoint = op,
            unspentAmount = outValue o,
            unspentScript = pk,
            unspentAddress = ma
          }
      forM_ ma $ replace_addr_unspent pk
    replace_addr_unspent pk a = do
      deleteAddrUnspent
        a
        Unspent
          { unspentBlock = old,
            unspentPoint = op,
            unspentAmount = outValue o,
            unspentScript = pk,
            unspentAddress = Just a
          }
      insertAddrUnspent
        a
        Unspent
          { unspentBlock = new,
            unspentPoint = op,
            unspentAmount = outValue o,
            unspentScript = pk,
            unspentAddress = Just a
          }
      decreaseBalance (confirmed old) a (outValue o)
      increaseBalance (confirmed new) a (outValue o)

confTx :: MonadImport m => TxData -> Maybe BlockRef -> m TxData
confTx t mbr = do
  replaceAddressTx t new
  forM_ (zip [0 ..] (txOut (txData t))) $ \(n, o) -> do
    let op = OutPoint (txHash (txData t)) n
    adjustAddressOutput op o old new
  rbf <- isRBF new (txData t)
  let td = t {txDataBlock = new, txDataRBF = rbf}
  insertTx td
  updateMempool td
  return td
  where
    new = fromMaybe (MemRef (txDataTime t)) mbr
    old = txDataBlock t

freeOutputs ::
  MonadImport m =>
  -- | only delete transaction if unconfirmed
  Bool ->
  -- | only delete RBF
  Bool ->
  Tx ->
  m ()
freeOutputs memonly rbfcheck tx = do
  let prevs = prevOuts tx
  unspents <- mapM getUnspent prevs
  let spents = [p | (p, Nothing) <- zip prevs unspents]
  spndrs <- catMaybes <$> mapM getSpender spents
  let txids = HashSet.fromList $ filter (/= txHash tx) $ map spenderHash spndrs
  mapM_ (deleteTx memonly rbfcheck) $ HashSet.toList txids

deleteConfirmedTx :: MonadImport m => TxHash -> m ()
deleteConfirmedTx = deleteTx False False

deleteUnconfirmedTx :: MonadImport m => Bool -> TxHash -> m ()
deleteUnconfirmedTx rbfcheck th =
  getActiveTxData th >>= \case
    Just _ -> deleteTx True rbfcheck th
    Nothing ->
      $(logDebugS) "BlockStore" $
        "Not found or already deleted: " <> txHashToHex th

deleteTx ::
  MonadImport m =>
  -- | only delete transaction if unconfirmed
  Bool ->
  -- | only delete RBF
  Bool ->
  TxHash ->
  m ()
deleteTx memonly rbfcheck th = do
  chain <- getChain memonly rbfcheck th
  $(logDebugS) "BlockStore" $
    "Deleting " <> cs (show (length chain))
      <> " txs from chain leading to "
      <> txHashToHex th
  mapM_ (\t -> let h = txHash t in deleteSingleTx h >> return h) chain

getChain ::
  (MonadImport m, MonadLoggerIO m) =>
  -- | only delete transaction if unconfirmed
  Bool ->
  -- | only delete RBF
  Bool ->
  TxHash ->
  m [Tx]
getChain memonly rbfcheck th' = do
  $(logDebugS) "BlockStore" $
    "Getting chain for tx " <> txHashToHex th'
  sort_clean <$> go HashSet.empty (HashSet.singleton th')
  where
    sort_clean = reverse . map snd . sortTxs
    get_tx th =
      getActiveTxData th >>= \case
        Nothing -> do
          $(logDebugS) "BlockStore" $
            "Transaction not found: " <> txHashToHex th
          return Nothing
        Just td
          | memonly && confirmed (txDataBlock td) -> do
            $(logErrorS) "BlockStore" $
              "Transaction already confirmed: "
                <> txHashToHex th
            throwError TxConfirmed
          | rbfcheck ->
            isRBF (txDataBlock td) (txData td) >>= \case
              True -> return $ Just td
              False -> do
                $(logErrorS) "BlockStore" $
                  "Double-spending transaction: "
                    <> txHashToHex th
                throwError DoubleSpend
          | otherwise -> return $ Just td
    go txs pdg = do
      tds <- catMaybes <$> mapM get_tx (HashSet.toList pdg)
      let txsn = HashSet.fromList $ fmap txData tds
          pdgn =
            HashSet.fromList
              . concatMap (map spenderHash . I.elems)
              $ fmap txDataSpenders tds
          txs' = txsn <> txs
          pdg' = pdgn `HashSet.difference` HashSet.map txHash txs'
      if HashSet.null pdg'
        then return $ HashSet.toList txs'
        else go txs' pdg'

deleteSingleTx :: MonadImport m => TxHash -> m ()
deleteSingleTx th =
  getActiveTxData th >>= \case
    Nothing -> do
      $(logErrorS) "BlockStore" $
        "Already deleted: " <> txHashToHex th
      throwError TxNotFound
    Just td ->
      if I.null (txDataSpenders td)
        then do
          $(logDebugS) "BlockStore" $
            "Deleting tx: " <> txHashToHex th
          commitDelTx td
        else do
          $(logErrorS) "BlockStore" $
            "Tried to delete spent tx: "
              <> txHashToHex th
          throwError TxSpent

commitDelTx :: MonadImport m => TxData -> m ()
commitDelTx = commitModTx False

commitAddTx :: MonadImport m => TxData -> m ()
commitAddTx = commitModTx True

commitModTx :: MonadImport m => Bool -> TxData -> m ()
commitModTx add tx_data = do
  mapM_ mod_addr_tx (txDataAddresses td)
  mod_outputs
  mod_unspent
  insertTx td
  updateMempool td
  where
    tx = txData td
    br = txDataBlock td
    td = tx_data {txDataDeleted = not add}
    tx_ref = TxRef br (txHash tx)
    mod_addr_tx a
      | add = do
        insertAddrTx a tx_ref
        modAddressCount add a
      | otherwise = do
        deleteAddrTx a tx_ref
        modAddressCount add a
    mod_unspent
      | add = spendOutputs tx
      | otherwise = unspendOutputs tx
    mod_outputs
      | add = addOutputs br tx
      | otherwise = delOutputs br tx

updateMempool :: MonadImport m => TxData -> m ()
updateMempool td@TxData {txDataDeleted = True} =
  deleteFromMempool (txHash (txData td))
updateMempool td@TxData {txDataBlock = MemRef t} =
  addToMempool (txHash (txData td)) t
updateMempool td@TxData {txDataBlock = BlockRef {}} =
  deleteFromMempool (txHash (txData td))

spendOutputs :: MonadImport m => Tx -> m ()
spendOutputs tx =
  zipWithM_ (spendOutput (txHash tx)) [0 ..] (prevOuts tx)

addOutputs :: MonadImport m => BlockRef -> Tx -> m ()
addOutputs br tx =
  zipWithM_ (addOutput br . OutPoint (txHash tx)) [0 ..] (txOut tx)

isRBF ::
  StoreReadBase m =>
  BlockRef ->
  Tx ->
  m Bool
isRBF br tx
  | confirmed br = return False
  | otherwise =
    getNetwork >>= \net ->
      if getReplaceByFee net
        then go
        else return False
  where
    go
      | any ((< 0xffffffff - 1) . txInSequence) (txIn tx) = return True
      | otherwise = carry_on
    carry_on =
      let hs = nub' $ map (outPointHash . prevOutput) (txIn tx)
          ck [] = return False
          ck (h : hs') =
            getActiveTxData h >>= \case
              Nothing -> return False
              Just t
                | confirmed (txDataBlock t) -> ck hs'
                | txDataRBF t -> return True
                | otherwise -> ck hs'
       in ck hs

addOutput :: MonadImport m => BlockRef -> OutPoint -> TxOut -> m ()
addOutput = modOutput True

delOutput :: MonadImport m => BlockRef -> OutPoint -> TxOut -> m ()
delOutput = modOutput False

modOutput :: MonadImport m => Bool -> BlockRef -> OutPoint -> TxOut -> m ()
modOutput add br op o = do
  mod_unspent
  forM_ ma $ \a -> do
    mod_addr_unspent a u
    modBalance (confirmed br) add a (outValue o)
    modifyReceived a v
  where
    v
      | add = (+ outValue o)
      | otherwise = subtract (outValue o)
    ma = eitherToMaybe (scriptToAddressBS (scriptOutput o))
    u =
      Unspent
        { unspentScript = scriptOutput o,
          unspentBlock = br,
          unspentPoint = op,
          unspentAmount = outValue o,
          unspentAddress = ma
        }
    mod_unspent
      | add = insertUnspent u
      | otherwise = deleteUnspent op
    mod_addr_unspent
      | add = insertAddrUnspent
      | otherwise = deleteAddrUnspent

delOutputs :: MonadImport m => BlockRef -> Tx -> m ()
delOutputs br tx =
  forM_ (zip [0 ..] (txOut tx)) $ \(i, o) -> do
    let op = OutPoint (txHash tx) i
    delOutput br op o

getImportTxData :: MonadImport m => TxHash -> m TxData
getImportTxData th =
  getActiveTxData th >>= \case
    Nothing -> do
      $(logDebugS) "BlockStore" $ "Tx not found: " <> txHashToHex th
      throwError TxNotFound
    Just d -> return d

getTxOut :: Word32 -> Tx -> Maybe TxOut
getTxOut i tx = do
  guard (fromIntegral i < length (txOut tx))
  return $ txOut tx !! fromIntegral i

insertSpender :: MonadImport m => OutPoint -> Spender -> m ()
insertSpender op s = do
  td <- getImportTxData (outPointHash op)
  let p = txDataSpenders td
      p' = I.insert (fromIntegral (outPointIndex op)) s p
      td' = td {txDataSpenders = p'}
  insertTx td'

deleteSpender :: MonadImport m => OutPoint -> m ()
deleteSpender op = do
  td <- getImportTxData (outPointHash op)
  let p = txDataSpenders td
      p' = I.delete (fromIntegral (outPointIndex op)) p
      td' = td {txDataSpenders = p'}
  insertTx td'

spendOutput :: MonadImport m => TxHash -> Word32 -> OutPoint -> m ()
spendOutput th ix op = do
  u <-
    getUnspent op >>= \case
      Just u -> return u
      Nothing -> error $ "Could not find UTXO to spend: " <> show op
  deleteUnspent op
  insertSpender op (Spender th ix)
  let pk = unspentScript u
  forM_ (scriptToAddressBS pk) $ \a -> do
    decreaseBalance
      (confirmed (unspentBlock u))
      a
      (unspentAmount u)
    deleteAddrUnspent a u

unspendOutputs :: MonadImport m => Tx -> m ()
unspendOutputs = mapM_ unspendOutput . prevOuts

unspendOutput :: MonadImport m => OutPoint -> m ()
unspendOutput op = do
  t <-
    getActiveTxData (outPointHash op) >>= \case
      Nothing -> error $ "Could not find tx data: " <> show (outPointHash op)
      Just t -> return t
  let o =
        fromMaybe
          (error ("Could not find output: " <> show op))
          (getTxOut (outPointIndex op) (txData t))
      m = eitherToMaybe (scriptToAddressBS (scriptOutput o))
      u =
        Unspent
          { unspentAmount = outValue o,
            unspentBlock = txDataBlock t,
            unspentScript = scriptOutput o,
            unspentPoint = op,
            unspentAddress = m
          }
  deleteSpender op
  insertUnspent u
  forM_ m $ \a -> do
    insertAddrUnspent a u
    increaseBalance (confirmed (unspentBlock u)) a (outValue o)

modifyReceived :: MonadImport m => Address -> (Word64 -> Word64) -> m ()
modifyReceived a f = do
  b <- getDefaultBalance a
  setBalance b {balanceTotalReceived = f (balanceTotalReceived b)}

decreaseBalance :: MonadImport m => Bool -> Address -> Word64 -> m ()
decreaseBalance conf = modBalance conf False

increaseBalance :: MonadImport m => Bool -> Address -> Word64 -> m ()
increaseBalance conf = modBalance conf True

modBalance ::
  MonadImport m =>
  -- | confirmed
  Bool ->
  -- | add
  Bool ->
  Address ->
  Word64 ->
  m ()
modBalance conf add a val = do
  b <- getDefaultBalance a
  setBalance $ (g . f) b
  where
    g b = b {balanceUnspentCount = m 1 (balanceUnspentCount b)}
    f b
      | conf = b {balanceAmount = m val (balanceAmount b)}
      | otherwise = b {balanceZero = m val (balanceZero b)}
    m
      | add = (+)
      | otherwise = subtract

modAddressCount :: MonadImport m => Bool -> Address -> m ()
modAddressCount add a = do
  b <- getDefaultBalance a
  setBalance b {balanceTxCount = f (balanceTxCount b)}
  where
    f
      | add = (+ 1)
      | otherwise = subtract 1

txOutAddrs :: [TxOut] -> [Address]
txOutAddrs = nub' . rights . map (scriptToAddressBS . scriptOutput)

txInAddrs :: [Prev] -> [Address]
txInAddrs = nub' . rights . map (scriptToAddressBS . prevScript)

txDataAddresses :: TxData -> [Address]
txDataAddresses t =
  nub' $ txInAddrs prevs <> txOutAddrs outs
  where
    prevs = I.elems (txDataPrevs t)
    outs = txOut (txData t)

prevOuts :: Tx -> [OutPoint]
prevOuts tx = filter (/= nullOutPoint) (map prevOutput (txIn tx))

testPresent :: StoreReadBase m => Tx -> m Bool
testPresent tx = isJust <$> getActiveTxData (txHash tx)