-- SPDX-FileCopyrightText: 2020 Tocqueville Group
-- SPDX-FileCopyrightText: 2018 obsidian.systems
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ
-- SPDX-License-Identifier: LicenseRef-MIT-obsidian-systems

-- | Module that define encoding and decoding function from Expression type
-- to binary format.
module Morley.Micheline.Binary
  ( decodeExpression
  , eitherDecodeExpression
  , encodeExpression
  ) where

import qualified Data.Binary.Builder as Bi
import qualified Data.Binary.Get as Bi
import Data.Bits (Bits, bit, setBit, shift, testBit, zeroBits, (.&.), (.|.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.Sequence (Seq((:<|)))
import qualified Data.Sequence as Seq
import qualified Data.Text.Encoding as TE

import Morley.Micheline.Expression
import Util.Binary (UnpackError(..), ensureEnd, launchGet)

newtype DynamicSize a = DynamicSize { unDynamicSize :: a }

-------------------------------------------------
-- Encode
-------------------------------------------------

-- | Encode 'Expression' to 'ByteString'.
encodeExpression :: Expression -> BS.ByteString
encodeExpression = LBS.toStrict . Bi.toLazyByteString . buildExpr

buildExpr :: Expression -> Bi.Builder
buildExpr = \case
  ExpressionSeq xs -> buildWord8 2 <> buildDynamic buildSeq (DynamicSize xs)
  ExpressionPrim (MichelinePrimAp prim args annots) -> case (args, annots) of
    (Seq.Empty, Seq.Empty) -> buildWord8 3 <> buildPrim prim
    (Seq.Empty, _) -> buildWord8 4 <> buildPrim prim <> buildAnnotationSeq annots
    (arg1 :<| Seq.Empty, Seq.Empty) -> buildWord8 5 <> buildPrim prim <> buildExpr arg1
    (arg1 :<| Seq.Empty, _) -> buildWord8 6 <> buildPrim prim <> buildExpr arg1 <> buildAnnotationSeq annots
    (arg1 :<| (arg2 :<| Seq.Empty), Seq.Empty) -> buildWord8 7 <> buildPrim prim <> buildExpr arg1 <> buildExpr arg2
    (arg1 :<| (arg2 :<| Seq.Empty), _) -> buildWord8 8 <> buildPrim prim <> buildExpr arg1 <> buildExpr arg2 <> buildAnnotationSeq annots
    _ -> buildWord8 9 <> buildPrim prim <> buildDynamic buildSeq (DynamicSize args) <> buildAnnotationSeq annots
  ExpressionString x -> buildWord8 1 <> buildDynamic buildText (DynamicSize x)
  ExpressionInt x -> buildWord8 0 <> buildInteger x
  ExpressionBytes x -> buildWord8 10 <> buildDynamic buildByteString (DynamicSize x)

buildWord8 :: Word8 -> Bi.Builder
buildWord8 = Bi.singleton

buildByteString :: ByteString -> Bi.Builder
buildByteString = Bi.fromByteString

buildInteger :: Integer -> Bi.Builder
buildInteger n =
  let signBit = if n < 0 then bit 6 else zeroBits
      ab = abs n
  in
    -- Refer to: https://gitlab.com/obsidian.systems/tezos-bake-monitor-lib/-/blob/2cf12e53072bcd966d471430ead25f597db5e23f/tezos-bake-monitor-lib/src/Tezos/Common/Binary.hs#L122
    if ab < 0x40 then Bi.singleton (fromIntegral ab .|. signBit)
    else Bi.singleton (fromIntegral (ab .&. 0x3f) .|. signBit .|. bit 7) <> writeZ (-6) ab

writeZ :: (Integral a, Bits a) => Int -> a -> Bi.Builder
writeZ offset n =
  if n < bit (7 - offset) then Bi.singleton $ fromIntegral $ n `shift` offset
    else Bi.singleton (fromIntegral (((n `shift` offset) .&. 0x7f) `setBit` 7)) <> writeZ (offset - 7) n

buildDynamic :: (a -> Bi.Builder) -> (DynamicSize a) -> Bi.Builder
buildDynamic build (DynamicSize x) =
  let b = build x
  in Bi.putWord32be (fromIntegral $ LBS.length $ Bi.toLazyByteString b) <> b

buildText :: Text -> Bi.Builder
buildText n =
  buildByteString $ TE.encodeUtf8 n

buildSeq :: Seq Expression -> Bi.Builder
buildSeq xs = foldMap buildExpr xs

buildPrim :: MichelinePrimitive -> Bi.Builder
buildPrim (MichelinePrimitive p) = case Seq.elemIndexL p michelsonPrimitive of
  Nothing -> error "unknown Michelson/Micheline primitive"
  Just ix -> buildWord8 (fromIntegral ix)

buildAnnotationSeq :: Seq Annotation -> Bi.Builder
buildAnnotationSeq seqAnn = buildDynamic buildText (DynamicSize $ unwords . toList . fmap annotToText $ seqAnn)

-------------------------------------------------
-- Decode
-------------------------------------------------

-- | Decode 'Expression' from 'ByteString'.
eitherDecodeExpression :: BS.ByteString -> Either UnpackError Expression
eitherDecodeExpression x = launchGet (getExpr <* ensureEnd) $ LBS.fromStrict x

-- | Partial version of 'eitherDecodeExpression'.
decodeExpression :: HasCallStack => BS.ByteString -> Expression
decodeExpression = either (error . unUnpackError) id . eitherDecodeExpression

getExpr :: Bi.Get Expression
getExpr = Bi.getWord8 >>= \case
  0 -> ExpressionInt <$> getInteger
  1 -> ExpressionString . unDynamicSize <$> (getDynamic getText)
  2 -> ExpressionSeq . unDynamicSize <$> (getDynamic getSeq)
  3 -> ExpressionPrim . (\pn -> MichelinePrimAp pn Seq.Empty Seq.Empty) <$> getPrim
  4 -> ExpressionPrim <$> (flip MichelinePrimAp Seq.Empty <$> getPrim <*> getAnnotationSeq)
  5 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (Seq.singleton <$> getExpr) <*> pure Seq.empty)
  6 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (Seq.singleton <$> getExpr) <*> getAnnotationSeq)
  7 -> ExpressionPrim <$> ((\n a -> MichelinePrimAp n a Seq.empty) <$> getPrim <*> Seq.replicateA 2 getExpr)
  8 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> Seq.replicateA 2 getExpr <*> getAnnotationSeq)
  9 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (unDynamicSize <$> (getDynamic getSeq)) <*> getAnnotationSeq)
  10 -> ExpressionBytes . unDynamicSize <$> (getDynamic getByteString)
  _ -> fail "invalid Micheline expression tag"

