{-# LANGUAGE OverloadedStrings #-}

module Zenith.Scanner where

import Control.Concurrent.Async (concurrently_, withAsync)
import Control.Exception (throwIO, try)
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Logger
  ( NoLoggingT
  , logErrorN
  , logInfoN
  , runNoLoggingT
  , runStderrLoggingT
  )
import Data.Aeson
import Data.HexString
import qualified Data.Text as T
import Data.Time (getCurrentTime)
import Database.Persist.Sqlite
import System.Console.AsciiProgress
import ZcashHaskell.Types
  ( BlockResponse(..)
  , RawZebraTx(..)
  , Transaction(..)
  , ZcashNet(..)
  , ZebraGetBlockChainInfo(..)
  , ZebraTxResponse(..)
  , fromRawOBundle
  , fromRawSBundle
  , fromRawTBundle
  )
import ZcashHaskell.Utils (getBlockTime, makeZebraCall, readZebraTransaction)
import Zenith.Core (checkBlockChain, syncWallet, updateCommitmentTrees)
import Zenith.DB
  ( ZcashBlock(..)
  , ZcashBlockId
  , clearWalletData
  , clearWalletTransactions
  , completeSync
  , getBlock
  , getMaxBlock
  , getMinBirthdayHeight
  , getUnconfirmedBlocks
  , getWallets
  , initDb
  , initPool
  , saveBlock
  , saveConfs
  , saveTransaction
  , startSync
  , updateWalletSync
  , upgradeAccountTable
  , upgradeQrTable
  )
import Zenith.Types
  ( Config(..)
  , HexStringDB(..)
  , ZcashNetDB(..)
  , ZenithStatus(..)
  )
import Zenith.Utils (jsonNumber)

-- | Function to scan the Zcash blockchain through the Zebra node and populate the Zenith database
rescanZebra ::
     T.Text -- ^ Host
  -> Int -- ^ Port
  -> T.Text -- ^ Path to database file
  -> IO ()
rescanZebra host port dbFilePath = do
  bc <-
    try $ checkBlockChain host port :: IO
      (Either IOError ZebraGetBlockChainInfo)
  case bc of
    Left e -> print e
    Right bStatus -> do
      let znet = ZcashNetDB $ zgb_net bStatus
      pool1 <- runNoLoggingT $ initPool dbFilePath
      {-pool2 <- runNoLoggingT $ initPool dbFilePath-}
      {-pool3 <- runNoLoggingT $ initPool dbFilePath-}
      _ <- runNoLoggingT $ initDb dbFilePath
      upgradeQrTable pool1
      clearWalletTransactions pool1
      clearWalletData pool1
      _ <- startSync pool1
      dbBlock <- getMaxBlock pool1 znet
      b <- liftIO $ getMinBirthdayHeight pool1 znet
      let sb = max dbBlock b
      if sb > zgb_blocks bStatus || sb < 1
        then liftIO $ throwIO $ userError "Invalid starting block for scan"
        else do
          print $
            "Scanning from " ++ show sb ++ " to " ++ show (zgb_blocks bStatus)
          let bList = [sb .. (zgb_blocks bStatus)]
          {-
          let batch = length bList `div` 3
          let bl1 = take batch bList
          let bl2 = take batch $ drop batch bList
          let bl3 = drop (2 * batch) bList
          -}
          _ <-
            displayConsoleRegions $ do
              pg1 <- newProgressBar def {pgTotal = fromIntegral $ length bList}
              {-pg2 <- newProgressBar def {pgTotal = fromIntegral $ length bl2}-}
              {-pg3 <- newProgressBar def {pgTotal = fromIntegral $ length bl3}-}
              mapM_ (processBlock host port pool1 pg1 znet) bList
              {-`concurrently_`-}
                {-mapM_ (processBlock host port pool2 pg2 znet) bl2 `concurrently_`-}
                {-mapM_ (processBlock host port pool3 pg3 znet) bl3-}
              print "Please wait..."
          _ <- completeSync pool1 Successful
          _ <- runNoLoggingT $ updateCommitmentTrees pool1 host port znet
          print "Rescan complete"

-- | Function to process a raw block and extract the transaction information
processBlock ::
     T.Text -- ^ Host name for `zebrad`
  -> Int -- ^ Port for `zebrad`
  -> ConnectionPool -- ^ DB file path
  -> ProgressBar -- ^ Progress bar
  -> ZcashNetDB -- ^ the network
  -> Int -- ^ The block number to process
  -> IO ()
processBlock host port pool pg net b = do
  r <-
    liftIO $
    makeZebraCall
      host
      port
      "getblock"
      [Data.Aeson.String $ T.pack $ show b, jsonNumber 1]
  case r of
    Left e -> do
      _ <- completeSync pool Failed
      liftIO $ throwIO $ userError e
    Right blk -> do
      bi <-
        saveBlock pool $
        ZcashBlock
          (fromIntegral $ bl_height blk)
          (HexStringDB $ bl_hash blk)
          (fromIntegral $ bl_confirmations blk)
          (fromIntegral $ bl_time blk)
          net
      mapM_ (processTx host port bi pool) $ bl_txs blk
      liftIO $ tick pg

-- | Function to process a raw transaction
processTx ::
     T.Text -- ^ Host name for `zebrad`
  -> Int -- ^ Port for `zebrad`
  -> ZcashBlockId -- ^ Block ID
  -> ConnectionPool -- ^ DB file path
  -> HexString -- ^ transaction id
  -> IO ()
processTx host port bt pool t = do
  r <-
    liftIO $
    makeZebraCall
      host
      port
      "getrawtransaction"
      [Data.Aeson.String $ toText t, jsonNumber 1]
  case r of
    Left e -> do
      _ <- completeSync pool Failed
      liftIO $ throwIO $ userError e
    Right rawTx -> do
      case readZebraTransaction (ztr_hex rawTx) of
        Nothing -> return ()
        Just rzt -> do
          _ <-
            runNoLoggingT $
            saveTransaction pool bt $
            Transaction
              t
              (ztr_blockheight rawTx)
              (ztr_conf rawTx)
              (fromIntegral $ zt_expiry rzt)
              (fromRawTBundle $ zt_tBundle rzt)
              (fromRawSBundle $ zt_sBundle rzt)
              (fromRawOBundle $ zt_oBundle rzt)
          return ()

-- | Function to update unconfirmed transactions
updateConfs ::
     T.Text -- ^ Host name for `zebrad`
  -> Int -- ^ Port for `zebrad`
  -> ConnectionPool
  -> IO ()
updateConfs host port pool = do
  targetBlocks <- getUnconfirmedBlocks pool
  mapM_ updateTx targetBlocks
  where
    updateTx :: Int -> IO ()
    updateTx b = do
      r <-
        makeZebraCall
          host
          port
          "getblock"
          [Data.Aeson.String $ T.pack $ show b, jsonNumber 1]
      case r of
        Left e -> throwIO $ userError e
        Right blk -> do
          saveConfs pool b $ fromInteger $ bl_confirmations blk

clearSync :: Config -> IO ()
clearSync config = do
  let zHost = c_zebraHost config
  let zPort = c_zebraPort config
  let dbPath = c_dbPath config
  pool <- runNoLoggingT $ initPool dbPath
  bc <-
    try $ checkBlockChain zHost zPort :: IO
      (Either IOError ZebraGetBlockChainInfo)
  case bc of
    Left e1 -> throwIO e1
    Right chainInfo -> do
      x <- runNoLoggingT $ initDb dbPath
      _ <- upgradeQrTable pool
      _ <- runNoLoggingT $ upgradeAccountTable pool
      case x of
        Left e2 -> throwIO $ userError e2
        Right x' -> do
          when x' $ rescanZebra zHost zPort dbPath
          _ <- clearWalletTransactions pool
          w <- getWallets pool $ zgb_net chainInfo
          liftIO $ mapM_ (updateWalletSync pool 0 . entityKey) w
          w' <- liftIO $ getWallets pool $ zgb_net chainInfo
          r <- runNoLoggingT $ mapM (syncWallet config) w'
          liftIO $ print r

-- | Detect chain re-orgs
checkIntegrity ::
     T.Text -- ^ Database path
  -> T.Text -- ^ Zebra host
  -> Int -- ^ Zebra port
  -> ZcashNet -- ^ the network to scan
  -> Int -- ^ The block to start the check
  -> Int -- ^ depth
  -> IO Int
checkIntegrity dbP zHost zPort znet b d =
  if b < 1
    then return 1
    else do
      r <-
        makeZebraCall
          zHost
          zPort
          "getblock"
          [Data.Aeson.String $ T.pack $ show b, jsonNumber 1]
      case r of
        Left e -> throwIO $ userError e
        Right blk -> do
          pool <- runNoLoggingT $ initPool dbP
          dbBlk <- getBlock pool b $ ZcashNetDB znet
          case dbBlk of
            Nothing -> return 1
            Just dbBlk' ->
              if bl_hash blk == getHex (zcashBlockHash $ entityVal dbBlk')
                then return b
                else checkIntegrity dbP zHost zPort znet (b - 5 * d) (d + 1)
