-- 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
  , encodeExpression'
  ) where

import qualified Data.Binary.Builder as Bi
import qualified Data.Binary.Get as Bi
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Sequence as Seq
import qualified Unsafe (fromIntegral)

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

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

-- | Encode 'Expression' to 'ByteString'.
encodeExpression :: Expression -> LByteString
encodeExpression = Bi.toLazyByteString . buildExpr

-- | Same as 'encodeExpression', for strict bytestring.
encodeExpression' :: Expression -> BS.ByteString
encodeExpression' = LBS.toStrict . encodeExpression

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

buildList :: [Expression] -> Bi.Builder
buildList = foldMap buildExpr

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

buildAnnotationList :: [Annotation] -> Bi.Builder
buildAnnotationList listAnn = buildDynamic buildText (DynamicSize $ unwords . fmap annotToText $ listAnn)

-------------------------------------------------
-- 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 getList)
  3 -> ExpressionPrim . (\pn -> MichelinePrimAp pn [] []) <$> getPrim
  4 -> ExpressionPrim <$> (flip MichelinePrimAp [] <$> getPrim <*> getAnnotationList)
  5 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (one <$> getExpr) <*> pure [])
  6 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (one <$> getExpr) <*> getAnnotationList)
  7 -> ExpressionPrim <$> ((\n a -> MichelinePrimAp n a []) <$> getPrim <*> replicateM 2 getExpr)
  8 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> replicateM 2 getExpr <*> getAnnotationList)
  9 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (unDynamicSize <$> (getDynamic getList)) <*> getAnnotationList)
  10 -> ExpressionBytes . unDynamicSize <$> (getDynamic getByteString)
  _ -> fail "invalid Micheline expression tag"

getList :: Bi.Get [Expression]
getList = 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

getAnnotationList :: Bi.Get [Annotation]
getAnnotationList = mapM annotFromText . words . unDynamicSize =<< (getDynamic getText)