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

-- | Module that provides type classes for converting to and from low-level
-- Micheline representation.
module Morley.Micheline.Class
  ( ToExpression (..)
  , FromExpError (..)
  , FromExpressionError
  , FromExp (..)
  , FromExpression
  , fromExpression
  ) where

import Control.Lens ((<>~))
import Data.Bits (Bits)
import Data.Default (def)
import Data.Singletons (SingI(..), demote)
import Fmt (Buildable(..), indentF, pretty, unlinesF)

import Morley.Micheline.Expression
import Morley.Michelson.Text (mkMText, unMText)
import Morley.Michelson.TypeCheck
  (TypeCheckMode(..), TypeCheckOptions(..), runTypeCheck, typeCheckingWith)
import Morley.Michelson.TypeCheck.Instr (typeCheckValue)
import Morley.Michelson.Typed
  (Contract, HasNoOp, Instr, Notes(..), T(..), Value, Value'(..), fromUType, mkUType, rfAnyInstr,
  toUType)
import Morley.Michelson.Typed.Convert (convertContract, instrToOpsOptimized, untypeValueOptimized)
import Morley.Michelson.Untyped qualified as Untyped
import Morley.Michelson.Untyped.Annotation
  (AnnotationSet(..), FieldAnn, FieldTag, RootAnn, TypeAnn, TypeTag, VarAnn, VarTag, annsCount,
  emptyAnnSet, firstAnn, noAnn, secondAnn)
import Morley.Michelson.Untyped.Contract (ContractBlock(..), orderContractBlock)
import Morley.Michelson.Untyped.Instr (ExpandedInstr, ExpandedOp(..), InstrAbstract(..))
import Morley.Michelson.Untyped.Type (Ty(..))
import Morley.Michelson.Untyped.View

-- ToExpression
----------------------------------------------------------------------------

-- | Type class that provides an ability to convert
-- something to Micheline Expression.
class ToExpression a where
  toExpression :: a -> Expression

instance (HasNoOp t) => ToExpression (Value t) where
  toExpression = toExpression . untypeValueOptimized

instance ToExpression Untyped.Value where
  toExpression = \case
    Untyped.ValueInt v -> expressionInt v
    Untyped.ValueString s -> expressionString $ unMText s
    Untyped.ValueBytes (Untyped.InternalByteString bs) -> expressionBytes bs
    Untyped.ValueUnit -> expressionPrim' "Unit" [] []
    Untyped.ValueTrue -> expressionPrim' "True" [] []
    Untyped.ValueFalse -> expressionPrim' "False" [] []
    Untyped.ValuePair l r ->
      expressionPrim' "Pair" [toExpression l, toExpression r] []
    Untyped.ValueLeft v -> expressionPrim' "Left" [toExpression v] []
    Untyped.ValueRight v -> expressionPrim' "Right" [toExpression v] []
    Untyped.ValueSome v -> expressionPrim' "Some" [toExpression v] []
    Untyped.ValueNone -> expressionPrim' "None" [] []
    Untyped.ValueNil -> expressionSeq []
    Untyped.ValueSeq vs -> toExpression vs
    Untyped.ValueMap elts -> toExpression $ eltToExpr <$> elts
    Untyped.ValueLambda ops -> toExpression ops
    where
      eltToExpr :: Untyped.Elt ExpandedOp -> Expression
      eltToExpr (Untyped.Elt l r) = expressionPrim' "Elt"
        [toExpression l, toExpression r] []


instance ToExpression (Instr inp out) where
  toExpression = toExpression . instrToOpsOptimized

instance ToExpression T where
  toExpression = toExpression . toUType

instance ToExpression (Notes t) where
  toExpression = toExpression . mkUType

instance ToExpression Untyped.T where
  toExpression = \case
    Untyped.TKey -> expressionPrim' "key" [] []
    Untyped.TUnit -> expressionPrim' "unit" [] []
    Untyped.TSignature -> expressionPrim' "signature" [] []
    Untyped.TChainId -> expressionPrim' "chain_id" [] []
    Untyped.TOption arg -> expressionPrim' "option" [toExpression arg] []
    Untyped.TList arg -> expressionPrim' "list" [toExpression arg] []
    Untyped.TSet arg -> expressionPrim' "set" [toExpression arg] []
    Untyped.TOperation -> expressionPrim' "operation" [] []
    Untyped.TContract arg -> expressionPrim' "contract" [toExpression arg] []
    Untyped.TTicket arg -> expressionPrim' "ticket" [toExpression arg] []
    t@Untyped.TPair{} -> expressionPrim' "pair"
      (rightCombedPairToList (Ty t noAnn) (noAnn, noAnn)) []
    Untyped.TOr fa1 fa2 l r ->
      let exprL = addTrimmedAnns (toExpression l) [] [fa1] []
          exprR = addTrimmedAnns (toExpression r) [] [fa2] []
      in expressionPrim' "or" [exprL, exprR] []
    Untyped.TLambda inp out ->
      expressionPrim' "lambda" [toExpression inp, toExpression out] []
    Untyped.TMap k v ->
      expressionPrim' "map" [toExpression k, toExpression v] []
    Untyped.TBigMap k v ->
      expressionPrim' "big_map" [toExpression k, toExpression v] []
    Untyped.TInt -> expressionPrim' "int" [] []
    Untyped.TNat -> expressionPrim' "nat" [] []
    Untyped.TString -> expressionPrim' "string" [] []
    Untyped.TBytes -> expressionPrim' "bytes" [] []
    Untyped.TMutez -> expressionPrim' "mutez" [] []
    Untyped.TBool -> expressionPrim' "bool" [] []
    Untyped.TKeyHash -> expressionPrim' "key_hash" [] []
    Untyped.TBls12381Fr -> expressionPrim' "bls12_381_fr" [] []
    Untyped.TBls12381G1 -> expressionPrim' "bls12_381_g1" [] []
    Untyped.TBls12381G2 -> expressionPrim' "bls12_381_g2" [] []
    Untyped.TTimestamp -> expressionPrim' "timestamp" [] []
    Untyped.TAddress -> expressionPrim' "address" [] []
    Untyped.TChest -> expressionPrim' "chest" [] []
    Untyped.TChestKey -> expressionPrim' "chest_key" [] []
    Untyped.TTxRollupL2Address -> expressionPrim' "tx_rollup_l2_address" [] []
    Untyped.TNever -> expressionPrim' "never" [] []
    Untyped.TSaplingState n -> expressionPrim' "sapling_state" [integralToExpr n] []
    Untyped.TSaplingTransaction n -> expressionPrim' "sapling_transaction" [integralToExpr n] []

    where
      addAnns :: Expression -> [Annotation] -> Expression
      addAnns e anns =
        e & _ExpressionPrim . mpaAnnotsL <>~ anns

      rightCombedPairToList :: Ty -> (FieldAnn, VarAnn) -> [Expression]
      rightCombedPairToList ty (fa, va) = case (ty, fa) of
        (Ty (Untyped.TPair fa1 fa2 va1 va2 l r) (Untyped.Annotation ""), Untyped.Annotation "") ->
          let annsL = mkAnns [] [fa1] [va1]
              exprL = toExpression l `addAnns` annsL
          in exprL : (rightCombedPairToList r (fa2, va2))
        _ ->
          let anns = mkAnns [] [fa] [va]
          in one $ toExpression ty `addAnns` anns