getInteger :: Bi.Get Integer
getInteger = do
  b <- Bi.getWord8
  n <- if b `testBit` 7 then readZ 6 (fromIntegral $ b .&. 0x3f) else pure (fromIntegral $ b .&. 0x3f)
  pure $ if b `testBit` 6 then negate n else n

readZ :: (Num a, Bits a) => Int -> a -> Bi.Get a
readZ offset n = do
  b <- Bi.getWord8
  if (b == 0) && (offset > 0) then fail "trailing zero" else pure ()
  let n' = (fromIntegral (b .&. 0x7f) `shift` offset) .|. n
  if b `testBit` 7 then readZ (offset + 7) n' else pure n'

getDynamic :: (Bi.Get a) -> (Bi.Get (DynamicSize a))
getDynamic getter = do
  len <- fromIntegral <$> Bi.getWord32be
  DynamicSize <$> Bi.isolate len getter

{-# ANN getText ("HLint: ignore Redundant fmap" :: Text) #-}
getText :: Bi.Get Text
getText =
  fmap decodeUtf8' getByteString >>= \case
    Left err -> fail $ show err
    Right answer -> pure answer

getByteString :: Bi.Get ByteString
getByteString =
  LBS.toStrict <$> Bi.getRemainingLazyByteString

getSeq :: Bi.Get (Seq Expression)
getSeq =
   Seq.fromList <$> many getExpr

getPrim :: Bi.Get MichelinePrimitive
getPrim =
  Bi.getWord8 >>= \ix -> case Seq.lookup (fromIntegral ix) michelsonPrimitive of
    Nothing -> fail "unknown Michelson/Micheline opcode"
    Just str -> pure $ MichelinePrimitive str

getAnnotationSeq :: Bi.Get (Seq Annotation)
getAnnotationSeq = mapM annotFromText . Seq.fromList . words . unDynamicSize =<< (getDynamic getText)