-- SPDX-FileCopyrightText: 2020 Tocqueville Group
-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA
-- 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 Data.Binary.Builder qualified as Bi
import Data.Binary.Get qualified as Bi
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.Sequence qualified as Seq
import Unsafe qualified (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
  ExpSeq () xs -> buildWord8 2 <> buildDynamic buildList (DynamicSize xs)
  ExpPrim () (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
  ExpString () x -> buildWord8 1 <> buildDynamic buildText (DynamicSize x)
  ExpInt () x -> buildWord8 0 <> buildInteger x
  ExpBytes () 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 -> ExpInt () <$> getInteger
  1 -> ExpString () . unDynamicSize <$> (getDynamic getText)
  2 -> ExpSeq () . unDynamicSize <$> (getDynamic getList)
  3 -> ExpPrim () . (\pn -> MichelinePrimAp pn [] []) <$> getPrim
  4 -> ExpPrim () <$> (flip MichelinePrimAp [] <$> getPrim <*> getAnnotationList)
  5 -> ExpPrim () <$> (MichelinePrimAp <$> getPrim <*> (one <$> getExpr) <*> pure [])
  6 -> ExpPrim () <$> (MichelinePrimAp <$> getPrim <*> (one <$> getExpr) <*> getAnnotationList)
  7 -> ExpPrim () <$> ((\n a -> MichelinePrimAp n a []) <$> getPrim <*> replicateM 2 getExpr)
  8 -> ExpPrim () <$> (MichelinePrimAp <$> getPrim <*> replicateM 2 getExpr <*> getAnnotationList)
  9 -> ExpPrim () <$> (MichelinePrimAp <$> getPrim <*> (unDynamicSize <$> (getDynamic getList)) <*> getAnnotationList)
  10 -> ExpBytes () . 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)