instance ToExpression Ty where
  toExpression (Ty t ta) = addTrimmedAnns (toExpression t) [ta] [] []

instance (ToExpression a) => ToExpression [a] where
  toExpression xs = expressionSeq $ toExpression <$> xs

instance (ToExpression a) => ToExpression (NonEmpty a) where
  toExpression = toExpression . toList

instance ToExpression Expression where
  toExpression = id

instance ToExpression ExpandedOp where
  toExpression = \case
    PrimEx instr   -> toExpression instr
    SeqEx s        -> expressionSeq $ toExpression <$> s
    WithSrcEx _ op -> toExpression op

instance ToExpression ViewName where
  toExpression (ViewName s) = expressionString s

instance ToExpression ExpandedInstr where
  toExpression = \case
    PUSH va ty v -> expressionPrim' "PUSH" [toExpression ty, toExpression v] $
      mkAnns [] [] [va]
    DROP -> expressionPrim' "DROP" [] []
    DROPN n -> expressionPrim' "DROP" [integralToExpr n] []
    DUP va -> expressionPrim' "DUP" [] $ mkAnns [] [] [va]
    DUPN va n -> expressionPrim' "DUP" [integralToExpr n] $ mkAnns [] [] [va]
    SWAP -> expressionPrim' "SWAP" [] []
    DIG n -> expressionPrim' "DIG" [integralToExpr n] []
    DUG n -> expressionPrim' "DUG" [integralToExpr n] []
    SOME ta va ->
      expressionPrim' "SOME" [] $ mkAnns [ta] [] [va]
    NONE ta va ty -> expressionPrim' "NONE" [toExpression ty] $ mkAnns [ta] [] [va]
    UNIT ta va -> expressionPrim' "UNIT" [] $ mkAnns [ta] [] [va]
    IF_NONE ops1 ops2 ->
      expressionPrim' "IF_NONE" [toExpression ops1, toExpression ops2] []
    PAIR ta va fa1 fa2 -> expressionPrim' "PAIR" [] $ mkAnns [ta] [fa1, fa2] [va]
    UNPAIR va1 va2 fa1 fa2 -> expressionPrim' "UNPAIR" [] $
      mkAnns [] [fa1, fa2] [va1, va2]
    PAIRN va n -> expressionPrim' "PAIR" [integralToExpr n] $ mkAnns [] [] [va]
    UNPAIRN n -> expressionPrim' "UNPAIR" [integralToExpr n] []
    CAR va fa -> expressionPrim' "CAR" [] $ mkAnns [] [fa] [va]
    CDR va fa -> expressionPrim' "CDR" [] $ mkAnns [] [fa] [va]
    LEFT ta va fa1 fa2 ty -> expressionPrim' "LEFT" [toExpression ty] $
      mkAnns [ta] [fa1, fa2] [va]
    RIGHT ta va fa1 fa2 ty -> expressionPrim' "RIGHT" [toExpression ty] $
      mkAnns [ta] [fa1, fa2] [va]
    IF_LEFT ops1 ops2 ->
      expressionPrim' "IF_LEFT" [toExpression ops1, toExpression ops2] []
    NIL ta va ty -> expressionPrim' "NIL" [toExpression ty] $
      mkAnns [ta] [] [va]
    CONS va -> expressionPrim' "CONS" [] $ mkAnns [] [] [va]
    IF_CONS ops1 ops2 ->
      expressionPrim' "IF_CONS" [toExpression ops1, toExpression ops2] []
    SIZE va -> expressionPrim' "SIZE" [] $ mkAnns [] [] [va]
    EMPTY_SET ta va ty -> expressionPrim' "EMPTY_SET" [toExpression ty] $
      mkAnns [ta] [] [va]
    EMPTY_MAP ta va kty vty ->
      expressionPrim' "EMPTY_MAP" [toExpression kty, toExpression vty] $
      mkAnns [ta] [] [va]
    EMPTY_BIG_MAP ta va kty vty ->
      expressionPrim' "EMPTY_BIG_MAP" [toExpression kty, toExpression vty] $
      mkAnns [ta] [] [va]
    MAP va ops -> expressionPrim' "MAP" [toExpression ops] $ mkAnns [] [] [va]
    ITER ops -> expressionPrim' "ITER" [toExpression ops] []
    MEM va -> expressionPrim' "MEM" [] $ mkAnns [] [] [va]
    GET va -> expressionPrim' "GET" [] $ mkAnns [] [] [va]
    GETN va n -> expressionPrim' "GET" [integralToExpr n] $ mkAnns [] [] [va]
    UPDATE va -> expressionPrim' "UPDATE" [] $ mkAnns [] [] [va]
    UPDATEN va n -> expressionPrim' "UPDATE" [integralToExpr n] $ mkAnns [] [] [va]
    GET_AND_UPDATE va -> expressionPrim' "GET_AND_UPDATE" [] $ mkAnns [] [] [va]
    IF ops1 ops2 ->
      expressionPrim' "IF" [toExpression ops1, toExpression ops2] []
    LOOP ops -> expressionPrim' "LOOP" [toExpression ops] []
    LOOP_LEFT ops -> expressionPrim' "LOOP_LEFT" [toExpression ops] []
    LAMBDA va tyin tyout ops ->
      expressionPrim' "LAMBDA" [ toExpression tyin
                            , toExpression tyout
                            , toExpression ops
                            ] $ mkAnns [] [] [va]
    EXEC va -> expressionPrim' "EXEC" [] $ mkAnns [] [] [va]
    APPLY va -> expressionPrim' "APPLY" [] $ mkAnns [] [] [va]
    DIP ops -> expressionPrim' "DIP" [toExpression ops] []
    DIPN n ops -> expressionPrim' "DIP" [integralToExpr n, toExpression ops] []
    FAILWITH -> expressionPrim' "FAILWITH" [] []
    CAST va ty -> expressionPrim' "CAST" [toExpression ty] $ mkAnns [] [] [va]
    RENAME va -> expressionPrim' "RENAME" [] $ mkAnns [] [] [va]
    PACK va -> expressionPrim' "PACK" [] $ mkAnns [] [] [va]
    UNPACK ta va ty -> expressionPrim' "UNPACK" [toExpression ty] $
      mkAnns [ta] [] [va]
    CONCAT va -> expressionPrim' "CONCAT" [] $ mkAnns [] [] [va]
    SLICE va -> expressionPrim' "SLICE" [] $ mkAnns [] [] [va]
    ISNAT va -> expressionPrim' "ISNAT" [] $ mkAnns [] [] [va]
    ADD va -> expressionPrim' "ADD" [] $ mkAnns [] [] [va]
    SUB va -> expressionPrim' "SUB" [] $ mkAnns [] [] [va]
    SUB_MUTEZ va -> expressionPrim' "SUB_MUTEZ" [] $ mkAnns [] [] [va]
    MUL va -> expressionPrim' "MUL" [] $ mkAnns [] [] [va]
    EDIV va -> expressionPrim' "EDIV" [] $ mkAnns [] [] [va]
    ABS va -> expressionPrim' "ABS" [] $ mkAnns [] [] [va]
    NEG va -> expressionPrim' "NEG" [] $ mkAnns [] [] [va]
    LSL va -> expressionPrim' "LSL" [] $ mkAnns [] [] [va]
    LSR va -> expressionPrim' "LSR" [] $ mkAnns [] [] [va]
    OR va -> expressionPrim' "OR" [] $ mkAnns [] [] [va]
    AND va -> expressionPrim' "AND" [] $ mkAnns [] [] [va]
    XOR va -> expressionPrim' "XOR" [] $ mkAnns [] [] [va]
    NOT va -> expressionPrim' "NOT" [] $ mkAnns [] [] [va]
    COMPARE va -> expressionPrim' "COMPARE" [] $ mkAnns [] [] [va]
    Untyped.EQ va -> expressionPrim' "EQ" [] $ mkAnns [] [] [va]
    NEQ va -> expressionPrim' "NEQ" [] $ mkAnns [] [] [va]
    Untyped.LT va -> expressionPrim' "LT" [] $ mkAnns [] [] [va]
    Untyped.GT va -> expressionPrim' "GT" [] $ mkAnns [] [] [va]
    LE va -> expressionPrim' "LE" [] $ mkAnns [] [] [va]
    GE va -> expressionPrim' "GE" [] $ mkAnns [] [] [va]
    INT va -> expressionPrim' "INT" [] $ mkAnns [] [] [va]
    VIEW va n t -> expressionPrim' "VIEW" [toExpression n, toExpression t] $
      mkAnns [] [] [va]
    SELF va fa -> expressionPrim' "SELF" [] $ mkAnns [] [fa] [va]
    CONTRACT va fa ty -> expressionPrim' "CONTRACT" [toExpression ty] $
      mkAnns [] [fa] [va]
    TRANSFER_TOKENS va -> expressionPrim' "TRANSFER_TOKENS" [] $ mkAnns [] [] [va]
    SET_DELEGATE va -> expressionPrim' "SET_DELEGATE" [] $ mkAnns [] [] [va]
    CREATE_CONTRACT va1 va2 c ->
      expressionPrim' "CREATE_CONTRACT" [toExpression c] $
      mkAnns [] [] [va1, va2]
    IMPLICIT_ACCOUNT va -> expressionPrim' "IMPLICIT_ACCOUNT" [] $
      mkAnns [] [] [va]
    NOW va -> expressionPrim' "NOW" [] $ mkAnns [] [] [va]
    AMOUNT va -> expressionPrim' "AMOUNT" [] $ mkAnns [] [] [va]
    BALANCE va -> expressionPrim' "BALANCE" [] $ mkAnns [] [] [va]
    VOTING_POWER va -> expressionPrim' "VOTING_POWER" [] $ mkAnns [] [] [va]
    TOTAL_VOTING_POWER va -> expressionPrim' "TOTAL_VOTING_POWER" [] $
      mkAnns [] [] [va]
    CHECK_SIGNATURE va -> expressionPrim' "CHECK_SIGNATURE" [] $ mkAnns [] [] [va]
    SHA256 va -> expressionPrim' "SHA256" [] $ mkAnns [] [] [va]
    SHA512 va -> expressionPrim' "SHA512" [] $ mkAnns [] [] [va]
    BLAKE2B va -> expressionPrim' "BLAKE2B" [] $ mkAnns [] [] [va]
    SHA3 va -> expressionPrim' "SHA3" [] $ mkAnns [] [] [va]
    KECCAK va -> expressionPrim' "KECCAK" [] $ mkAnns [] [] [va]
    HASH_KEY va -> expressionPrim' "HASH_KEY" [] $ mkAnns [] [] [va]
    PAIRING_CHECK va -> expressionPrim' "PAIRING_CHECK" [] $ mkAnns [] [] [va]
    SOURCE va -> expressionPrim' "SOURCE" [] $ mkAnns [] [] [va]
    SENDER va -> expressionPrim' "SENDER" [] $ mkAnns [] [] [va]
    ADDRESS va -> expressionPrim' "ADDRESS" [] $ mkAnns [] [] [va]
    CHAIN_ID va -> expressionPrim' "CHAIN_ID" [] $ mkAnns [] [] [va]
    LEVEL va -> expressionPrim' "LEVEL" [] $ mkAnns [] [] [va]
    SELF_ADDRESS va -> expressionPrim' "SELF_ADDRESS" [] $ mkAnns [] [] [va]
    TICKET va -> expressionPrim' "TICKET" [] $ mkAnns [] [] [va]
    READ_TICKET va -> expressionPrim' "READ_TICKET" [] $ mkAnns [] [] [va]
    SPLIT_TICKET va -> expressionPrim' "SPLIT_TICKET" [] $ mkAnns [] [] [va]
    JOIN_TICKETS va -> expressionPrim' "JOIN_TICKETS" [] $ mkAnns [] [] [va]
    OPEN_CHEST va -> expressionPrim' "OPEN_CHEST" [] $ mkAnns [] [] [va]
    NEVER -> expressionPrim' "NEVER" [] []
    EXT _ -> expressionSeq []
    SAPLING_EMPTY_STATE va n -> expressionPrim' "SAPLING_EMPTY_STATE" [integralToExpr n] $ mkAnns [] [] [va]
    SAPLING_VERIFY_UPDATE va -> expressionPrim' "SAPLING_VERIFY_UPDATE" [] $ mkAnns [] [] [va]
    MIN_BLOCK_TIME va -> expressionPrim' "MIN_BLOCK_TIME" [] $ mkAnnsFromAny va

