content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  
  -- |
  -- Module      : Haskoin.Util
  -- Copyright   : No rights reserved
  -- License     : MIT
  -- Maintainer  : jprupp@protonmail.ch
  -- Stability   : experimental
  -- Portability : POSIX
  --
  -- This module defines various utility functions used across the library.
  module Haskoin.Util.Helpers
    ( -- * ByteString Helpers
      bsToInteger,
      integerToBS,
      hexEncoding,
      hexBuilder,
      encodeHex,
      encodeHexLazy,
      decodeHex,
      decodeHexLazy,
      getBits,
  
      -- * Maybe & Either Helpers
      eitherToMaybe,
      maybeToEither,
      liftEither,
      liftMaybe,
  
      -- * Other Helpers
      updateIndex,
      matchTemplate,
      convertBits,
  
      -- * Triples
      fst3,
      snd3,
      lst3,
  
      -- * JSON Utilities
      dropFieldLabel,
      dropSumLabels,
  
      -- * Serialization Helpers
      putList,
      getList,
      putMaybe,
      getMaybe,
      putLengthBytes,
      getLengthBytes,
      putInteger,
      getInteger,
      putInt32be,
      getInt32be,
      putInt64be,
      getInt64be,
      getIntMap,
      putIntMap,
      getTwo,
      putTwo,
  
      -- * Test Helpers
      prepareContext,
      customCerealID,
      readTestFile,
      readTestFileParser,
    )
  where
  
  import Control.Monad
  import Control.Monad.Except (ExceptT (..), liftEither)
  import Crypto.Secp256k1
  import Data.Aeson (eitherDecodeFileStrict)
  import Data.Aeson.Encoding
  import Data.Aeson.Types
  import Data.Base16.Types
  import Data.Bits
  import Data.ByteString (ByteString)
  import Data.ByteString qualified as B
  import Data.ByteString.Base16
  import Data.ByteString.Builder
  import Data.ByteString.Lazy qualified as LB
  import Data.ByteString.Lazy.Base16 qualified as LB16
  import Data.Bytes.Get
  import Data.Bytes.Put
  import Data.Bytes.Serial
  import Data.Char (toLower)
  import Data.Int
  import Data.IntMap (IntMap)
  import Data.IntMap qualified as IntMap
  import Data.List
  import Data.Serialize qualified as S
  import Data.Text (Text)
  import Data.Text.Encoding qualified as T
  import Data.Text.Lazy qualified as LT
  import Data.Text.Lazy.Encoding qualified as LT
  import Data.Word
  import Test.Hspec
  
  -- ByteString helpers
  
  -- | Decode a big endian 'Integer' from a 'ByteString'.
  bsToInteger :: ByteString -> Integer
  bsToInteger = B.foldr f 0 . B.reverse
    where
      f w n = toInteger w .|. shiftL n 8
  
  -- | Encode an 'Integer' to a 'ByteString' as big endian.
  integerToBS :: Integer -> ByteString
  integerToBS 0 = B.pack [0]
  integerToBS i
    | i > 0 = B.reverse $ B.unfoldr f i
    | otherwise = error "integerToBS not defined for negative values"
    where
      f 0 = Nothing
      f x = Just (fromInteger x :: Word8, x `shiftR` 8)
  
  hexEncoding :: LB.ByteString -> Encoding
  hexEncoding b =
    unsafeToEncoding $
      char7 '"' <> hexBuilder b <> char7 '"'
  
  hexBuilder :: LB.ByteString -> Builder
  hexBuilder = lazyByteStringHex
  
  encodeHex :: ByteString -> Text
  encodeHex = extractBase16 . encodeBase16
  
  -- | Encode as string of human-readable hex characters.
  encodeHexLazy :: LB.ByteString -> LT.Text
  encodeHexLazy = extractBase16 . LB16.encodeBase16
  
  decodeHex :: Text -> Maybe ByteString
  decodeHex t =
    if isBase16 u8
      then Just . decodeBase16 $ assertBase16 u8
      else Nothing
    where
      u8 = T.encodeUtf8 t
  
  -- | Decode string of human-readable hex characters.
  decodeHexLazy :: LT.Text -> Maybe LB.ByteString
  decodeHexLazy t =
    if LB16.isBase16 u8
      then Just . LB16.decodeBase16 $ assertBase16 u8
      else Nothing
    where
      u8 = LT.encodeUtf8 t
  
  -- | Obtain 'Int' bits from beginning of 'ByteString'. Resulting 'ByteString'
  -- will be smallest required to hold that many bits, padded with zeroes to the
  -- right.
  getBits :: Int -> ByteString -> ByteString
  getBits b bs
    | r == 0 = B.take q bs
    | otherwise = i `B.snoc` l
    where
      (q, r) = b `quotRem` 8
      s = B.take (q + 1) bs
      i = B.init s
      l = B.last s .&. (0xff `shiftL` (8 - r)) -- zero unneeded bits
  
  -- Maybe and Either monad helpers
  
  -- | Transform an 'Either' value into a 'Maybe' value. 'Right' is mapped to
  -- 'Just' and 'Left' is mapped to 'Nothing'. The value inside 'Left' is lost.
  eitherToMaybe :: Either a b -> Maybe b
  eitherToMaybe (Right b) = Just b
  eitherToMaybe _ = Nothing
  
  -- | Transform a 'Maybe' value into an 'Either' value. 'Just' is mapped to
  -- 'Right' and 'Nothing' is mapped to 'Left'. Default 'Left' required.
  maybeToEither :: b -> Maybe a -> Either b a
  maybeToEither err = maybe (Left err) Right
  
  -- | Lift a 'Maybe' computation into the 'ExceptT' monad.
  liftMaybe :: (Monad m) => b -> Maybe a -> ExceptT b m a
  liftMaybe err = liftEither . maybeToEither err
  
  -- Various helpers
  
  -- | Applies a function to only one element of a list defined by its index.  If
  -- the index is out of the bounds of the list, the original list is returned.
  updateIndex ::
    -- | index of the element to change
    Int ->
    -- | list of elements
    [a] ->
    -- | function to apply
    (a -> a) ->
    -- | result with one element changed
    [a]
  updateIndex i xs f
    | i < 0 || i >= length xs = xs
    | otherwise = l ++ (f h : r)
    where
      (l, h : r) = splitAt i xs
  
  -- | Use the list @[b]@ as a template and try to match the elements of @[a]@
  -- against it. For each element of @[b]@ return the (first) matching element of
  -- @[a]@, or 'Nothing'. Output list has same size as @[b]@ and contains results
  -- in same order. Elements of @[a]@ can only appear once.
  matchTemplate ::
    -- | input list
    [a] ->
    -- | list to serve as a template
    [b] ->
    -- | comparison function
    (a -> b -> Bool) ->
    [Maybe a]
  matchTemplate [] bs _ = replicate (length bs) Nothing
  matchTemplate _ [] _ = []
  matchTemplate as (b : bs) f = case break (`f` b) as of
    (l, r : rs) -> Just r : matchTemplate (l ++ rs) bs f
    _ -> Nothing : matchTemplate as bs f
  
  -- | Returns the first value of a triple.
  fst3 :: (a, b, c) -> a
  fst3 (a, _, _) = a
  
  -- | Returns the second value of a triple.
  snd3 :: (a, b, c) -> b
  snd3 (_, b, _) = b
  
  -- | Returns the last value of a triple.
  lst3 :: (a, b, c) -> c
  lst3 (_, _, c) = c
  
  -- | Field label goes lowercase and first @n@ characters get removed.
  dropFieldLabel :: Int -> Options
  dropFieldLabel n =
    defaultOptions
      { fieldLabelModifier = map toLower . drop n
      }
  
  -- | Transformation from 'dropFieldLabel' is applied with argument @f@, plus
  -- constructor tags are lowercased and first @c@ characters removed. @tag@ is
  -- used as the name of the object field name that will hold the transformed
  -- constructor tag as its value.
  dropSumLabels :: Int -> Int -> String -> Options
  dropSumLabels c f tag =
    (dropFieldLabel f)
      { constructorTagModifier = map toLower . drop c,
        sumEncoding = defaultTaggedObject {tagFieldName = tag}
      }
  
  -- | Convert from one power-of-two base to another, as long as it fits in a
  -- 'Word'.
  convertBits :: Bool -> Int -> Int -> [Word] -> ([Word], Bool)
  convertBits pad frombits tobits i = (reverse yout, rem')
    where
      (xacc, xbits, xout) = foldl' outer (0, 0, []) i
      (yout, rem')
        | pad && xbits /= 0 =
            let xout' = (xacc `shiftL` (tobits - xbits)) .&. maxv : xout
             in (xout', False)
        | pad = (xout, False)
        | xbits /= 0 = (xout, True)
        | otherwise = (xout, False)
      maxv = 1 `shiftL` tobits - 1
      max_acc = 1 `shiftL` (frombits + tobits - 1) - 1
      outer (acc, bits, out) it =
        let acc' = ((acc `shiftL` frombits) .|. it) .&. max_acc
            bits' = bits + frombits
            (out', bits'') = inner acc' out bits'
         in (acc', bits'', out')
      inner acc out bits
        | bits >= tobits =
            let bits' = bits - tobits
                out' = ((acc `shiftR` bits') .&. maxv) : out
             in inner acc out' bits'
        | otherwise = (out, bits)
  
  --
  -- Serialization helpers
  --
  
  putInt32be :: (MonadPut m) => Int32 -> m ()
  putInt32be n
    | n < 0 = putWord32be (complement (fromIntegral (abs n)) + 1)
    | otherwise = putWord32be (fromIntegral (abs n))
  
  getInt32be :: (MonadGet m) => m Int32
  getInt32be = do
    n <- getWord32be
    if testBit n 31
      then return (negate (complement (fromIntegral n) + 1))
      else return (fromIntegral n)
  
  putInt64be :: (MonadPut m) => Int64 -> m ()
  putInt64be n
    | n < 0 = putWord64be (complement (fromIntegral (abs n)) + 1)
    | otherwise = putWord64be (fromIntegral (abs n))
  
  getInt64be :: (MonadGet m) => m Int64
  getInt64be = do
    n <- getWord64be
    if testBit n 63
      then return (negate (complement (fromIntegral n) + 1))
      else return (fromIntegral n)
  
  putInteger :: (MonadPut m) => Integer -> m ()
  putInteger n
    | n >= lo && n <= hi = do
        putWord8 0x00
        putInt32be (fromIntegral n)
    | otherwise = do
        putWord8 0x01
        putWord8 (fromIntegral (signum n))
        let len = (nrBits (abs n) + 7) `div` 8
        putWord64be (fromIntegral len)
        mapM_ putWord8 (unroll (abs n))
    where
      lo = fromIntegral (minBound :: Int32)
      hi = fromIntegral (maxBound :: Int32)
  
  getInteger :: (MonadGet m) => m Integer
  getInteger =
    getWord8 >>= \case
      0 -> fromIntegral <$> getInt32be
      _ -> do
        sign <- getWord8
        bytes <- getList getWord8
        let v = roll bytes
        return $! if sign == 0x01 then v else -v
  
  putMaybe :: (MonadPut m) => (a -> m ()) -> Maybe a -> m ()
  putMaybe f Nothing = putWord8 0x00
  putMaybe f (Just x) = putWord8 0x01 >> f x
  
  getMaybe :: (MonadGet m) => m a -> m (Maybe a)
  getMaybe f =
    getWord8 >>= \case
      0x00 -> return Nothing
      0x01 -> Just <$> f
      _ -> fail "Not a Maybe"
  
  putLengthBytes :: (MonadPut m) => ByteString -> m ()
  putLengthBytes bs = do
    putWord64be (fromIntegral (B.length bs))
    putByteString bs
  
  getLengthBytes :: (MonadGet m) => m ByteString
  getLengthBytes = do
    len <- fromIntegral <$> getWord64be
    getByteString len
  
  --
  -- Fold and unfold an Integer to and from a list of its bytes
  --
  unroll :: (Integral a, Bits a) => a -> [Word8]
  unroll = unfoldr step
    where
      step 0 = Nothing
      step i = Just (fromIntegral i, i `shiftR` 8)
  
  roll :: (Integral a, Bits a) => [Word8] -> a
  roll = foldr unstep 0
    where
      unstep b a = a `shiftL` 8 .|. fromIntegral b
  
  nrBits :: (Ord a, Integral a) => a -> Int
  nrBits k =
    let expMax = until (\e -> 2 ^ e > k) (* 2) 1
        findNr :: Int -> Int -> Int
        findNr lo hi
          | mid == lo = hi
          | 2 ^ mid <= k = findNr mid hi
          | 2 ^ mid > k = findNr lo mid
          where
            mid = (lo + hi) `div` 2
     in findNr (expMax `div` 2) expMax
  
  -- | Read as a list of pairs of int and element.
  getIntMap :: (MonadGet m) => m Int -> m a -> m (IntMap a)
  getIntMap i m = IntMap.fromList <$> getList (getTwo i m)
  
  putIntMap :: (MonadPut m) => (Int -> m ()) -> (a -> m ()) -> IntMap a -> m ()
  putIntMap f g = putList (putTwo f g) . IntMap.toAscList
  
  putTwo :: (MonadPut m) => (a -> m ()) -> (b -> m ()) -> (a, b) -> m ()
  putTwo f g (x, y) = f x >> g y
  
  getTwo :: (MonadGet m) => m a -> m b -> m (a, b)
  getTwo f g = (,) <$> f <*> g
  
  putList :: (MonadPut m) => (a -> m ()) -> [a] -> m ()
  putList f ls = do
    putWord64be (fromIntegral (length ls))
    mapM_ f ls
  
  getList :: (MonadGet m) => m a -> m [a]
  getList f = do
    l <- fromIntegral <$> getWord64be
    replicateM l f
  
  --
  -- Test Helpers
  --
  
  prepareContext :: (Ctx -> SpecWith a) -> SpecWith a
  prepareContext go = do
    ctx <- runIO $ do
      ctx <- createContext
      randomizeContext ctx
      return ctx
    afterAll_ (destroyContext ctx) (go ctx)
  
  customCerealID :: (Eq a) => S.Get a -> S.Putter a -> a -> Bool
  customCerealID g p a = (S.runGet g . S.runPut . p) a == Right a
  
  readTestFile :: (FromJSON a) => FilePath -> IO a
  readTestFile fp =
    eitherDecodeFileStrict ("data/" <> fp) >>= either (error . message) return
    where
      message aesonErr = "Could not read test file " <> fp <> ": " <> aesonErr
  
  readTestFileParser :: (Value -> Parser a) -> FilePath -> IO a
  readTestFileParser p fp = do
    v <- readTestFile fp
    case parse p v of
      Error s -> error s
      Success x -> return x
  