-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | This module contains helper functions to deal with encoding
-- and decoding of binary data.
module Morley.Util.Binary
  ( UnpackError (..)
  , ensureEnd
  , launchGet
  , TaggedDecoder
  , TaggedDecoderM(..)
  , (#:)
  , (##:)
  , decodeBytesLike
  , decodeWithTag
  , decodeWithTagM
  , getByteStringCopy
  , getRemainingByteStringCopy
  , unknownTag
  ) where

import Prelude hiding (EQ, Ordering(..), get)

import Data.Binary (Get)
import Data.Binary.Get qualified as Get
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.List qualified as List
import Fmt (Buildable, build, fmt, hexF, pretty, (+|), (|+))
import Text.Hex (encodeHex)

----------------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------------

-- | Any decoding error.
newtype UnpackError = UnpackError { unUnpackError :: Text }
  deriving stock (Show, Eq)

instance Buildable UnpackError where
  build (UnpackError msg) = build msg

instance Exception UnpackError where
  displayException = pretty

ensureEnd :: Get ()
ensureEnd =
  unlessM Get.isEmpty $ do
    remainder <- Get.getRemainingLazyByteString
    fail $ "Expected end of entry, unconsumed bytes \
           \(" +| length remainder |+ "): \""
           +| encodeHex (LBS.toStrict remainder) |+ "\""

launchGet :: Get a -> LByteString -> Either UnpackError a
launchGet decoder bs =
  case Get.runGetOrFail decoder bs of
    Left (_remainder, _offset, err) -> Left . UnpackError $ toText err
    Right (_remainder, _offset, res) -> Right res

-- | Specialization of 'TaggedDecoderM' to 'IdentityT' transformer.
type TaggedDecoder a = TaggedDecoderM IdentityT a

-- | Describes how 'decodeWithTagM' should decode tag-dependent data.
-- We expect bytes of such structure: 'tdTag' followed by a bytestring
-- which will be parsed with 'tdDecoder'.
data TaggedDecoderM t a = TaggedDecoder
  { tdTag :: Word8
  , tdDecoder :: t Get a
  }

-- | Alias for v'TaggedDecoder' constructor specialized to 'Get'
(#:) :: Word8 -> Get a -> TaggedDecoder a
(#:) t = TaggedDecoder t . lift
infixr 0 #:

-- | Alias for v'TaggedDecoder' constructor.
(##:) :: Word8 -> t Get a -> TaggedDecoderM t a
(##:) = TaggedDecoder
infixr 0 ##:

-- | Get a bytestring of the given length leaving no references to the
-- original data in serialized form.
getByteStringCopy :: Int -> Get ByteString
getByteStringCopy = fmap BS.copy . Get.getByteString

-- | Get remaining available bytes.
--
-- Note that reading all remaining decoded input may be expensive and is thus
-- discouraged, use can use this function only when you know that amount
-- of data to be consumed is limited, e.g. within 'decodeBytesLike' call.
getRemainingByteStringCopy :: Get ByteString
getRemainingByteStringCopy = do
  lbs <- Get.getRemainingLazyByteString
  -- Avoiding memory overflows in case bad length to 'Get.isolate' was provided.
  -- Normally this function is used only to decode primitives, 'Signature' in
  -- the worst case, so we could set little length, but since this is a hack
  -- anyway let's make sure it never obstructs our work.
  when (length lbs > 640000) $ fail "Too big length for an entity"
  return (LBS.toStrict lbs)

-- | Fail with "unknown tag" error.
unknownTag :: String -> Word8 -> Get a
unknownTag desc tag =
  fail . fmt $ "Unknown " <> build desc <> " tag: 0x" <> hexF tag

-- | Common decoder for the case when packed data starts with a tag (1 byte)
-- that specifies how to decode remaining data.
--
-- This is a version of 'decodeWithTagM' specialized to naked 'Get' monad.
decodeWithTag :: String -> [TaggedDecoder a] -> Get a
decodeWithTag what decoders =
  runIdentityT $ decodeWithTagM what (lift . unknownTag what) decoders

-- | Common decoder for the case when packed data starts with a tag (1 byte)
-- that specifies how to decode remaining data.
--
-- This is a general version of 'decodeWithTag' that allows 'Get' to be wrapped
-- in a monad transformer.
decodeWithTagM
  :: (MonadTrans t, Monad (t Get))
  => String
  -> (Word8 -> t Get a)
  -> [TaggedDecoderM t a]
  -> t Get a
decodeWithTagM what unknownTagFail decoders = do
  tag <- lift $ Get.label (what <> " tag") Get.getWord8
  -- Number of decoders is usually small, so linear runtime lookup should be ok.
  case List.find ((tag ==) . tdTag) decoders of
    Nothing -> unknownTagFail tag
    Just TaggedDecoder{..} -> tdDecoder

decodeBytesLike
  :: (Buildable err)
  => String -> (ByteString -> Either err a) -> Get a
decodeBytesLike what constructor = do
  bs <- getRemainingByteStringCopy
  case constructor bs of
    Left err -> fail $ "Wrong " +| what |+ ": " +| err |+ ""
    Right res -> pure res