instance ToExpression Untyped.Contract where
  toExpression contract
    = expressionSeq $ Untyped.mapEntriesOrdered contract
          (\(Untyped.ParameterType ty rootAnn) -> expressionPrim' "parameter"
            [insertRootAnn (toExpression ty) rootAnn] [])
          (\storage -> expressionPrim' "storage" [toExpression storage] [])
          (\code -> expressionPrim' "code" [toExpression code] [])
          (\Untyped.View{..} -> expressionPrim' "view"
            [toExpression viewName, toExpression viewArgument, toExpression viewReturn, toExpression viewCode] []
          )

instance ToExpression (Contract cp st) where
  toExpression = toExpression . convertContract

-- FromExpression
----------------------------------------------------------------------------

-- | Errors that can happen when we convert an 'Exp' to our
-- data type.
data FromExpError x = FromExpError (Exp x) Text

deriving stock instance Show (Exp x) => Show (FromExpError x)
deriving stock instance Eq (Exp x) => Eq (FromExpError x)

-- | Error in case of vanilla expression.
type FromExpressionError = FromExpError RegularExp

instance Buildable FromExpressionError where
  build (FromExpError expr err) =
    unlinesF
      [ "Failed to convert expression:"
      , indentF 2 $ build expr
      , ""
      , "Error:"
      , indentF 2 $ build err
      ]

instance Exception FromExpressionError where
  displayException = pretty

-- | Type class that provides the ability to convert
-- something from a Micheline Expression.
class FromExp x a where
  fromExp :: Exp x -> Either (FromExpError x) a

type FromExpression = FromExp RegularExp

-- | Parse vanilla expression to something.
fromExpression
  :: FromExp RegularExp a
  => Expression -> Either FromExpressionError a
fromExpression = fromExp

instance (FromExp x Untyped.Value, SingI t) => FromExp x (Value t) where
  fromExp expr =
    case fromExp @_ @Untyped.Value expr of
    Right uv -> case typeCheck uv of
      Left tcErr -> Left $ FromExpError expr $
        pretty $ unlinesF
          [ "Failed to typecheck expression as a value of type:"
          , indentF 2 $ build $ demote @t
          , ""
          , "Typechecker error:"
          , indentF 2 $ build tcErr
          ]
      Right tv -> Right tv
    Left e -> Left e
    where
      typeCheck uv = typeCheckingWith (TypeCheckOptions False False) $
        (runTypeCheck $ TypeCheckValue (uv, demote @t)) $
        usingReaderT def $
        typeCheckValue uv

instance FromExp x op => FromExp x (Untyped.Value' op) where
  fromExp e = case e of
    ExpInt _ v -> pure $ Untyped.ValueInt v
    ExpString _ s -> first (FromExpError e)
      (Untyped.ValueString <$> mkMText s)
    ExpBytes _ bs -> pure $ Untyped.ValueBytes $ Untyped.InternalByteString bs
    ExpPrim' _ "Unit" [] [] -> pure Untyped.ValueUnit
    ExpPrim' _ "True" [] [] -> pure Untyped.ValueTrue
    ExpPrim' _ "False" [] [] -> pure Untyped.ValueFalse
    ExpPrim' _ "Pair" args [] ->
      case nonEmpty args >>= forbidSingletonList of
        Nothing -> Left $ FromExpError e
          "Expected a pair with at least 2 arguments"
        Just args' -> do
          tys <- mapM fromExp args'
          return $ foldr1 Untyped.ValuePair tys
    ExpPrim' _ "Left" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.ValueLeft arg'
    ExpPrim' _ "Right" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.ValueRight arg'
    ExpPrim' _ "Some" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.ValueSome arg'
    ExpPrim' _ "None" [] [] -> pure Untyped.ValueNone
    ExpSeq _ [] -> pure Untyped.ValueNil
    ExpSeq _ (h : t) ->
      case fromExp @x @op h of
      Right op -> do
        ops <- traverse (fromExp @x @op) t
        pure . Untyped.ValueLambda $ op :| ops
      Left _ -> case exprToElt h of
        Right elt -> do
          elts <- traverse exprToElt t
          pure . Untyped.ValueMap $ elt :| elts
        Left _ -> case fromExp h of
          Left (FromExpError err _) -> Left $ FromExpError err
            "Value, instruction or 'Elt' expression expected"
          Right h' -> do
            t' <- traverse fromExp t
            pure . Untyped.ValueSeq $ h' :| t'
    _ -> Left $ FromExpError e "Expected a value"
    where
      exprToElt :: Exp x -> Either (FromExpError x) (Untyped.Elt op)
      exprToElt ex = case ex of
        ExpPrim' _ "Elt" [l, r] [] -> do
          l' <- fromExp l
          r' <- fromExp r
          pure $ Untyped.Elt l' r'
        ExpPrim' _ "Elt" _ [] -> Left $ FromExpError ex
          "Expected 'Elt' expression with exactly 2 elements"
        ExpPrim' _ "Elt" _ _ -> Left $ FromExpError ex
          "Expected 'Elt' expression without annotations"
        _ -> Left $ FromExpError ex "Expected 'Elt' expression"

instance FromExp x a => FromExp x [a] where
  fromExp = \case
    ExpSeq _ exprs -> traverse fromExp exprs
    e -> Left $ FromExpError e "'ExpressionSeq' expected"

instance FromExp RegularExp ExpandedOp where
  fromExp = \case
    ExpSeq _ s -> SeqEx <$> traverse fromExp s
    e          -> PrimEx <$> fromExp e

instance FromExp x op =>
         FromExp x (InstrAbstract op) where
  fromExp e = let annSet = getAnnSet e in case e of
    ExpPrim' _ "DROP" [n] [] -> do
      n' <- integralFromExpr n
      pure $ DROPN n'
    ExpPrim' _ "DROP" [] _ -> pure $ DROP
    ExpPrim' _ "DUP" [n] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      let va = firstAnn @VarTag annSet
      n' <- integralFromExpr n
      pure $ DUPN va n'
    ExpPrim' _ "DUP" [] _ ->
      let va = firstAnn @VarTag annSet
      in checkAnnsCount e annSet (0, 0, 1) $> DUP va
    ExpPrim' _ "SWAP" [] [] -> pure $ SWAP
    ExpPrim' _ "DIG" [n] [] -> do
      n' <- integralFromExpr n
      pure $ DIG $ n'
    ExpPrim' _ "DUG" [n] [] -> do
      n' <- integralFromExpr n
      pure $ DUG n'
    ExpPrim' _ "PUSH" [t, v] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      let va = firstAnn @VarTag annSet
      t' <- fromExp @x @Ty t
      v' <- fromExp @x @(Untyped.Value' op) v
      pure $ PUSH va t' v'
    ExpPrim' _ "SOME" [] _ ->
      let ta = firstAnn @TypeTag annSet
          va = firstAnn @VarTag annSet
      in checkAnnsCount e annSet (1, 0, 1) $> SOME ta va
    ExpPrim' _ "NONE" [t] _ -> do
      checkAnnsCount e annSet (1, 0, 1)
      let ta = firstAnn @TypeTag annSet
      let va = firstAnn @VarTag annSet
      t' <- fromExp @x @Ty t
      pure $ NONE ta va t'
    ExpPrim' _ "UNIT" [] _ ->
      let ta = firstAnn @TypeTag annSet
          va = firstAnn @VarTag annSet
      in checkAnnsCount e annSet (1, 0, 1) $> UNIT ta va
    ExpPrim' _ "IF_NONE" [ops1, ops2] [] -> do
      ops1' <- fromExp @x @[op] ops1
      ops2' <- fromExp @x @[op] ops2
      pure $ IF_NONE ops1' ops2'
    ExpPrim' _ "PAIR" [] _ ->
      let ta = firstAnn @TypeTag annSet
          va = firstAnn @VarTag annSet
          fa1 = firstAnn @FieldTag annSet
          fa2 = secondAnn @FieldTag annSet
      in (checkAnnsCount e annSet (1, 2, 1)) $> PAIR ta va fa1 fa2
    ExpPrim' _ "UNPAIR" [] _ ->
      let va1 = firstAnn @VarTag annSet
          va2 = secondAnn @VarTag annSet
          fa1 = firstAnn @FieldTag annSet
          fa2 = secondAnn @FieldTag annSet
      in checkAnnsCount e annSet (0, 2, 2) $> UNPAIR va1 va2 fa1 fa2
    ExpPrim' _ "PAIR" [n] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      n' <- integralFromExpr n
      let va = firstAnn @VarTag annSet
      pure $ PAIRN va n'
    ExpPrim' _ "UNPAIR" [n] [] -> do
      n' <- integralFromExpr n
      pure $ UNPAIRN n'
    ExpPrim' _ "CAR" [] _ ->
      let va = firstAnn @VarTag annSet
          fa = firstAnn @FieldTag annSet
      in checkAnnsCount e annSet (0, 1, 1) $> CAR va fa
    ExpPrim' _ "CDR" [] _ ->
      let va = firstAnn @VarTag annSet
          fa = firstAnn @FieldTag annSet
      in checkAnnsCount e annSet (0, 1, 1) $> CDR va fa
    ExpPrim' _ "LEFT" [t] _ -> do
      checkAnnsCount e annSet (1, 2, 1)
      t' <- fromExp @x @Ty t
      let ta = firstAnn @TypeTag annSet
      let va = firstAnn @VarTag annSet
      let fa1 = firstAnn @FieldTag annSet
      let fa2 = secondAnn @FieldTag annSet
      pure $ LEFT ta va fa1 fa2 t'
    ExpPrim' _ "RIGHT" [t] _ -> do
      checkAnnsCount e annSet (1, 2, 1)
      t' <- fromExp @x @Ty t
      let ta = firstAnn @TypeTag annSet
      let va = firstAnn @VarTag annSet
      let fa1 = firstAnn @FieldTag annSet
      let fa2 = secondAnn @FieldTag annSet
      pure $ RIGHT ta va fa1 fa2 t'
    ExpPrim' _ "IF_LEFT" [ops1, ops2] [] -> do
      ops1' <- fromExp @x @[op] ops1
      ops2' <- fromExp @x @[op] ops2
      pure $ IF_LEFT ops1' ops2'
    ExpPrim' _ "NIL" [t] _ -> do
      checkAnnsCount e annSet (1, 0, 1)
      t' <- fromExp @x @Ty t
      let va = firstAnn @VarTag annSet
      let ta = firstAnn @TypeTag annSet
      pure $ NIL ta va t'
    ExpPrim' _ "CONS" [] anns -> mkInstrWithVarAnn CONS anns
    ExpPrim' _ "IF_CONS" [ops1, ops2] [] -> do
      ops1' <- fromExp @x @[op] ops1
      ops2' <- fromExp @x @[op] ops2
      pure $ IF_CONS ops1' ops2'
    ExpPrim' _ "SIZE" [] anns -> mkInstrWithVarAnn SIZE anns
    ExpPrim' _ "EMPTY_SET" [t] _ -> do
      checkAnnsCount e annSet (1, 0, 1)
      t' <- fromExp @x @Ty t
      let va = firstAnn @VarTag annSet
      let ta = firstAnn @TypeTag annSet
      pure $ EMPTY_SET ta va t'
    ExpPrim' _ "EMPTY_MAP" [kt, vt] _ -> do
      checkAnnsCount e annSet (1, 0, 1)
      kt' <- fromExp @x @Ty kt
      vt' <- fromExp @x @Ty vt
      let va = firstAnn @VarTag annSet
      let ta = firstAnn @TypeTag annSet
      pure $ EMPTY_MAP ta va kt' vt'
    ExpPrim' _ "EMPTY_BIG_MAP" [kt, vt] _ -> do
      checkAnnsCount e annSet (1, 0, 1)
      kt' <- fromExp @x @Ty kt
      vt' <- fromExp @x @Ty vt
      let va = firstAnn @VarTag annSet
      let ta = firstAnn @TypeTag annSet
      pure $ EMPTY_BIG_MAP ta va kt' vt'
    ExpPrim' _ "MAP" [ops] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      ops' <- fromExp @x @[op] ops
      let va = firstAnn @VarTag annSet
      pure $ MAP va ops'
    ExpPrim' _ "ITER" [ops] [] -> do
      ops' <- fromExp @x @[op] ops
      pure $ ITER ops'
    ExpPrim' _ "MEM" [] anns -> mkInstrWithVarAnn MEM anns
    ExpPrim' _ "GET" [] anns -> mkInstrWithVarAnn GET anns
    ExpPrim' _ "GET" [n] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      n' <- integralFromExpr n
      let va = firstAnn @VarTag annSet
      pure $ GETN va n'
    ExpPrim' _ "UPDATE" [] anns -> mkInstrWithVarAnn UPDATE anns
    ExpPrim' _ "UPDATE" [n] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      n' <- integralFromExpr n
      let va = firstAnn @VarTag annSet
      pure $ UPDATEN va n'
    ExpPrim' _ "GET_AND_UPDATE" [] anns ->
      mkInstrWithVarAnn GET_AND_UPDATE anns
    ExpPrim' _ "IF" [ops1, ops2] [] -> do
      ops1' <- fromExp @x @[op] ops1
      ops2' <- fromExp @x @[op] ops2
      pure $ IF ops1' ops2'
    ExpPrim' _ "LOOP" [ops] [] -> do
      ops' <- fromExp @x @[op] ops
      pure $ LOOP ops'
    ExpPrim' _ "LOOP_LEFT" [ops] [] -> do
      ops' <- fromExp @x @[op] ops
      pure $ LOOP_LEFT ops'
    ExpPrim' _ "LAMBDA" [inp, out, ops] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      inp' <- fromExp @x @Ty inp
      out' <- fromExp @x @Ty out
      ops' <- fromExp @x @[op] ops
      let va = firstAnn @VarTag annSet
      pure $ LAMBDA va inp' out' ops'
    ExpPrim' _ "EXEC" [] anns -> mkInstrWithVarAnn EXEC anns
    ExpPrim' _ "APPLY" [] anns -> mkInstrWithVarAnn APPLY anns
    ExpPrim' _ "DIP" [ops] [] -> do
      ops' <- fromExp @x @[op] ops
      pure $ DIP ops'
    ExpPrim' _ "DIP" [n, ops] [] -> do
      n' <- integralFromExpr n
      ops' <- fromExp @x @[op] ops
      pure $ DIPN n' ops'
    ExpPrim' _ "FAILWITH" [] [] -> pure FAILWITH
    ExpPrim' _ "CAST" [t] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      t' <- fromExp @x @Ty t
      let va = firstAnn @VarTag annSet
      pure $ CAST va t'
    ExpPrim' _ "RENAME" [] anns -> mkInstrWithVarAnn RENAME anns
    ExpPrim' _ "PACK" [] anns -> mkInstrWithVarAnn PACK anns
    ExpPrim' _ "UNPACK" [t] _ -> do
      checkAnnsCount e annSet (1, 0, 1)
      t' <- fromExp @x @Ty t
      let ta = firstAnn @TypeTag annSet
      let va = firstAnn @VarTag annSet
      pure $ UNPACK ta va t'
    ExpPrim' _ "CONCAT" [] anns -> mkInstrWithVarAnn CONCAT anns
    ExpPrim' _ "SLICE" [] anns -> mkInstrWithVarAnn SLICE anns
    ExpPrim' _ "ISNAT" [] anns -> mkInstrWithVarAnn ISNAT anns
    ExpPrim' _ "ADD" [] anns -> mkInstrWithVarAnn ADD anns
    ExpPrim' _ "SUB" [] anns -> mkInstrWithVarAnn SUB anns
    ExpPrim' _ "SUB_MUTEZ" [] anns -> mkInstrWithVarAnn SUB_MUTEZ anns
    ExpPrim' _ "MUL" [] anns -> mkInstrWithVarAnn MUL anns
    ExpPrim' _ "EDIV" [] anns -> mkInstrWithVarAnn EDIV anns
    ExpPrim' _ "ABS" [] anns -> mkInstrWithVarAnn ABS anns
    ExpPrim' _ "NEG" [] anns -> mkInstrWithVarAnn NEG anns
    ExpPrim' _ "LSL" [] anns -> mkInstrWithVarAnn LSL anns
    ExpPrim' _ "LSR" [] anns -> mkInstrWithVarAnn LSR anns
    ExpPrim' _ "OR" [] anns -> mkInstrWithVarAnn OR anns
    ExpPrim' _ "AND" [] anns -> mkInstrWithVarAnn AND anns
    ExpPrim' _ "XOR" [] anns -> mkInstrWithVarAnn XOR anns
    ExpPrim' _ "NOT" [] anns -> mkInstrWithVarAnn NOT anns
    ExpPrim' _ "COMPARE" [] anns -> mkInstrWithVarAnn COMPARE anns
    ExpPrim' _ "EQ" [] anns -> mkInstrWithVarAnn Untyped.EQ anns
    ExpPrim' _ "NEQ" [] anns -> mkInstrWithVarAnn NEQ anns
    ExpPrim' _ "LT" [] anns -> mkInstrWithVarAnn Untyped.LT anns
    ExpPrim' _ "GT" [] anns -> mkInstrWithVarAnn Untyped.GT anns
    ExpPrim' _ "LE" [] anns -> mkInstrWithVarAnn LE anns
    ExpPrim' _ "GE" [] anns -> mkInstrWithVarAnn GE anns
    ExpPrim' _ "INT" [] anns -> mkInstrWithVarAnn INT anns
    ExpPrim' _ "VIEW" [name, t] _ -> do
      let va = firstAnn @VarTag annSet
      name' <- fromExp @x @ViewName name
      t' <- fromExp @x @Ty t
      checkAnnsCount e annSet (0, 0, 1) $> VIEW va name' t'
    ExpPrim' _ "SELF" [] _ ->
      let fa = firstAnn @FieldTag annSet
          va = firstAnn @VarTag annSet
      in checkAnnsCount e annSet (0, 1, 1) $> SELF va fa
    ExpPrim' _ "CONTRACT" [t] _ -> do
      checkAnnsCount e annSet (0, 1, 1)
      t' <- fromExp @x @Ty t
      let va = firstAnn @VarTag annSet
      let fa = firstAnn @FieldTag annSet
      pure $ CONTRACT va fa t'
    ExpPrim' _ "TRANSFER_TOKENS" [] anns ->
      mkInstrWithVarAnn TRANSFER_TOKENS anns
    ExpPrim' _ "SET_DELEGATE" [] anns ->
      mkInstrWithVarAnn SET_DELEGATE anns
    ExpPrim' _ "CREATE_CONTRACT" [c] _ -> do
      checkAnnsCount e annSet (0, 0, 2)
      c' <- fromExp @x @(Untyped.Contract' op) c
      let va1 = firstAnn @VarTag annSet
      let va2 = secondAnn @VarTag annSet
      pure $ CREATE_CONTRACT va1 va2 c'
    ExpPrim' _ "IMPLICIT_ACCOUNT" [] anns ->
      mkInstrWithVarAnn IMPLICIT_ACCOUNT anns
    ExpPrim' _ "NOW" [] anns -> mkInstrWithVarAnn NOW anns
    ExpPrim' _ "AMOUNT" [] anns -> mkInstrWithVarAnn AMOUNT anns
    ExpPrim' _ "BALANCE" [] anns -> mkInstrWithVarAnn BALANCE anns
    ExpPrim' _ "VOTING_POWER" [] anns ->
      mkInstrWithVarAnn VOTING_POWER anns
    ExpPrim' _ "TOTAL_VOTING_POWER" [] anns ->
      mkInstrWithVarAnn TOTAL_VOTING_POWER anns
    ExpPrim' _ "CHECK_SIGNATURE" [] anns ->
      mkInstrWithVarAnn CHECK_SIGNATURE anns
    ExpPrim' _ "SHA256" [] anns -> mkInstrWithVarAnn SHA256 anns
    ExpPrim' _ "SHA512" [] anns -> mkInstrWithVarAnn SHA512 anns
    ExpPrim' _ "BLAKE2B" [] anns -> mkInstrWithVarAnn BLAKE2B anns
    ExpPrim' _ "SHA3" [] anns -> mkInstrWithVarAnn SHA3 anns
    ExpPrim' _ "KECCAK" [] anns -> mkInstrWithVarAnn KECCAK anns
    ExpPrim' _ "HASH_KEY" [] anns -> mkInstrWithVarAnn HASH_KEY anns
    ExpPrim' _ "PAIRING_CHECK" [] anns ->
      mkInstrWithVarAnn PAIRING_CHECK anns
    ExpPrim' _ "SOURCE" [] anns -> mkInstrWithVarAnn SOURCE anns
    ExpPrim' _ "SENDER" [] anns -> mkInstrWithVarAnn SENDER anns
    ExpPrim' _ "ADDRESS" [] anns -> mkInstrWithVarAnn ADDRESS anns
    ExpPrim' _ "CHAIN_ID" [] anns -> mkInstrWithVarAnn CHAIN_ID anns
    ExpPrim' _ "LEVEL" [] anns -> mkInstrWithVarAnn LEVEL anns
    ExpPrim' _ "SELF_ADDRESS" [] anns ->
      mkInstrWithVarAnn SELF_ADDRESS anns
    ExpPrim' _ "NEVER" [] [] -> pure NEVER
    ExpPrim' _ "TICKET" [] anns -> mkInstrWithVarAnn TICKET anns
    ExpPrim' _ "READ_TICKET" [] anns ->
      mkInstrWithVarAnn READ_TICKET anns
    ExpPrim' _ "SPLIT_TICKET" [] anns ->
      mkInstrWithVarAnn SPLIT_TICKET anns
    ExpPrim' _ "JOIN_TICKETS" [] anns ->
      mkInstrWithVarAnn JOIN_TICKETS anns
    ExpPrim' _ "OPEN_CHEST" [] anns ->
      mkInstrWithVarAnn OPEN_CHEST anns
    ExpPrim' _ "SAPLING_EMPTY_STATE" [n] _ -> do
      checkAnnsCount e annSet (0, 0, 1)
      n' <- integralFromExpr n
      let va = firstAnn @VarTag annSet
      pure $ SAPLING_EMPTY_STATE va n'
    ExpPrim' _ "SAPLING_VERIFY_UPDATE" [] anns ->
      mkInstrWithVarAnn SAPLING_VERIFY_UPDATE anns
    ExpPrim' _ "MIN_BLOCK_TIME" [] anns -> pure $ mkInstrWithAnyAnns MIN_BLOCK_TIME anns
    _ -> Left $ FromExpError e "Expected an instruction"

    where
      mkInstrWithVarAnn
        :: (VarAnn -> InstrAbstract op)
        -> [Annotation]
        -> Either (FromExpError x) (InstrAbstract op)
      mkInstrWithVarAnn ctor anns =
        let annSet = toAnnSet anns
            va = firstAnn @VarTag annSet
        in checkAnnsCount e annSet (0, 0, 1) $> ctor va

      mkInstrWithAnyAnns
        :: ([Untyped.AnyAnn] -> InstrAbstract op)
        -> [Annotation]
        -> InstrAbstract op
      mkInstrWithAnyAnns ctor anns = ctor $ anns <&> \case
        AnnotationType x -> Untyped.AnyAnnType x
        AnnotationField x -> Untyped.AnyAnnField x
        AnnotationVariable x -> Untyped.AnyAnnVar x

      getAnnSet :: Exp d -> AnnotationSet
      getAnnSet = \case
        ExpPrim' _ _ _ anns -> toAnnSet anns
        _                   -> emptyAnnSet

instance (FromExp x op) => FromExp x (Untyped.Contract' op) where
  fromExp blocks = case blocks of
    ExpSeq _ bs -> do
      bs' <- mapM exprToCB bs
      maybeToRight (FromExpError blocks "Something's wrong with top-level contract blocks")
        (orderContractBlock bs')
    expr -> Left $ FromExpError expr "Failed to parse contract, expected sequence"
    where
      exprToCB
        :: Exp x
        -> Either (FromExpError x) (ContractBlock op)
      exprToCB e = case e of
        ExpPrim' _ "parameter" args anns -> mkCbParam e args anns
        ExpPrim' _ "storage"   args anns -> mkCBStorage e args anns
        ExpPrim' _ "code"      args anns -> mkCBCode e args anns
        ExpPrim' _ "view"      args anns -> mkCBView e args anns
        _                              ->
          Left $ FromExpError e "Unexpected primitive at contract top-level"

      mkCbParam
        :: Exp x
        -> [Exp x]
        -> [Annotation]
        -> Either (FromExpError x) (ContractBlock op)
      mkCbParam e args anns = case (args, anns) of
        ([p], []) -> do
          let annSet = toAnnSet (p ^. _ExpPrim . _2 . mpaAnnotsL)
          let rootAnn = firstAnn @FieldTag annSet
          unless (secondAnn @FieldTag annSet == noAnn) $
            Left $ FromExpError p
              "Expected parameter with at most 1 root annotation"
          p' <- fromExp @x @Ty
            (p & _ExpPrim . _2 . mpaAnnotsL %~ filter (not . isAnnotationField))
          pure $ CBParam $ Untyped.ParameterType p' rootAnn
        _ -> Left $ FromExpError e
          "Expected 'parameter' block without annotations and exactly 1 argument"

      mkCBStorage
        :: Exp x
        -> [Exp x]
        -> [Annotation]
        -> Either (FromExpError x) (ContractBlock op)
      mkCBStorage e args anns = case (args, anns) of
        ([s], []) -> do
          s' <- fromExp @x @Ty s
          pure $ CBStorage s'
        _ -> Left $ FromExpError e
               "Expected 'storage' block without annotations and exactly 1 argument"

      mkCBCode
        :: Exp x
        -> [Exp x]
        -> [Annotation]
        -> Either (FromExpError x) (ContractBlock op)
      mkCBCode e args anns = case (args, anns) of
        ([ops], []) -> do
          ops' <- fromExp @x @[op] ops
          pure $ CBCode ops'
        _ -> Left $ FromExpError e
               "Expected 'code' block without annotations"

      mkCBView
        :: Exp x
        -> [Exp x]
        -> [Annotation]
        -> Either (FromExpError x) (ContractBlock op)
      mkCBView e args anns = case (args, anns) of
        ([name, arg, ret, ops], []) -> do
          name' <- fromExp name
          arg' <- fromExp arg
          ret' <- fromExp ret
          ops' <- fromExp @x @[op] ops
          pure $ CBView $ Untyped.View name' arg' ret' ops'
        (_, _ : _) ->
          Left $ FromExpError e
               "Expected 'view' block without annotations"
        (_, []) ->
          Left $ FromExpError e
               "Invalid 'view' block, expected 4 expressions in it"

instance FromExp x Untyped.T where
  fromExp e = case e of
    ExpPrim' _ "key" [] [] -> pure Untyped.TKey
    ExpPrim' _ "unit" [] [] -> pure Untyped.TUnit
    ExpPrim' _ "signature" [] [] -> pure Untyped.TSignature
    ExpPrim' _ "chain_id" [] [] -> pure Untyped.TChainId
    ExpPrim' _ "option" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.TOption arg'
    ExpPrim' _ "list" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.TList arg'
    ExpPrim' _ "set" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.TSet arg'
    ExpPrim' _ "operation" [] [] -> pure Untyped.TOperation
    ExpPrim' _ "contract" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.TContract arg'
    ExpPrim' _ "ticket" [arg] [] -> do
      arg' <- fromExp arg
      pure $ Untyped.TTicket arg'
    ExpPrim' _ "or" [arg1, arg2] [] -> do
      let as1 = toAnnSet $ arg1 ^. _ExpPrim . _2 . mpaAnnotsL
      let as2 = toAnnSet $ arg2 ^. _ExpPrim . _2 . mpaAnnotsL
      checkAnnsCount e as1 (1, 1, 0)
      checkAnnsCount e as2 (1, 1, 0)
      let fa1 = firstAnn @FieldTag as1
      let fa2 = firstAnn @FieldTag as2
      l <- fromExp $ removeAnns arg1 isAnnotationField
      r <- fromExp $ removeAnns arg2 isAnnotationField
      pure $ Untyped.TOr fa1 fa2 l r
    ExpPrim' _ "pair" args [] -> do
      args2 <- case nonEmpty args >>= forbidSingletonList of
        Nothing -> Left $ FromExpError e
          "Expected a pair with at least 2 arguments"
        Just as -> pure as
      -- Check and extract annotations info
      tyInfos <- forM args2 \arg -> do
        let as = toAnnSet $ arg ^. _ExpPrim . _2 . mpaAnnotsL
        checkAnnsCount e as (1, 1, 1)
        let fa = firstAnn @FieldTag as
        let va = firstAnn @VarTag as
        ty <- fromExp $ removeAnns arg
          (isAnnotationField || isAnnotationVariable)
        return (ty, fa, va)
      -- Make a right-comb pairs tree
      let combiner (ty1, fa1, va1) (ty2, fa2, va2) =
            ( Ty (Untyped.TPair fa1 fa2 va1 va2 ty1 ty2) noAnn
            , noAnn
            , noAnn
            )
      let (Ty tRes _, _, _) = foldr1 combiner tyInfos
      return tRes
    ExpPrim' _ "lambda" args [] -> mkDoubleParamType Untyped.TLambda args e
      "Expected a lambda with input and output types"
    ExpPrim' _ "map" args [] -> mkDoubleParamType Untyped.TMap args e
      "Expected a map with key and value types"
    ExpPrim' _ "big_map" args [] -> mkDoubleParamType Untyped.TBigMap args e
      "Expected a big_map with key and value types"
    ExpPrim' _ "int" [] [] -> pure Untyped.TInt
    ExpPrim' _ "nat" [] [] -> pure Untyped.TNat
    ExpPrim' _ "string" [] [] -> pure Untyped.TString
    ExpPrim' _ "bytes" [] [] -> pure Untyped.TBytes
    ExpPrim' _ "mutez" [] [] -> pure Untyped.TMutez
    ExpPrim' _ "bool" [] [] -> pure Untyped.TBool
    ExpPrim' _ "key_hash" [] [] -> pure Untyped.TKeyHash
    ExpPrim' _ "bls12_381_fr" [] [] -> pure Untyped.TBls12381Fr
    ExpPrim' _ "bls12_381_g1" [] [] -> pure Untyped.TBls12381G1
    ExpPrim' _ "bls12_381_g2" [] [] -> pure Untyped.TBls12381G2
    ExpPrim' _ "timestamp" [] [] -> pure Untyped.TTimestamp
    ExpPrim' _ "address" [] [] -> pure Untyped.TAddress
    ExpPrim' _ "chest" [] [] -> pure Untyped.TChest
    ExpPrim' _ "chest_key" [] [] -> pure Untyped.TChestKey
    ExpPrim' _ "tx_rollup_l2_address" [] [] -> pure Untyped.TTxRollupL2Address
    ExpPrim' _ "never" [] [] -> pure Untyped.TNever
    ExpPrim' _ "sapling_state" [n] [] -> do
      n' <- integralFromExpr n
      pure $ Untyped.TSaplingState n'
    ExpPrim' _ "sapling_transaction" [n] [] -> do
      n' <- integralFromExpr n
      pure $ Untyped.TSaplingTransaction n'
    ExpPrim' _ "sapling_transaction_deprecated" _ _ -> do
      Left $ FromExpError e "Use of deprecated type: sapling_transaction_deprecated"
    _ -> Left $ FromExpError e "Expected a type"
    where
      mkDoubleParamType
        :: (Ty -> Ty -> Untyped.T)
        -> [Exp x]
        -> Exp x
        -> Text
        -> Either (FromExpError x) Untyped.T
      mkDoubleParamType ctor args expr msg = do
        case args of
          [arg1, arg2] -> do
            arg1' <- fromExp arg1
            arg2' <- fromExp arg2
            pure $ ctor arg1' arg2'
          _ -> Left $ FromExpError expr msg

      removeAnns :: Exp x -> (Annotation -> Bool) -> Exp x
      removeAnns expr p =
        expr & _ExpPrim . _2 . mpaAnnotsL %~ filter (not . p)

instance FromExp x Ty where
  fromExp e = case e of
    ExpPrim' ex primName args anns -> do
      let annSet = toAnnSet anns
      let ta = firstAnn @TypeTag annSet
      when (secondAnn @TypeTag annSet /= noAnn) $
        Left $ FromExpError e
          "Expected expression with at most 1 type annotation"
      t <- fromExp @x @Untyped.T $ ExpPrim' ex primName args $ filter
        (not . isAnnotationType) anns
      pure $ Ty t ta
    _ -> Left $ FromExpError e "Expected a type"

instance FromExp x T where
  fromExp =
    second fromUType . fromExp @x @Untyped.Ty

-- Note: we should generalize this to work for any instruction,
-- not just lambdas (i.e. instructions with one input and one output).
instance (SingI inp, SingI out) => FromExp RegularExp (Instr '[inp] '[out]) where
  fromExp expr =
    fromExpression @(Value ('TLambda inp out)) expr <&> \case
      VLam instr -> rfAnyInstr instr

instance FromExp x ViewName where
  fromExp e = case e of
    ExpString _ s ->
      first (FromExpError e . pretty) $ mkViewName s
    _ -> Left $ FromExpError e "Expected view name"

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

-- | Adds annotations to the expression, after removing empty annotations
-- at the end of each list.
addTrimmedAnns
  :: Exp x
  -> [TypeAnn]
  -> [FieldAnn]
  -> [VarAnn]
  -> Exp x
addTrimmedAnns e tas fas vas =
  e & _ExpPrim . _2 . mpaAnnotsL <>~ mkAnns tas fas vas

-- | Inserts the root annotation into the contract parameter.
insertRootAnn :: HasCallStack => Expression -> RootAnn -> Expression
insertRootAnn expr rootAnn = case expr of
  ExpPrim () p
    -- The order of annotations is important iff there are
    -- multiple annotations of the same kind or there are
    -- other kinds of annotations in the list.
    -- Prepending root field annotation is okay because
    -- there can not be more than one root annotation.
    | rootAnn /= noAnn -> expressionPrim p
      { mpaAnnots = AnnotationField rootAnn : mpaAnnots p
      }
    | otherwise -> expr
  -- Currently this error can't happen because parameter type
  -- must be a Micheline primitive. If it ever changes, we
  -- would like to notice it ASAP and update this place.
  _ -> error $ "parameter is not a primitive: " <> pretty expr

-- | Checks for a given expression that the number of annotations
-- of each type in it doesn't exceed the specified threshold.
checkAnnsCount
  :: Exp x
  -> AnnotationSet
  -> (Int, Int, Int)
  -> Either (FromExpError x) ()
checkAnnsCount e annSet maxCount@(maxTas, maxFas, maxVas) =
  let actualCount@(tasCnt, fasCnt, vasCnt) = annsCount annSet
  in unless (actualCount <= maxCount) $
    Left $ FromExpError e $ pretty $ unlinesF
      [ "Expected at most"
      , indentF 2 $ build maxTas <> " type annotations,"
      , indentF 2 $ build maxFas <> " field annotations,"
      , indentF 2 $ build maxVas <> " variable annotations"
      , "but found:"
      , indentF 2 $ build tasCnt <> " type annotations,"
      , indentF 2 $ build fasCnt <> " field annotations,"
      , indentF 2 $ build vasCnt <> " variable annotations."
      ]

forbidSingletonList :: NonEmpty a -> Maybe (NonEmpty a)
forbidSingletonList = \case
  _ :| [] -> Nothing
  x       -> Just x

integralToExpr :: Integral i => i -> Expression
integralToExpr = expressionInt . toInteger

integralFromExpr :: (Integral i, Bits i) => Exp x -> Either (FromExpError x) i
integralFromExpr e = case e of
  ExpInt _ v ->
    maybeToRight (FromExpError e "Value is out of bounds")
      (fromIntegralMaybe @Integer v)
  _ -> Left $ FromExpError e "Expected a number here"